diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts index 3e567b09c7..23355aee71 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts +++ b/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); diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts index 4c4035a53b..8b1ce61b89 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts +++ b/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 , 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 { @@ -850,7 +875,7 @@ export class CodexAppServerAgent extends BaseAcpAgent { } /** Start one codex turn and await its completion. */ - private async runTurn(input: CodexUserInput[]): Promise { + private async runTurn(input: CodexUserInput[]): Promise { 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 { - let reason = stopReason; + response: PromptResponse, + ): Promise { + 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 { + private async runFollowUpTurn(text: string): Promise { 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 { if (!this.sessionId || !this.taskRunId) return; await this.client diff --git a/packages/agent/src/adapters/codex-app-server/ext-notifications.ts b/packages/agent/src/adapters/codex-app-server/ext-notifications.ts index a692eaa421..f397285f21 100644 --- a/packages/agent/src/adapters/codex-app-server/ext-notifications.ts +++ b/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, }, }; } diff --git a/packages/agent/src/adapters/codex-app-server/turn-controller.test.ts b/packages/agent/src/adapters/codex-app-server/turn-controller.test.ts index 9e7398ffe0..e8d1f6ab8f 100644 --- a/packages/agent/src/adapters/codex-app-server/turn-controller.test.ts +++ b/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); diff --git a/packages/agent/src/adapters/codex-app-server/turn-controller.ts b/packages/agent/src/adapters/codex-app-server/turn-controller.ts index 350ba0b1ba..e6f6483a9a 100644 --- a/packages/agent/src/adapters/codex-app-server/turn-controller.ts +++ b/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; + private completion?: Promise; private generation = 0; private readonly cancelled = new Set(); - begin(): { completion: Promise; turn: number } { + begin(): { completion: Promise; turn: number } { const turn = ++this.generation; - this.completion = new Promise((resolve, reject) => { + this.completion = new Promise((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 { - return this.completion ?? Promise.resolve("end_turn"); + awaitCompletion(): Promise { + 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(); diff --git a/packages/agent/src/adapters/codex-app-server/usage-tracker.test.ts b/packages/agent/src/adapters/codex-app-server/usage-tracker.test.ts index aba636c727..d8c579bd18 100644 --- a/packages/agent/src/adapters/codex-app-server/usage-tracker.test.ts +++ b/packages/agent/src/adapters/codex-app-server/usage-tracker.test.ts @@ -46,6 +46,8 @@ describe("UsageTracker", () => { outputTokens: 80, cachedReadTokens: 50, cachedWriteTokens: 0, + thoughtTokens: 20, + totalTokens: 500, }); }); @@ -54,7 +56,7 @@ describe("UsageTracker", () => { const update = tracker.ingest(payload({ last: undefined })); expect(update?.used).toBe(1200); - expect(tracker.perTurnUsage().inputTokens).toBe(900); + expect(tracker.perTurnUsage()?.inputTokens).toBe(900); }); it("derives `used` from inputTokens when totalTokens is absent (same order as the gauge)", () => { @@ -78,17 +80,12 @@ describe("UsageTracker", () => { expect(tracker.contextTokens()).toBe(500); }); - it("resetForTurn zeroes the per-turn view so a token-less turn reports 0, not stale data", () => { + it("resetForTurn clears stale per-turn usage", () => { const tracker = new UsageTracker(); tracker.ingest(payload()); tracker.resetForTurn(); expect(tracker.contextTokens()).toBeUndefined(); - expect(tracker.perTurnUsage()).toEqual({ - inputTokens: 0, - outputTokens: 0, - cachedReadTokens: 0, - cachedWriteTokens: 0, - }); + expect(tracker.perTurnUsage()).toBeUndefined(); }); }); diff --git a/packages/agent/src/adapters/codex-app-server/usage-tracker.ts b/packages/agent/src/adapters/codex-app-server/usage-tracker.ts index 154d7d1159..fd107a7c13 100644 --- a/packages/agent/src/adapters/codex-app-server/usage-tracker.ts +++ b/packages/agent/src/adapters/codex-app-server/usage-tracker.ts @@ -1,8 +1,8 @@ +import type { Usage } from "@agentclientprotocol/sdk"; import { type ContextBreakdownBaseline, emptyBaseline, } from "../claude/context-breakdown"; -import type { AccumulatedUsage } from "./ext-notifications"; import { readTokenUsage } from "./token-usage"; /** The live `_posthog/usage_update` fields (context-window occupancy). */ @@ -25,7 +25,7 @@ export interface UsageUpdate { */ export class UsageTracker { private baseline: ContextBreakdownBaseline = emptyBaseline(); - private lastTurn?: AccumulatedUsage; + private lastTurn?: Usage; private contextUsed?: number; setBaseline(baseline: ContextBreakdownBaseline): void { @@ -36,7 +36,6 @@ export class UsageTracker { return this.baseline; } - /** Zero the per-turn view at turn start so a token-less turn reports 0. */ resetForTurn(): void { this.lastTurn = undefined; this.contextUsed = undefined; @@ -49,12 +48,17 @@ export class UsageTracker { const { context, used, size } = reading; // Drives the per-source breakdown's "conversation" bucket on turn complete. this.contextUsed = used; + const inputTokens = context.inputTokens ?? 0; + const outputTokens = context.outputTokens ?? 0; + const cachedReadTokens = context.cachedInputTokens ?? 0; this.lastTurn = { - inputTokens: context.inputTokens ?? 0, - outputTokens: context.outputTokens ?? 0, - cachedReadTokens: context.cachedInputTokens ?? 0, - // codex's TokenUsageBreakdown has no cache-write field; 0 is authoritative. + inputTokens, + outputTokens, + cachedReadTokens, cachedWriteTokens: 0, + thoughtTokens: context.reasoningOutputTokens, + totalTokens: + context.totalTokens ?? inputTokens + outputTokens + cachedReadTokens, }; return { used, @@ -69,16 +73,8 @@ export class UsageTracker { }; } - /** Per-turn usage for `_posthog/turn_complete` — codex's `last`, not a delta. */ - perTurnUsage(): AccumulatedUsage { - return ( - this.lastTurn ?? { - inputTokens: 0, - outputTokens: 0, - cachedReadTokens: 0, - cachedWriteTokens: 0, - } - ); + perTurnUsage(): Usage | undefined { + return this.lastTurn ? { ...this.lastTurn } : undefined; } /** Live context occupancy (same derivation as the renderer gauge), or undefined pre-usage. */