From 6696d79a0d399ec17b5dbd95864d07dd12db3ab1 Mon Sep 17 00:00:00 2001 From: Rob Colbert Date: Mon, 18 May 2026 16:40:19 -0400 Subject: [PATCH] repairs for OpenAI token stats/economics --- gadget-code/frontend/src/lib/api.ts | 4 ++++ .../frontend/src/pages/ChatSessionView.tsx | 6 +++-- gadget-code/src/lib/code-session.ts | 8 ++++--- gadget-code/src/lib/drone-session.ts | 18 +++++++++++---- gadget-drone/src/services/agent.ts | 15 +++++-------- packages/api/src/messages/drone.ts | 4 ++++ repairs.md | 22 +++++++++++++++++++ 7 files changed, 58 insertions(+), 19 deletions(-) create mode 100644 repairs.md diff --git a/gadget-code/frontend/src/lib/api.ts b/gadget-code/frontend/src/lib/api.ts index cb67ced..e8849b1 100644 --- a/gadget-code/frontend/src/lib/api.ts +++ b/gadget-code/frontend/src/lib/api.ts @@ -408,6 +408,10 @@ export interface IWorkOrderCompleteStats { masterThinkingTokens: number; toolCallCount: number; durationMs: number; + /** Aggregate (master + subagent) totals for session-level accounting */ + totalInputTokens?: number; + totalOutputTokens?: number; + totalToolCallCount?: number; } export interface ChatTurnPrompts { diff --git a/gadget-code/frontend/src/pages/ChatSessionView.tsx b/gadget-code/frontend/src/pages/ChatSessionView.tsx index bf03f8a..af9c28c 100644 --- a/gadget-code/frontend/src/pages/ChatSessionView.tsx +++ b/gadget-code/frontend/src/pages/ChatSessionView.tsx @@ -756,13 +756,15 @@ export default function ChatSessionView() { })); // Update session contextWindowUsage for fuel gauge + // Uses aggregate totals (master+subagent) when available, falls back to master-only 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, + inputTokens: prev.stats.inputTokens + (stats.totalInputTokens ?? stats.masterInputTokens), + outputTokens: prev.stats.outputTokens + (stats.totalOutputTokens ?? stats.masterOutputTokens), + toolCallCount: prev.stats.toolCallCount + (stats.totalToolCallCount ?? stats.toolCallCount), }, } : prev); } diff --git a/gadget-code/src/lib/code-session.ts b/gadget-code/src/lib/code-session.ts index 0da3463..095f331 100644 --- a/gadget-code/src/lib/code-session.ts +++ b/gadget-code/src/lib/code-session.ts @@ -675,10 +675,12 @@ export class CodeSession extends SocketSession { this.socket.emit("workOrderComplete", turnId, success, message, stats); // Update in-memory session stats for UI reactivity + // Uses aggregate totals (master + subagent) when available, falls back to master-only if (stats && this.chatSession) { - this.chatSession.stats.inputTokens += stats.masterInputTokens; - this.chatSession.stats.outputTokens += stats.masterOutputTokens; - this.chatSession.stats.toolCallCount += stats.toolCallCount; + this.chatSession.stats.inputTokens += stats.totalInputTokens ?? stats.masterInputTokens; + this.chatSession.stats.outputTokens += stats.totalOutputTokens ?? stats.masterOutputTokens; + this.chatSession.stats.toolCallCount += stats.totalToolCallCount ?? stats.toolCallCount; + // contextWindowUsage is intentionally master-only (fuel gauge tracks master context) this.chatSession.contextWindowUsage = (this.chatSession.contextWindowUsage ?? 0) + stats.masterInputTokens; } } diff --git a/gadget-code/src/lib/drone-session.ts b/gadget-code/src/lib/drone-session.ts index 5b7801e..f27daa0 100644 --- a/gadget-code/src/lib/drone-session.ts +++ b/gadget-code/src/lib/drone-session.ts @@ -306,6 +306,8 @@ export class DroneSession extends SocketSession { await this.flushBuffer(turnId); this.streamingBuffers.delete(turnId); + let enrichedStats: IWorkOrderCompleteStats | undefined; + const turn = await ChatTurn.findById(turnId); if (turn) { if (success && message === "aborted") { @@ -333,22 +335,30 @@ export class DroneSession extends SocketSession { // COMPUTE AGGREGATE: walk the turn's subagent records for aggregate totals let totalInputTokens = stats?.masterInputTokens ?? 0; let totalOutputTokens = stats?.masterOutputTokens ?? 0; + let totalToolCallCount = stats?.toolCallCount ?? 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; + const subStats = (tc as any).subagent.stats; + totalInputTokens += subStats.inputTokens; + totalOutputTokens += subStats.responseTokens + subStats.thinkingTokenCount; + totalToolCallCount += subStats.toolCallCount; } } } + // Build enriched stats with aggregate values for in-memory consumers + enrichedStats = stats + ? { ...stats, totalInputTokens, totalOutputTokens, totalToolCallCount } + : undefined; + // 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, + "stats.toolCallCount": totalToolCallCount, // Master-only field for the fuel gauge "contextWindowUsage": stats?.masterInputTokens ?? 0, }}, @@ -358,7 +368,7 @@ export class DroneSession extends SocketSession { const codeSession = SocketService.getCodeSessionByChatSessionId( this.chatSessionId, ); - codeSession.onWorkOrderComplete(turnId, success, message, stats); + codeSession.onWorkOrderComplete(turnId, success, message, enrichedStats); this.currentTurnId = undefined; } catch (error) { diff --git a/gadget-drone/src/services/agent.ts b/gadget-drone/src/services/agent.ts index f131cb6..b7fd7a0 100644 --- a/gadget-drone/src/services/agent.ts +++ b/gadget-drone/src/services/agent.ts @@ -30,6 +30,7 @@ import { ChatSessionMode, ChatTurnStatus, type IProject, + type IWorkOrderCompleteStats, } from "@gadget/api"; import AiService from "./ai.ts"; @@ -59,13 +60,6 @@ import { 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 { @@ -269,9 +263,9 @@ class AgentService extends GadgetService { } // Accumulate real token counts from the AI API response - masterInputTokens += response.stats.tokenCounts.input; - masterOutputTokens += response.stats.tokenCounts.response; - masterThinkingTokens += response.stats.tokenCounts.thinking; + masterInputTokens += response.stats?.tokenCounts?.input ?? 0; + masterOutputTokens += response.stats?.tokenCounts?.response ?? 0; + masterThinkingTokens += response.stats?.tokenCounts?.thinking ?? 0; // Process tool calls if present if (response.toolCalls && response.toolCalls.length > 0) { @@ -319,6 +313,7 @@ class AgentService extends GadgetService { agentId: toolCall.callId, response: responseText, subagent: subagentData, + stats: subagentData?.stats as Record | undefined, }); } else { socket.emit("toolCall", toolCall.callId, toolCall.function.name, toolCall.function.arguments, result); diff --git a/packages/api/src/messages/drone.ts b/packages/api/src/messages/drone.ts index b82f05f..95fbc50 100644 --- a/packages/api/src/messages/drone.ts +++ b/packages/api/src/messages/drone.ts @@ -47,6 +47,10 @@ export interface IWorkOrderCompleteStats { masterThinkingTokens: number; toolCallCount: number; durationMs: number; + /** Aggregate (master + subagent) totals for session-level accounting */ + totalInputTokens?: number; + totalOutputTokens?: number; + totalToolCallCount?: number; } export type WorkOrderCompleteMessage = ( diff --git a/repairs.md b/repairs.md new file mode 100644 index 0000000..464edb4 --- /dev/null +++ b/repairs.md @@ -0,0 +1,22 @@ +Summary of Changes (6 files, +36/-19 lines) +Root Cause +The token-economics commit (00969cf) correctly added subagent token aggregation to the DB write path (drone-session.ts), but the in-memory and frontend paths only used master-only values. This meant: +- +During a session: UI showed lower token counts than reality +- +After page refresh: numbers jumped up (DB values loaded) +- +Session toolCallCount missed subagent tool calls entirely +Fixes +1. packages/api/src/messages/drone.ts — Extended IWorkOrderCompleteStats with optional aggregate fields (totalInputTokens?, totalOutputTokens?, totalToolCallCount?) +2. gadget-code/src/lib/drone-session.ts — Now computes totalToolCallCount (master + subagent tool calls) alongside existing token aggregates, passes enrichedStats (with aggregate fields populated) to codeSession.onWorkOrderComplete() +3. gadget-code/src/lib/code-session.ts — In-memory session stats use stats.totalInputTokens ?? stats.masterInputTokens (aggregate when available, master fallback) +4. gadget-code/frontend/src/lib/api.ts — Frontend type matches the shared interface +5. gadget-code/frontend/src/pages/ChatSessionView.tsx — Frontend session stats update uses aggregate values; also now correctly increments toolCallCount (was previously missing) +6. gadget-drone/src/services/agent.ts — Three fixes: +- +Removed local IWorkOrderCompleteStats redefinition, now imports from @gadget/api +- +Master loop token accumulation uses optional chaining (response.stats?.tokenCounts?.input ?? 0) matching the subagent loop +- +agent:complete emit now includes stats field from subagent data