fix(agent): persist Codex prompt token usage (#3506)

Author: tatoalo

packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts

@@ -1961,9 +1961,23 @@ describe("CodexAppServerAgent", () => {
     } as unknown as PromptRequest);
 
     // The single turn/completed resolves both the original and the folded prompt.
+    stub.emit("thread/tokenUsage/updated", {
+      tokenUsage: {
+        last: {
+          totalTokens: 45,
+          inputTokens: 30,
+          cachedInputTokens: 5,
+          outputTokens: 10,
+        },
+      },
+    });
     stub.emit("turn/completed", { turn: { status: "completed" } });
-    expect((await first).stopReason).toBe("end_turn");
-    expect((await second).stopReason).toBe("end_turn");
+    const [firstResult, secondResult] = await Promise.all([first, second]);
+    expect(firstResult).toMatchObject({
+      stopReason: "end_turn",
+      usage: { totalTokens: 45 },
+    });
+    expect(secondResult).toEqual({ stopReason: "end_turn" });
 
     const steer = stub.requests.find((r) => r.method === "turn/steer");
     expect(steer?.params).toMatchObject({
@@ -2126,7 +2140,19 @@ describe("CodexAppServerAgent", () => {
       },
     });
     stub.emit("turn/completed", { turn: { status: "completed" } });
-    await done;
+    const result = await done;
+
+    expect(result).toEqual({
+      stopReason: "end_turn",
+      usage: {
+        inputTokens: 60,
+        outputTokens: 30,
+        cachedReadTokens: 10,
+        cachedWriteTokens: 0,
+        thoughtTokens: 5,
+        totalTokens: 100,
+      },
+    });
 
     const turnComplete = extNotifications.find(
       (n) => n.method === "_posthog/turn_complete",
@@ -2818,6 +2844,16 @@ describe("CodexAppServerAgent", () => {
         text: "The implementation plan is ready.",
       },
     });
+    stub.emit("thread/tokenUsage/updated", {
+      tokenUsage: {
+        last: {
+          totalTokens: 30,
+          inputTokens: 20,
+          outputTokens: 10,
+          reasoningOutputTokens: 2,
+        },
+      },
+    });
     stub.emit("turn/completed", {
       turn: { id: "turn_1", status: "completed" },
     });
@@ -2826,10 +2862,31 @@ describe("CodexAppServerAgent", () => {
     await waitUntil(
       () => stub.requests.filter((r) => r.method === "turn/start").length >= 2,
     );
+    stub.emit("thread/tokenUsage/updated", {
+      tokenUsage: {
+        last: {
+          totalTokens: 50,
+          inputTokens: 35,
+          cachedInputTokens: 5,
+          outputTokens: 10,
+          reasoningOutputTokens: 3,
+        },
+      },
+    });
     stub.emit("turn/completed", {
       turn: { id: "turn_2", status: "completed" },
     });
-    expect((await done).stopReason).toBe("end_turn");
+    expect(await done).toEqual({
+      stopReason: "end_turn",
+      usage: {
+        inputTokens: 55,
+        outputTokens: 20,
+        cachedReadTokens: 5,
+        cachedWriteTokens: 0,
+        thoughtTokens: 5,
+        totalTokens: 80,
+      },
+    });
 
     // The approval renders as the plan-approval UI (switch_mode + the plan text).
     expect(permissionRequests).toHaveLength(1);

packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts

@@ -52,7 +52,6 @@ import {
 } from "./app-server-client";
 import { handleServerRequest } from "./approvals";
 import {
-  type AccumulatedUsage,
   buildSdkSessionParams,
   buildTurnCompleteParams,
   buildUsageBreakdownParams,
@@ -162,6 +161,31 @@ function parseGoalCommand(prompt: PromptRequest["prompt"]): GoalCommand | null {
   }
 }
 
+function mergePromptUsage(
+  left: PromptResponse["usage"],
+  right: PromptResponse["usage"],
+): PromptResponse["usage"] {
+  if (!left) return right;
+  if (!right) return left;
+  return {
+    inputTokens: left.inputTokens + right.inputTokens,
+    outputTokens: left.outputTokens + right.outputTokens,
+    cachedReadTokens:
+      (left.cachedReadTokens ?? 0) + (right.cachedReadTokens ?? 0),
+    cachedWriteTokens:
+      (left.cachedWriteTokens ?? 0) + (right.cachedWriteTokens ?? 0),
+    thoughtTokens: (left.thoughtTokens ?? 0) + (right.thoughtTokens ?? 0),
+    totalTokens: left.totalTokens + right.totalTokens,
+  };
+}
+
+function mergePromptResponses(
+  left: PromptResponse,
+  right: PromptResponse,
+): PromptResponse {
+  return { ...right, usage: mergePromptUsage(left.usage, right.usage) };
+}
+
 // The native app-server owns its config; BaseAcpAgent only calls dispose() on this.
 class NoopSettingsManager implements BaseSettingsManager {
   constructor(private cwd: string) {}
@@ -219,7 +243,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
   /** The in-flight turn's <proposed_plan>, streamed or completed (drives the implement handoff). */
   private planProposal?: { itemId: string; text: string };
   /** Idle signal deferred while the plan handoff keeps this prompt busy. */
-  private deferredTurnComplete?: { usage: AccumulatedUsage };
+  private deferredTurnComplete?: { usage: PromptResponse["usage"] };
   /** Settles the pending plan-approval race on cancel/close/preempting prompt. */
   private planHandoffCancel?: () => void;
   private readonly mcp = new McpManager();
@@ -727,15 +751,16 @@ export class CodexAppServerAgent extends BaseAcpAgent {
           return undefined;
         });
       this.turns.onSteered(steerRes?.turnId);
-      return { stopReason: await this.turns.awaitCompletion() };
+      const response = await this.turns.awaitCompletion();
+      return { stopReason: response.stopReason };
     }
     if (this.turns.isPending) {
       // A turn is pending but has no turnId yet, so we can't steer; fail fast.
       throw new Error("prompt() called while a turn is already in progress");
     }
 
-    const stopReason = await this.runTurn(input);
-    return { stopReason: await this.maybeOfferPlanImplementation(stopReason) };
+    const response = await this.runTurn(input);
+    return this.maybeOfferPlanImplementation(response);
   }
 
   private async handleGoalCommand(command: GoalCommand): Promise<void> {
@@ -850,7 +875,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
   }
 
   /** Start one codex turn and await its completion. */
-  private async runTurn(input: CodexUserInput[]): Promise<StopReason> {
+  private async runTurn(input: CodexUserInput[]): Promise<PromptResponse> {
     this.lastAgentMessage = "";
     this.resetUsage();
     this.planProposal = undefined;
@@ -895,12 +920,12 @@ export class CodexAppServerAgent extends BaseAcpAgent {
    * back into another plan turn, whose revised plan prompts again.
    */
   private async maybeOfferPlanImplementation(
-    stopReason: StopReason,
-  ): Promise<StopReason> {
-    let reason = stopReason;
+    response: PromptResponse,
+  ): Promise<PromptResponse> {
+    let result = response;
     try {
       while (
-        reason === "end_turn" &&
+        result.stopReason === "end_turn" &&
         this.config.mode === "plan" &&
         this.planProposal &&
         !this.session.cancelled
@@ -911,7 +936,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
         // Re-check after the await: a cancel that raced the response wins, so a
         // late accept can never start implementation on a cancelled prompt.
         if (this.session.cancelled) {
-          reason = "cancelled";
+          result = { ...result, stopReason: "cancelled" };
           break;
         }
         // A picker change while approval was open owns the mode. Never let a
@@ -921,19 +946,25 @@ export class CodexAppServerAgent extends BaseAcpAgent {
           this.config.setOption("mode", outcome.mode);
           this.emitCurrentMode(outcome.mode);
           this.emitConfigOptions();
-          reason = await this.runFollowUpTurn(IMPLEMENT_PLAN_MESSAGE);
+          result = mergePromptResponses(
+            result,
+            await this.runFollowUpTurn(IMPLEMENT_PLAN_MESSAGE),
+          );
           break;
         }
         if (outcome.kind === "feedback") {
-          reason = await this.runFollowUpTurn(outcome.feedback);
+          result = mergePromptResponses(
+            result,
+            await this.runFollowUpTurn(outcome.feedback),
+          );
           continue;
         }
         break;
       }
     } finally {
-      await this.flushDeferredTurnComplete(reason);
+      await this.flushDeferredTurnComplete(result.stopReason);
     }
-    return reason;
+    return result;
   }
 
   /**
@@ -949,7 +980,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
   }
 
   /** Run an adapter-initiated turn, echoed as a user message like a host prompt. */
-  private async runFollowUpTurn(text: string): Promise<StopReason> {
+  private async runFollowUpTurn(text: string): Promise<PromptResponse> {
     this.broadcastUserInput([{ type: "text", text }]);
     return this.runTurn(toCodexInput([{ type: "text", text }]));
   }
@@ -1448,7 +1479,10 @@ export class CodexAppServerAgent extends BaseAcpAgent {
       await this.emitTurnCompleteSignal(reason, usage);
       await this.emitUsageBreakdown(contextUsed);
     }
-    pending.resolve(reason);
+    pending.resolve({
+      stopReason: reason,
+      ...(usage ? { usage } : {}),
+    });
   }
 
   /** Whether maybeOfferPlanImplementation will run for a turn that ended this way. */
@@ -1464,7 +1498,7 @@ export class CodexAppServerAgent extends BaseAcpAgent {
   /** Emit the cloud idle signal `_posthog/turn_complete` (only with a taskRunId). */
   private async emitTurnCompleteSignal(
     reason: StopReason,
-    usage: AccumulatedUsage,
+    usage: PromptResponse["usage"],
   ): Promise<void> {
     if (!this.sessionId || !this.taskRunId) return;
     await this.client

packages/agent/src/adapters/codex-app-server/ext-notifications.ts

@@ -50,13 +50,12 @@ export interface TurnCompleteParams {
   usage: TurnCompleteUsage;
 }
 
-/** The four component counts the caller accumulates; total is computed here. */
-export interface AccumulatedUsage {
-  inputTokens: number;
-  outputTokens: number;
-  cachedReadTokens: number;
-  cachedWriteTokens: number;
-}
+type TurnCompleteUsageInput = {
+  inputTokens?: number | null;
+  outputTokens?: number | null;
+  cachedReadTokens?: number | null;
+  cachedWriteTokens?: number | null;
+};
 
 /**
  * `_posthog/turn_complete` — fired when a prompt turn finishes. `totalTokens` is the
@@ -65,21 +64,22 @@ export interface AccumulatedUsage {
 export function buildTurnCompleteParams(
   sessionId: string,
   stopReason: StopReason,
-  usage: AccumulatedUsage,
+  usage?: TurnCompleteUsageInput | null,
 ): TurnCompleteParams {
+  const inputTokens = usage?.inputTokens ?? 0;
+  const outputTokens = usage?.outputTokens ?? 0;
+  const cachedReadTokens = usage?.cachedReadTokens ?? 0;
+  const cachedWriteTokens = usage?.cachedWriteTokens ?? 0;
   return {
     sessionId,
     stopReason,
     usage: {
-      inputTokens: usage.inputTokens,
-      outputTokens: usage.outputTokens,
-      cachedReadTokens: usage.cachedReadTokens,
-      cachedWriteTokens: usage.cachedWriteTokens,
+      inputTokens,
+      outputTokens,
+      cachedReadTokens,
+      cachedWriteTokens,
       totalTokens:
-        usage.inputTokens +
-        usage.outputTokens +
-        usage.cachedReadTokens +
-        usage.cachedWriteTokens,
+        inputTokens + outputTokens + cachedReadTokens + cachedWriteTokens,
     },
   };
 }

packages/agent/src/adapters/codex-app-server/turn-controller.test.ts

@@ -18,8 +18,8 @@ describe("TurnController", () => {
     expect(turns.activeTurnId).toBeUndefined();
     expect(turns.claim()).toBeUndefined();
 
-    pending?.resolve("end_turn");
-    await expect(completion).resolves.toBe("end_turn");
+    pending?.resolve({ stopReason: "end_turn" });
+    await expect(completion).resolves.toEqual({ stopReason: "end_turn" });
   });
 
   it("finishPrompt for an older turn does not wipe a newer turn's pending state", async () => {
@@ -36,15 +36,15 @@ describe("TurnController", () => {
     turns.onStarted("turn-b");
 
     // Turn A's prompt() resolves and its finally runs; it must not clear turn B.
-    claimedA?.resolve("end_turn");
-    await expect(a.completion).resolves.toBe("end_turn");
+    claimedA?.resolve({ stopReason: "end_turn" });
+    await expect(a.completion).resolves.toEqual({ stopReason: "end_turn" });
     turns.finishPrompt(a.turn);
 
     expect(turns.isRunning).toBe(true);
     const claimedB = turns.claim();
     expect(claimedB).toBeDefined();
-    claimedB?.resolve("end_turn");
-    await expect(b.completion).resolves.toBe("end_turn");
+    claimedB?.resolve({ stopReason: "end_turn" });
+    await expect(b.completion).resolves.toEqual({ stopReason: "end_turn" });
   });
 
   it("finishPrompt with the current turn token clears the pending slot", () => {
@@ -84,7 +84,7 @@ describe("TurnController", () => {
     turns.markInterrupted();
 
     turns.close("cancelled");
-    await expect(completion).resolves.toBe("cancelled");
+    await expect(completion).resolves.toEqual({ stopReason: "cancelled" });
     expect(turns.isPending).toBe(false);
     expect(turns.activeTurnId).toBeUndefined();
     expect(turns.shouldDropCompletion("turn-1")).toBe(false);

packages/agent/src/adapters/codex-app-server/turn-controller.ts

@@ -1,7 +1,7 @@
-import type { StopReason } from "@agentclientprotocol/sdk";
+import type { PromptResponse, StopReason } from "@agentclientprotocol/sdk";
 
 interface PendingTurn {
-  resolve: (reason: StopReason) => void;
+  resolve: (response: PromptResponse) => void;
   reject: (err: Error) => void;
 }
 
@@ -13,13 +13,13 @@ interface PendingTurn {
 export class TurnController {
   private turnId?: string;
   private pending?: PendingTurn;
-  private completion?: Promise<StopReason>;
+  private completion?: Promise<PromptResponse>;
   private generation = 0;
   private readonly cancelled = new Set<string>();
 
-  begin(): { completion: Promise<StopReason>; turn: number } {
+  begin(): { completion: Promise<PromptResponse>; turn: number } {
     const turn = ++this.generation;
-    this.completion = new Promise<StopReason>((resolve, reject) => {
+    this.completion = new Promise<PromptResponse>((resolve, reject) => {
       this.pending = { resolve, reject };
     });
     return { completion: this.completion, turn };
@@ -49,8 +49,8 @@ export class TurnController {
   }
 
   /** Await the in-flight turn's completion (the steer path reuses the original). */
-  awaitCompletion(): Promise<StopReason> {
-    return this.completion ?? Promise.resolve("end_turn");
+  awaitCompletion(): Promise<PromptResponse> {
+    return this.completion ?? Promise.resolve({ stopReason: "end_turn" });
   }
 
   /** Atomically claim the pending turn (clears the slot + turnId synchronously), or undefined if already claimed. */
@@ -96,7 +96,7 @@ export class TurnController {
   /** Resolve and clear everything on session close. */
   close(reason: StopReason): void {
     this.turnId = undefined;
-    this.pending?.resolve(reason);
+    this.pending?.resolve({ stopReason: reason });
     this.pending = undefined;
     this.completion = undefined;
     this.cancelled.clear();