diff --git a/gadget-code/frontend/src/components/SessionPanel.tsx b/gadget-code/frontend/src/components/SessionPanel.tsx index da6f59e..3f46f62 100644 --- a/gadget-code/frontend/src/components/SessionPanel.tsx +++ b/gadget-code/frontend/src/components/SessionPanel.tsx @@ -41,6 +41,7 @@ interface SessionPanelProps { truncateModelName: (name: string, maxLength?: number) => string; getSelectedModelCapabilities: () => { hasThinking: boolean } | null; getSelectedModel: () => AiModel | null; + contextWindowUsage: number; } export default function SessionPanel({ @@ -82,7 +83,9 @@ export default function SessionPanel({ truncateModelName, getSelectedModelCapabilities, getSelectedModel, + contextWindowUsage, }: SessionPanelProps) { + const usagePercent = sessionNumCtx > 0 ? (contextWindowUsage / sessionNumCtx) * 100 : 0; return (
@@ -292,6 +295,33 @@ export default function SessionPanel({
)} +
+
Context Window Usage
+
+
+ E +
+
= 80 ? 'bg-red-500' : + usagePercent >= 50 ? 'bg-yellow-500' : + 'bg-green-500' + }`} + style={{ width: `${Math.min(usagePercent, 100)}%` }} + /> +
+ F +
+
+ + {contextWindowUsage.toLocaleString()} tokens + + + {sessionNumCtx.toLocaleString()} max + +
+
+
); diff --git a/gadget-code/frontend/src/lib/api.ts b/gadget-code/frontend/src/lib/api.ts index 2d15bd2..cb67ced 100644 --- a/gadget-code/frontend/src/lib/api.ts +++ b/gadget-code/frontend/src/lib/api.ts @@ -388,6 +388,7 @@ export interface ChatSession { selectedModel: string; reasoningEffort?: string; numCtx?: number; + contextWindowUsage: number; stats: ChatSessionStats; pins: Array<{ _id?: string; content: string }>; } @@ -401,6 +402,14 @@ export interface ChatTurnStats { durationLabel: string; } +export interface IWorkOrderCompleteStats { + masterInputTokens: number; + masterOutputTokens: number; + masterThinkingTokens: number; + toolCallCount: number; + durationMs: number; +} + export interface ChatTurnPrompts { user: string; system?: string; diff --git a/gadget-code/frontend/src/pages/ChatSessionView.tsx b/gadget-code/frontend/src/pages/ChatSessionView.tsx index 32179eb..bf03f8a 100644 --- a/gadget-code/frontend/src/pages/ChatSessionView.tsx +++ b/gadget-code/frontend/src/pages/ChatSessionView.tsx @@ -2,7 +2,7 @@ import { useState, useEffect, useRef, useContext, useCallback } from 'react'; import { useParams, useNavigate, useLocation } from 'react-router-dom'; import { BrushCleaning, Maximize2 } from 'lucide-react'; import { socketClient } from '../lib/socket'; -import { chatSessionApi, projectApi, providerApi, type ChatSession, type ChatTurn, type ChatTurnBlock, ChatTurnStats, ChatSessionMode, type AiProvider, type Project } from '../lib/api'; +import { chatSessionApi, projectApi, providerApi, type ChatSession, type ChatTurn, type ChatTurnBlock, ChatTurnStats, ChatSessionMode, type AiProvider, type Project, type IWorkOrderCompleteStats } from '../lib/api'; import { WorkspaceMode } from '../lib/types'; import SessionPanel from '../components/SessionPanel'; import ProjectPanel from '../components/ProjectPanel'; @@ -686,7 +686,7 @@ export default function ChatSessionView() { }, 4000); }, []); - const handleWorkOrderComplete = useCallback((turnId: string, success: boolean, message?: string) => { + const handleWorkOrderComplete = useCallback((turnId: string, success: boolean, message?: string, stats?: IWorkOrderCompleteStats) => { // Backend has already flushed and persisted all streaming content // Just clean up frontend streaming state and update status if (streamingStateRef.current.has(turnId)) { @@ -734,6 +734,38 @@ export default function ChatSessionView() { setIsAborting(false); setIsProcessing(false); currentTurnIdRef.current = null; + + // Update turn stats from completion event + if (stats) { + setTurns(prev => prev.map(t => { + if (t._id !== turnId) return t; + return { + ...t, + stats: { + ...t.stats, + inputTokens: stats.masterInputTokens, + thinkingTokenCount: stats.masterThinkingTokens, + responseTokens: stats.masterOutputTokens, + toolCallCount: stats.toolCallCount, + durationMs: stats.durationMs, + durationLabel: stats.durationMs < 1000 + ? `${stats.durationMs}ms` + : `${(stats.durationMs / 1000).toFixed(1)}s`, + }, + }; + })); + + // Update session contextWindowUsage for fuel gauge + setSession(prev => prev ? { + ...prev, + contextWindowUsage: (prev.contextWindowUsage ?? 0) + stats.masterInputTokens, + stats: { + ...prev.stats, + inputTokens: prev.stats.inputTokens + stats.masterInputTokens, + outputTokens: prev.stats.outputTokens + stats.masterOutputTokens, + }, + } : prev); + } }, [showToast]); const handleWorkspaceModeChanged = useCallback((mode: string) => { @@ -1441,6 +1473,7 @@ export default function ChatSessionView() { truncateModelName={truncateModelName} getSelectedModelCapabilities={getSelectedModelCapabilities} getSelectedModel={getSelectedModel} + contextWindowUsage={session?.contextWindowUsage ?? 0} /> { if (!this.chatSessionId) { this.log.warn("workOrderComplete event received but no chat session is active"); @@ -314,20 +317,55 @@ export class DroneSession extends SocketSession { turn.errorMessage = message; } } + + // UPDATE TURN STATS — master agent tokens only + if (stats) { + turn.stats.inputTokens = stats.masterInputTokens; + turn.stats.thinkingTokenCount = stats.masterThinkingTokens; + turn.stats.responseTokens = stats.masterOutputTokens; + turn.stats.toolCallCount = stats.toolCallCount; + turn.stats.durationMs = stats.durationMs; + turn.stats.durationLabel = this.formatDurationLabel(stats.durationMs); + } + await turn.save(); + + // COMPUTE AGGREGATE: walk the turn's subagent records for aggregate totals + let totalInputTokens = stats?.masterInputTokens ?? 0; + let totalOutputTokens = stats?.masterOutputTokens ?? 0; + if (stats) { + for (const tc of turn.toolCalls) { + if ((tc as any).subagent?.stats) { + totalInputTokens += (tc as any).subagent.stats.inputTokens; + totalOutputTokens += (tc as any).subagent.stats.responseTokens + (tc as any).subagent.stats.thinkingTokenCount; + } + } + } + + // INCREMENT SESSION STATS — aggregate totals (master + subagents) + await ChatSession.updateOne( + { _id: this.chatSessionId }, + { $inc: { + "stats.inputTokens": totalInputTokens, + "stats.outputTokens": totalOutputTokens, + "stats.toolCallCount": stats?.toolCallCount ?? 0, + // Master-only field for the fuel gauge + "contextWindowUsage": stats?.masterInputTokens ?? 0, + }}, + ); } const codeSession = SocketService.getCodeSessionByChatSessionId( this.chatSessionId, ); - codeSession.onWorkOrderComplete(turnId, success, message); + codeSession.onWorkOrderComplete(turnId, success, message, stats); this.currentTurnId = undefined; } catch (error) { // Routing failed - queue to Redis await MessageQueue.enqueue(this.chatSessionId, { type: 'workOrderComplete', - args: [turnId, success, message], + args: [turnId, success, message, stats], timestamp: Date.now(), }); this.log.debug("queued workOrderComplete message", { chatSessionId: this.chatSessionId }); @@ -343,6 +381,17 @@ export class DroneSession extends SocketSession { this.streamingBuffers.clear(); } + private formatDurationLabel(durationMs: number): string { + if (durationMs < 1000) return `${durationMs}ms`; + const seconds = Math.floor(durationMs / 1000); + const minutes = Math.floor(seconds / 60); + const secs = seconds % 60; + if (minutes > 0) { + return `${minutes}m ${secs}s`; + } + return `${seconds}.${Math.floor((durationMs % 1000) / 100)}s`; + } + /** * Gets or creates a streaming buffer for a turn. */ @@ -669,7 +718,8 @@ export class DroneSession extends SocketSession { codeSession.onWorkOrderComplete( msg.args[0] as string, msg.args[1] as boolean, - msg.args[2] as string, + msg.args[2] as string | undefined, + msg.args[3] as IWorkOrderCompleteStats | undefined, ); break; case 'log': diff --git a/gadget-code/src/models/chat-session.ts b/gadget-code/src/models/chat-session.ts index e55ec58..22b44e5 100644 --- a/gadget-code/src/models/chat-session.ts +++ b/gadget-code/src/models/chat-session.ts @@ -31,6 +31,7 @@ export const ChatSessionSchema = new Schema({ default: "off", }, numCtx: { type: Number }, + contextWindowUsage: { type: Number, default: 0, required: true }, stats: { turnCount: { type: Number, default: 0, required: true }, toolCallCount: { type: Number, default: 0, required: true }, diff --git a/gadget-code/src/services/chat-session.ts b/gadget-code/src/services/chat-session.ts index cc82740..31d326a 100644 --- a/gadget-code/src/services/chat-session.ts +++ b/gadget-code/src/services/chat-session.ts @@ -131,6 +131,7 @@ class ChatSessionService extends DtpService { inputTokens: 0, outputTokens: 0, }, + contextWindowUsage: 0, pins: [], }); diff --git a/gadget-drone/src/services/agent.test.ts b/gadget-drone/src/services/agent.test.ts index b0d7d4d..c4bebfd 100644 --- a/gadget-drone/src/services/agent.test.ts +++ b/gadget-drone/src/services/agent.test.ts @@ -39,6 +39,7 @@ describe("AgentService", () => { inputTokens: 0, outputTokens: 0, }, + contextWindowUsage: 0, pins: [], }; const workOrder: IAgentWorkOrder = { @@ -150,6 +151,7 @@ describe("AgentService", () => { inputTokens: 0, outputTokens: 0, }, + contextWindowUsage: 0, pins: [], }; @@ -314,6 +316,7 @@ describe("AgentService", () => { inputTokens: 0, outputTokens: 0, }, + contextWindowUsage: 0, pins: [], }; const rawResponse = "Subagent completed the task successfully."; diff --git a/gadget-drone/src/services/agent.ts b/gadget-drone/src/services/agent.ts index a4aea64..f131cb6 100644 --- a/gadget-drone/src/services/agent.ts +++ b/gadget-drone/src/services/agent.ts @@ -58,6 +58,14 @@ import { ReadSkillTool, type DroneToolboxEnvironment, } from "../tools/index.ts"; + +interface IWorkOrderCompleteStats { + masterInputTokens: number; + masterOutputTokens: number; + masterThinkingTokens: number; + toolCallCount: number; + durationMs: number; +} import type { IChatExportData } from "@gadget/ai-toolbox"; export interface IAgentWorkOrder { @@ -164,8 +172,10 @@ class AgentService extends GadgetService { const { turn } = workOrder; let toolCallCount = 0; - let inputTokens = 0; - let outputTokens = 0; + let masterInputTokens = 0; + let masterOutputTokens = 0; + let masterThinkingTokens = 0; + const startTime = Date.now(); // Build the full message array that grows with each iteration const messages: IContextChatMessage[] = []; @@ -258,6 +268,11 @@ class AgentService extends GadgetService { throw new Error("Model failed to respond (still loading or error)"); } + // Accumulate real token counts from the AI API response + masterInputTokens += response.stats.tokenCounts.input; + masterOutputTokens += response.stats.tokenCounts.response; + masterThinkingTokens += response.stats.tokenCounts.thinking; + // Process tool calls if present if (response.toolCalls && response.toolCalls.length > 0) { continueLoop = true; @@ -338,17 +353,30 @@ class AgentService extends GadgetService { toolName: toolCall.function.name, }); - inputTokens += Math.ceil(toolCall.function.arguments.length / 4); - outputTokens += Math.ceil(result.length / 4); } } } - socket.emit("workOrderComplete", turn._id, true); + const durationMs = Date.now() - startTime; + const stats: IWorkOrderCompleteStats = { + masterInputTokens, + masterOutputTokens, + masterThinkingTokens, + toolCallCount, + durationMs, + }; + socket.emit("workOrderComplete", turn._id, true, undefined, stats); } catch (cause) { if (cause instanceof Error && cause.name === "AbortError") { this.log.info("work order aborted by user", { turnId: turn._id }); - socket.emit("workOrderComplete", turn._id, true, "aborted"); + const abortStats: IWorkOrderCompleteStats = { + masterInputTokens, + masterOutputTokens, + masterThinkingTokens, + toolCallCount, + durationMs: Date.now() - startTime, + }; + socket.emit("workOrderComplete", turn._id, true, "aborted", abortStats); return; } const msg = cause instanceof Error ? cause.message : String(cause); @@ -356,7 +384,14 @@ class AgentService extends GadgetService { turnId: turn._id, error: msg, }); - socket.emit("workOrderComplete", turn._id, false, msg); + const errorStats: IWorkOrderCompleteStats = { + masterInputTokens, + masterOutputTokens, + masterThinkingTokens, + toolCallCount, + durationMs: Date.now() - startTime, + }; + socket.emit("workOrderComplete", turn._id, false, msg, errorStats); throw cause; } finally { this.abortController = null; @@ -740,11 +775,11 @@ class AgentService extends GadgetService { inputTokens += response.stats?.tokenCounts?.input ?? 0; if (response.thinking) { - thinkingTokenCount += Math.ceil(response.thinking.length / 4); + thinkingTokenCount += response.stats?.tokenCounts?.thinking ?? 0; } if (response.response) { fullResponse += response.response; - outputTokens += Math.ceil(response.response.length / 4); + outputTokens += response.stats?.tokenCounts?.response ?? 0; } if (response.thinking) { fullThinking = (fullThinking + response.thinking).trim(); @@ -800,8 +835,6 @@ class AgentService extends GadgetService { toolName: toolCall.function.name, }); - inputTokens += Math.ceil(toolArgsRaw.length / 4); - outputTokens += Math.ceil(result.length / 4); } } } diff --git a/packages/ai/src/openai.ts b/packages/ai/src/openai.ts index 1dc5a42..1719ea3 100644 --- a/packages/ai/src/openai.ts +++ b/packages/ai/src/openai.ts @@ -80,6 +80,11 @@ interface OpenAiChatIterationResult { chunkCount: number; contentDeltaCount: number; toolDeltaCount: number; + // Usage data from OpenAI API (prompt_tokens, completion_tokens) + usage?: { + promptTokens: number; + completionTokens: number; + }; } export class OpenAiApi extends AiApi { @@ -206,6 +211,7 @@ export class OpenAiApi extends AiApi { { role: "user" as const, content: options.prompt }, ], stream: true, + stream_options: { include_usage: true }, ...(model.params.maxCompletionTokens ? { max_completion_tokens: model.params.maxCompletionTokens } : {}), @@ -223,12 +229,21 @@ export class OpenAiApi extends AiApi { let accumulatedResponse = ""; let accumulatedThinking = ""; + let usage: { promptTokens: number; completionTokens: number } | undefined; for await (const chunk of response) { if (options.signal?.aborted) { throw new DOMException("The operation was aborted", "AbortError"); } + // Capture usage from the final streaming chunk + if (chunk.usage) { + usage = { + promptTokens: chunk.usage.prompt_tokens, + completionTokens: chunk.usage.completion_tokens, + }; + } + const delta = chunk.choices[0]?.delta; if (delta) { if (delta.content) { @@ -265,8 +280,8 @@ export class OpenAiApi extends AiApi { text: numeral(durationMs / 1000).format("hh:mm:ss"), }, tokenCounts: { - input: 0, - response: 0, + input: usage?.promptTokens ?? 0, + response: usage?.completionTokens ?? 0, thinking: 0, }, }, @@ -325,7 +340,7 @@ export class OpenAiApi extends AiApi { thinking: undefined, toolCalls: undefined, toolCallResults: undefined, - stats: this.buildStats(startTime), + stats: this.buildStats(startTime, iteration.usage), }); } @@ -334,7 +349,7 @@ export class OpenAiApi extends AiApi { response: iteration.response, thinking: iteration.thinking, toolCalls: iteration.toolCalls.length > 0 ? iteration.toolCalls : undefined, - stats: this.buildStats(startTime), + stats: this.buildStats(startTime, iteration.usage), }; this.assertNonEmptyChatResponse(finalResponse); return finalResponse; @@ -401,6 +416,7 @@ export class OpenAiApi extends AiApi { messages, tools, stream: true, + stream_options: { include_usage: true }, ...(model.params.maxCompletionTokens ? { max_completion_tokens: model.params.maxCompletionTokens } : {}), @@ -423,6 +439,7 @@ export class OpenAiApi extends AiApi { let toolDeltaCount = 0; let finishReason: string | null | undefined; const toolCallMap = new Map(); + let usage: { promptTokens: number; completionTokens: number } | undefined; for await (const chunk of response) { if (signal?.aborted) { @@ -431,6 +448,15 @@ export class OpenAiApi extends AiApi { chunkCount++; finishReason = chunk.choices[0]?.finish_reason ?? finishReason; + + // Capture usage from the final streaming chunk + if (chunk.usage) { + usage = { + promptTokens: chunk.usage.prompt_tokens, + completionTokens: chunk.usage.completion_tokens, + }; + } + const delta = chunk.choices[0]?.delta; if (!delta) continue; @@ -463,6 +489,7 @@ export class OpenAiApi extends AiApi { chunkCount, contentDeltaCount, toolDeltaCount, + usage, }; } @@ -498,6 +525,10 @@ export class OpenAiApi extends AiApi { const choice = response.choices[0]; const message = choice?.message; const content = typeof message?.content === "string" ? message.content : ""; + const usage = response.usage ? { + promptTokens: response.usage.prompt_tokens, + completionTokens: response.usage.completion_tokens, + } : undefined; const assistantToolCalls = (message?.tool_calls || []) .filter((tc) => tc.type === "function") .map((tc) => ({ @@ -524,6 +555,7 @@ export class OpenAiApi extends AiApi { chunkCount: 0, contentDeltaCount: content ? 1 : 0, toolDeltaCount: assistantToolCalls.length, + usage, }; } @@ -592,7 +624,10 @@ export class OpenAiApi extends AiApi { ); } - private buildStats(startTime: number): IAiChatResponse["stats"] { + private buildStats( + startTime: number, + usage?: { promptTokens: number; completionTokens: number }, + ): IAiChatResponse["stats"] { const seconds = (Date.now() - startTime) / 1000; return { duration: { @@ -600,8 +635,8 @@ export class OpenAiApi extends AiApi { text: numeral(seconds).format("hh:mm:ss"), }, tokenCounts: { - input: 0, - response: 0, + input: usage?.promptTokens ?? 0, + response: usage?.completionTokens ?? 0, thinking: 0, }, }; diff --git a/packages/api/src/interfaces/chat-session.ts b/packages/api/src/interfaces/chat-session.ts index da6bbd7..023f651 100644 --- a/packages/api/src/interfaces/chat-session.ts +++ b/packages/api/src/interfaces/chat-session.ts @@ -36,6 +36,7 @@ export interface IChatSession { selectedModel: string; reasoningEffort?: ReasoningEffort; numCtx?: number; + contextWindowUsage: number; stats: { turnCount: number; toolCallCount: number; diff --git a/packages/api/src/messages/drone.ts b/packages/api/src/messages/drone.ts index 3f8a3d4..b82f05f 100644 --- a/packages/api/src/messages/drone.ts +++ b/packages/api/src/messages/drone.ts @@ -41,10 +41,19 @@ export type ToolCallMessage = ( response: string, ) => void; +export interface IWorkOrderCompleteStats { + masterInputTokens: number; + masterOutputTokens: number; + masterThinkingTokens: number; + toolCallCount: number; + durationMs: number; +} + export type WorkOrderCompleteMessage = ( workOrderId: string, success: boolean, message?: string, + stats?: IWorkOrderCompleteStats, ) => void; export type RequestCrashRecoveryMessage = (data: {