repairs for OpenAI token stats/economics

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 {

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,
+          inputTokens: prev.stats.inputTokens + (stats.totalInputTokens ?? stats.masterInputTokens),
-          outputTokens: prev.stats.outputTokens + stats.masterOutputTokens,
+          outputTokens: prev.stats.outputTokens + (stats.totalOutputTokens ?? stats.masterOutputTokens),
+          toolCallCount: prev.stats.toolCallCount + (stats.totalToolCallCount ?? stats.toolCallCount),
         },
       } : prev);
     }

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.inputTokens += stats.totalInputTokens ?? stats.masterInputTokens;
-      this.chatSession.stats.outputTokens += stats.masterOutputTokens;
+      this.chatSession.stats.outputTokens += stats.totalOutputTokens ?? stats.masterOutputTokens;
-      this.chatSession.stats.toolCallCount += stats.toolCallCount;
+      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;
     }
   }

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;
+              const subStats = (tc as any).subagent.stats;
-              totalOutputTokens += (tc as any).subagent.stats.responseTokens + (tc as any).subagent.stats.thinkingTokenCount;
+              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) {

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;
+        masterInputTokens += response.stats?.tokenCounts?.input ?? 0;
-        masterOutputTokens += response.stats.tokenCounts.response;
+        masterOutputTokens += response.stats?.tokenCounts?.response ?? 0;
-        masterThinkingTokens += response.stats.tokenCounts.thinking;
+        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<string, unknown> | undefined,
               });
             } else {
               socket.emit("toolCall", toolCall.callId, toolCall.function.name, toolCall.function.arguments, result);

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 = (

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