diff --git a/README.md b/README.md index 88a2ad51..62d0b17a 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ Use [OpenAI Codex](https://github.com/openai/codex) from [Agent Client Protocol] - ChatGPT, API key, and client-provided custom gateway authentication. - Model, reasoning effort, fast mode, approval, and sandbox mode configuration. - Text prompts, embedded context, images, resource links, and additional workspace directories. -- Shell command, file change, [permission request](docs/permission-extension.md), MCP tool call, terminal output, reasoning, plan, web search, image generation, image view, token usage, and review events. +- Shell command, file change, [permission request](docs/permission-extension.md), MCP tool call, terminal output, reasoning, plan, web search, image generation, image view, [token usage](docs/usage-accounting.md), and review events. - [Native ACP subagent sessions](docs/subagent-sessions.md) (after capability negotiation) with separate child histories and root-routed permissions; a legacy tool-call fallback otherwise. - [Background terminal tasks](docs/async-tasks.md) in AIR, with task status and targeted stop support after capability negotiation. - Session-scoped long-running goals through the provider-neutral [goal extension](docs/goal-extension.md). diff --git a/docs/usage-accounting.md b/docs/usage-accounting.md new file mode 100644 index 00000000..b5e101c4 --- /dev/null +++ b/docs/usage-accounting.md @@ -0,0 +1,36 @@ +# Prompt usage + +`PromptResponse.usage` and `_meta.quota.token_count` report observed usage for +the root Codex thread during this ACP prompt. They include all model requests +in the prompt, including a plan and its approved implementation. They exclude +earlier prompts, native child threads, and background title generation. + +The adapter subtracts the thread total captured before the prompt from each +subsequent total. Repeated notifications add zero. New threads start at zero; +load/resume snapshots supply the historical baseline. A missing baseline is +unknown: the first snapshot establishes it without charging historical usage. +Later observations are retained, but the result is partial. + +`_meta.usageAccounting` identifies the contract: + +```json +{ + "version": 1, + "source": "codex/thread-token-usage-delta", + "scope": "root_thread_prompt", + "completeness": "reported" +} +``` + +`reported` means the observed counters were usable. It does not certify provider +billing or include work outside the stated scope. `partial` means a baseline or +usage was missing, a counter was replaced or invalid, or the prompt was cancelled +or returned a typed failure. Known counts remain available; no observations +produce `usage: null`. A transport error without a prompt response still has no +terminal usage result. + +Cached reads are separated from input. Cache writes remain included in non-read +input, as before; reasoning is a subset of output. Neither is added twice. +`session/update.usage_update.used` remains context occupancy and must not be +summed as token consumption. Reports from adapter versions before this change +describe only the last model request and cannot be repaired retroactively. diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 7040057f..011a366b 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -25,6 +25,7 @@ import type { import type {ServiceTier} from "./app-server/ServiceTier"; import type {JsonValue} from "./app-server/serde_json/JsonValue"; import {ModelId} from "./ModelId"; +import {toTokenCount, type TokenCount} from "./TokenCount"; import {AgentMode} from "./AgentMode"; import path from "node:path"; import {logger} from "./Logger"; @@ -126,6 +127,7 @@ export class CodexAcpClient { private readonly subagents: CodexSubagentSubscriptions; private skillExtraRoots: string[] = []; private configPath: string | null = null; + private readonly threadTokenUsage = new Map(); constructor(codexClient: CodexAppServerClient, codexConfig?: JsonObject, modelProvider?: string) { @@ -135,6 +137,16 @@ export class CodexAcpClient { this.gatewayConfig = null; this.gatewayConfigSource = null; this.subagents = new CodexSubagentSubscriptions(codexClient); + // Capture restored totals even before the ACP session handler is installed. + codexClient.onClientTransportEvent(event => { + if (event.eventType === "notification" && event.method === "thread/tokenUsage/updated") { + this.threadTokenUsage.set(event.params.threadId, toTokenCount(event.params.tokenUsage.total)); + } + }); + } + + getThreadTokenUsage(sessionId: string): TokenCount | null { + return this.threadTokenUsage.get(sessionId) ?? null; } get appServerClient(): CodexAppServerClient { @@ -638,7 +650,11 @@ export class CodexAcpClient { await this.codexClient.threadUnsubscribe({threadId: sessionId}); } finally { this.codexClient.clearThreadHandlers(sessionId); + for (const childSessionId of this.subagents.childSessionIds(sessionId)) { + this.threadTokenUsage.delete(childSessionId); + } this.subagents.clear(sessionId); + this.threadTokenUsage.delete(sessionId); } } diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index 4dc15e01..27bb1a6e 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -52,10 +52,10 @@ import { REASONING_EFFORT_CONFIG_ID, } from "./ModelConfigOption"; import type {TokenCount} from "./TokenCount"; +import {PromptTokenUsage, ZERO_TOKEN_COUNT} from "./PromptTokenUsage"; import {toPromptUsage} from "./TokenCount"; import {CodexCommands, GOAL_CONTINUATION_PROMPT} from "./CodexCommands"; import {SteeringQueue} from "./SteeringQueue"; -import type {QuotaMeta} from "./QuotaMeta"; import {logger} from "./Logger"; import {sanitizeMcpServerName} from "./McpServerName"; import {createResponseItemHistoryFallbackUpdates} from "./ResponseItemHistoryFallback"; @@ -162,6 +162,7 @@ export interface SessionState { currentTurnId: string | null; lastTokenUsage: TokenCount | null; totalTokenUsage: TokenCount | null; + promptTokenUsage?: PromptTokenUsage; modelContextWindow: number | null; rateLimits: RateLimitsMap | null; account: Account | null; @@ -672,7 +673,7 @@ export class CodexAcpServer { collaborationMode: sessionMetadata.collaborationMode, currentTurnId: null, lastTokenUsage: null, - totalTokenUsage: null, + totalTokenUsage: operation === "new" ? ZERO_TOKEN_COUNT : null, modelContextWindow: null, rateLimits: null, account: authState.account, @@ -2751,6 +2752,9 @@ export class CodexAcpServer { let agentFileChangeReportUnavailableReason: AgentFileChangeReportUnavailableReason = "providerError"; let promptWasCancelled = false; let recoverableSessionFailure = sessionState.sessionFailure; + sessionState.promptTokenUsage = new PromptTokenUsage( + this.codexAcpClient.getThreadTokenUsage(params.sessionId) ?? sessionState.totalTokenUsage, + ); sessionState.currentTurnId = null; const activePrompt = this.trackActivePrompt(params.sessionId); let pendingTurnStart: PendingTurnStart | null = null; @@ -2911,7 +2915,7 @@ export class CodexAcpServer { await clearRecoveredSessionFailure(eventHandler); return { stopReason: "end_turn", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), + usage: this.buildPromptUsage(sessionState.promptTokenUsage?.tokenCount() ?? null), _meta: this.buildQuotaMeta(sessionState), }; } @@ -3145,7 +3149,7 @@ export class CodexAcpServer { return { stopReason: "end_turn", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), + usage: this.buildPromptUsage(sessionState.promptTokenUsage?.tokenCount() ?? null), _meta: this.buildQuotaMeta(sessionState), }; } catch (err) { @@ -3246,8 +3250,8 @@ export class CodexAcpServer { private cancelledPromptResponse(sessionState: SessionState): acp.PromptResponse { return { stopReason: "cancelled", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), - _meta: this.buildQuotaMeta(sessionState), + usage: this.buildPromptUsage(sessionState.promptTokenUsage?.tokenCount() ?? null), + _meta: this.buildQuotaMeta(sessionState, true), }; } @@ -3263,16 +3267,16 @@ export class CodexAcpServer { } return { stopReason: "end_turn", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), + usage: this.buildPromptUsage(sessionState.promptTokenUsage?.tokenCount() ?? null), _meta: { - ...this.buildQuotaMeta(sessionState), + ...this.buildQuotaMeta(sessionState, true), ...failureMeta, }, }; } - private buildQuotaMeta(sessionState: SessionState): { quota: QuotaMeta } { - const lastTokenUsage = sessionState.lastTokenUsage; + private buildQuotaMeta(sessionState: SessionState, interrupted = false) { + const lastTokenUsage = sessionState.promptTokenUsage?.tokenCount() ?? null; // Remove the "[reasoning-level]" suffix from currentModelId if present const modelName = sessionState.currentModelId.replace(/\[.*?]$/, ''); @@ -3284,9 +3288,10 @@ export class CodexAcpServer { return { quota: { - token_count: sessionState.lastTokenUsage, + token_count: lastTokenUsage, model_usage: modelUsage - } + }, + usageAccounting: sessionState.promptTokenUsage?.accounting(interrupted) }; } diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts index 7b567541..fc6476cc 100644 --- a/src/CodexEventHandler.ts +++ b/src/CodexEventHandler.ts @@ -1280,6 +1280,12 @@ export class CodexEventHandler { } private handleTokenUsageUpdated(params: ThreadTokenUsageUpdatedNotification): void { + if (params.threadId === this.sessionState.sessionId + && params.turnId === this.sessionState.currentTurnId) { + this.sessionState.promptTokenUsage?.observe(params.tokenUsage); + } else if (params.threadId === this.sessionState.sessionId && this.sessionState.currentTurnId === null) { + this.sessionState.promptTokenUsage?.restoreBaseline(toTokenCount(params.tokenUsage.total)); + } this.sessionState.lastTokenUsage = toTokenCount(params.tokenUsage.last); this.sessionState.totalTokenUsage = toTokenCount(params.tokenUsage.total); this.sessionState.modelContextWindow = params.tokenUsage.modelContextWindow; diff --git a/src/PromptTokenUsage.ts b/src/PromptTokenUsage.ts new file mode 100644 index 00000000..525025ec --- /dev/null +++ b/src/PromptTokenUsage.ts @@ -0,0 +1,63 @@ +import type {ThreadTokenUsage} from "./app-server/v2"; +import {toTokenCount, type TokenCount} from "./TokenCount"; + +const fields = ["totalTokens", "inputTokens", "cachedInputTokens", "outputTokens", "reasoningOutputTokens"] as const; +export const ZERO_TOKEN_COUNT: TokenCount = { + totalTokens: 0, inputTokens: 0, cachedInputTokens: 0, outputTokens: 0, reasoningOutputTokens: 0, +}; + +/** Usage observed during one ACP prompt. Context occupancy is not usage. */ +export class PromptTokenUsage { + private counts: TokenCount | null = null; + private incomplete = false; + private started = false; + + constructor(private previous: TokenCount | null) {} + + restoreBaseline(total: TokenCount): void { + // A resume snapshot can arrive after session/resume but before turn/start. + if (!this.started) this.previous = total; + } + + observe(usage: ThreadTokenUsage): void { + this.started = true; + const total = toTokenCount(usage.total); + const previous = this.previous; + this.previous = total; + if (previous === null) { + // A first notification can repeat historical usage on a rate-limit + // update. Without a baseline, even `last` cannot safely be charged. + this.incomplete = true; + return; + } + const delta = {...total}; + for (const field of fields) delta[field] -= previous[field]; + if (fields.some(field => !Number.isSafeInteger(delta[field]) || delta[field] < 0) + || delta.reasoningOutputTokens > delta.outputTokens + || delta.totalTokens !== delta.inputTokens + delta.cachedInputTokens + delta.outputTokens) { + // Codex can replace counters with a synthetic context-window estimate. + this.incomplete = true; + return; + } + const next = {...(this.counts ?? ZERO_TOKEN_COUNT)}; + for (const field of fields) next[field] += delta[field]; + if (fields.some(field => !Number.isSafeInteger(next[field]))) { + this.incomplete = true; + return; + } + this.counts = next; + } + + tokenCount(): TokenCount | null { + return this.counts; + } + + accounting(interrupted = false) { + return { + version: 1, + source: "codex/thread-token-usage-delta", + scope: "root_thread_prompt", + completeness: interrupted || this.incomplete || this.counts === null ? "partial" : "reported", + }; + } +} diff --git a/src/__tests__/CodexACPAgent/data/command-status-with-rate-limits.json b/src/__tests__/CodexACPAgent/data/command-status-with-rate-limits.json index 13403fd0..d70e1a48 100644 --- a/src/__tests__/CodexACPAgent/data/command-status-with-rate-limits.json +++ b/src/__tests__/CodexACPAgent/data/command-status-with-rate-limits.json @@ -7,7 +7,7 @@ "sessionUpdate": "agent_message_chunk", "content": { "type": "text", - "text": "**Model:** model-id[effort] \n**Directory:** /test/cwd \n**Approval:** on-request \n**Sandbox:** workspace-write \n**Account:** not logged in \n**Session:** `session-id` \n \n**Token usage:** data not available yet \n**Context window:** data not available yet \n**Standard 1h limit:** 75% left \n**Fast 1d limit:** 20% left" + "text": "**Model:** model-id[effort] \n**Directory:** /test/cwd \n**Approval:** on-request \n**Sandbox:** workspace-write \n**Account:** not logged in \n**Session:** `session-id` \n \n**Token usage:** 0 total (0 input + 0 cached input, 0 output) \n**Context window:** data not available yet \n**Standard 1h limit:** 75% left \n**Fast 1d limit:** 20% left" } } } diff --git a/src/__tests__/CodexACPAgent/data/command-status.json b/src/__tests__/CodexACPAgent/data/command-status.json index 50260582..7fabb21e 100644 --- a/src/__tests__/CodexACPAgent/data/command-status.json +++ b/src/__tests__/CodexACPAgent/data/command-status.json @@ -7,7 +7,7 @@ "sessionUpdate": "agent_message_chunk", "content": { "type": "text", - "text": "**Model:** model-id[effort] \n**Directory:** /test/cwd \n**Approval:** on-request \n**Sandbox:** workspace-write \n**Account:** not logged in \n**Session:** `session-id` \n \n**Token usage:** data not available yet \n**Context window:** data not available yet \n**Limits:** data not available yet" + "text": "**Model:** model-id[effort] \n**Directory:** /test/cwd \n**Approval:** on-request \n**Sandbox:** workspace-write \n**Account:** not logged in \n**Session:** `session-id` \n \n**Token usage:** 0 total (0 input + 0 cached input, 0 output) \n**Context window:** data not available yet \n**Limits:** data not available yet" } } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json b/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json index 426011ab..c6b47b8a 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json @@ -1,33 +1,39 @@ { "stopReason": "cancelled", "usage": { - "totalTokens": 1500, - "inputTokens": 1200, + "totalTokens": 3000, + "inputTokens": 2500, "cachedReadTokens": 0, - "outputTokens": 300, + "outputTokens": 500, "thoughtTokens": 0 }, "_meta": { "quota": { "token_count": { - "totalTokens": 1500, - "inputTokens": 1200, + "totalTokens": 3000, + "inputTokens": 2500, "cachedInputTokens": 0, - "outputTokens": 300, + "outputTokens": 500, "reasoningOutputTokens": 0 }, "model_usage": [ { "model": "model-id", "token_count": { - "totalTokens": 1500, - "inputTokens": 1200, + "totalTokens": 3000, + "inputTokens": 2500, "cachedInputTokens": 0, - "outputTokens": 300, + "outputTokens": 500, "reasoningOutputTokens": 0 } } ] + }, + "usageAccounting": { + "version": 1, + "source": "codex/thread-token-usage-delta", + "scope": "root_thread_prompt", + "completeness": "partial" } } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json b/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json index 60ba2e2e..2ec69031 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json @@ -1,33 +1,39 @@ { "stopReason": "end_turn", "usage": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedReadTokens": 500, - "outputTokens": 450, - "thoughtTokens": 50 + "totalTokens": 5000, + "inputTokens": 3000, + "cachedReadTokens": 1000, + "outputTokens": 1000, + "thoughtTokens": 100 }, "_meta": { "quota": { "token_count": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedInputTokens": 500, - "outputTokens": 450, - "reasoningOutputTokens": 50 + "totalTokens": 5000, + "inputTokens": 3000, + "cachedInputTokens": 1000, + "outputTokens": 1000, + "reasoningOutputTokens": 100 }, "model_usage": [ { "model": "model-id", "token_count": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedInputTokens": 500, - "outputTokens": 450, - "reasoningOutputTokens": 50 + "totalTokens": 5000, + "inputTokens": 3000, + "cachedInputTokens": 1000, + "outputTokens": 1000, + "reasoningOutputTokens": 100 } } ] + }, + "usageAccounting": { + "version": 1, + "source": "codex/thread-token-usage-delta", + "scope": "root_thread_prompt", + "completeness": "reported" } } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json b/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json index aebc1783..5442c67f 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json @@ -1,33 +1,39 @@ { "stopReason": "end_turn", "usage": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2300, "cachedReadTokens": 500, - "outputTokens": 200, + "outputTokens": 700, "thoughtTokens": 100 }, "_meta": { "quota": { "token_count": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2300, "cachedInputTokens": 500, - "outputTokens": 200, + "outputTokens": 700, "reasoningOutputTokens": 100 }, "model_usage": [ { "model": "model-id", "token_count": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2300, "cachedInputTokens": 500, - "outputTokens": 200, + "outputTokens": 700, "reasoningOutputTokens": 100 } } ] + }, + "usageAccounting": { + "version": 1, + "source": "codex/thread-token-usage-delta", + "scope": "root_thread_prompt", + "completeness": "reported" } } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-null.json b/src/__tests__/CodexACPAgent/data/token-usage-null.json index 65196087..b5ac3242 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-null.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-null.json @@ -5,6 +5,12 @@ "quota": { "token_count": null, "model_usage": [] + }, + "usageAccounting": { + "version": 1, + "source": "codex/thread-token-usage-delta", + "scope": "root_thread_prompt", + "completeness": "partial" } } } diff --git a/src/__tests__/CodexACPAgent/plan-review-events.test.ts b/src/__tests__/CodexACPAgent/plan-review-events.test.ts index 67dbfc73..768414c0 100644 --- a/src/__tests__/CodexACPAgent/plan-review-events.test.ts +++ b/src/__tests__/CodexACPAgent/plan-review-events.test.ts @@ -60,6 +60,7 @@ describe("CodexACPAgent - plan review", () => { permissionOptionId: string | null, options: { typedFailures?: boolean; + reportUsage?: boolean; emitCompletionNotification?: boolean; implementationStart?: Promise; permissionResponse?: acp.RequestPermissionResponse | Promise; @@ -144,6 +145,15 @@ describe("CodexACPAgent - plan review", () => { }, }, }); + if (options.reportUsage) { + fixture.sendServerNotification({method: "thread/tokenUsage/updated", params: { + threadId: sessionId, turnId: "plan-turn", tokenUsage: { + total: {totalTokens: 120, inputTokens: 100, cachedInputTokens: 0, cacheWriteInputTokens: 0, outputTokens: 20, reasoningOutputTokens: 0}, + last: {totalTokens: 120, inputTokens: 100, cachedInputTokens: 0, cacheWriteInputTokens: 0, outputTokens: 20, reasoningOutputTokens: 0}, + modelContextWindow: 128000, + }, + }}); + } const completion: TurnCompletion = { threadId: sessionId, turn: { @@ -165,6 +175,24 @@ describe("CodexACPAgent - plan review", () => { return {promptPromise, sessionState, turnStart, implementationTurn}; } + it("includes both planning and implementation usage in one prompt", async () => { + const {promptPromise, sessionState, implementationTurn} = await startPlanPrompt("implement_plan", {reportUsage: true}); + await vi.waitFor(() => expect(sessionState.currentTurnId).toBe("implementation-turn")); + fixture.sendServerNotification({method: "thread/tokenUsage/updated", params: { + threadId: sessionId, turnId: "implementation-turn", tokenUsage: { + total: {totalTokens: 350, inputTokens: 300, cachedInputTokens: 100, cacheWriteInputTokens: 0, outputTokens: 50, reasoningOutputTokens: 0}, + last: {totalTokens: 230, inputTokens: 200, cachedInputTokens: 100, cacheWriteInputTokens: 0, outputTokens: 30, reasoningOutputTokens: 0}, + modelContextWindow: 128000, + }, + }}); + implementationTurn.resolve({threadId: sessionId, turn: { + id: "implementation-turn", items: [], itemsView: "notLoaded", status: "completed", + error: null, startedAt: null, completedAt: null, durationMs: null, + }}); + const result = await promptPromise; + expect(result.usage).toMatchObject({totalTokens: 350, inputTokens: 200, cachedReadTokens: 100, outputTokens: 50}); + }); + it("requests plan permission and starts one implementation turn when approved", async () => { const {promptPromise, sessionState, turnStart, implementationTurn} = await startPlanPrompt("implement_plan"); diff --git a/src/__tests__/CodexACPAgent/prompt-usage-accounting.test.ts b/src/__tests__/CodexACPAgent/prompt-usage-accounting.test.ts new file mode 100644 index 00000000..81de4aee --- /dev/null +++ b/src/__tests__/CodexACPAgent/prompt-usage-accounting.test.ts @@ -0,0 +1,132 @@ +import {beforeEach, describe, expect, it, vi} from "vitest"; +import type {ServerNotification} from "../../app-server"; +import type {TokenUsageBreakdown} from "../../app-server/v2"; +import {createCodexMockTestFixture, createTestSessionState} from "../acp-test-utils"; + +const threadId = "accounting-thread"; +const request = {sessionId: threadId, prompt: [{type: "text" as const, text: "test"}]}; + +function count(input: number, output: number, cached = 0): TokenUsageBreakdown { + return {totalTokens: input + output, inputTokens: input, cachedInputTokens: cached, + cacheWriteInputTokens: 0, outputTokens: output, reasoningOutputTokens: 0}; +} + +function usage(turnId: string, total: TokenUsageBreakdown, last = total): ServerNotification { + return {method: "thread/tokenUsage/updated", params: { + threadId, turnId, tokenUsage: {total, last, modelContextWindow: 128000}, + }}; +} + +describe("prompt usage accounting", () => { + let fixture: ReturnType; + beforeEach(() => { + vi.clearAllMocks(); + fixture = createCodexMockTestFixture(); + vi.spyOn(fixture.getCodexAcpAgent(), "getSessionState") + .mockReturnValue(createTestSessionState({sessionId: threadId})); + }); + + function turn(turnId: string, events: ServerNotification[], status = "completed") { + fixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue({ + turn: {id: turnId, items: [], status: "inProgress", error: null}, + }); + fixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockImplementation(async () => { + for (const event of events) fixture.sendServerNotification(event); + return {threadId, turn: {id: turnId, items: [], status, error: null}}; + }); + return fixture.getCodexAcpAgent().prompt(request); + } + + it("counts identical requests separately and replayed snapshots once", async () => { + const first = usage("one", count(100, 20, 30)); + const second = usage("one", count(200, 40, 60), count(100, 20, 30)); + const result = await turn("one", [first, first, second, second]); + expect(result.usage).toEqual({totalTokens: 240, inputTokens: 140, + cachedReadTokens: 60, outputTokens: 40, thoughtTokens: 0}); + expect(result._meta?.["quota"]).toMatchObject({token_count: {totalTokens: 240}}); + expect(result._meta?.["usageAccounting"]).toMatchObject({scope: "root_thread_prompt", completeness: "reported"}); + }); + + it("does not charge a previous prompt again", async () => { + await turn("one", [usage("one", count(100, 20))]); + const result = await turn("two", [ + usage("one", count(100, 20)), // delayed historical notification + usage("two", count(300, 50), count(200, 30)), + usage("two", count(600, 90), count(300, 40)), + ]); + expect(result.usage?.totalTokens).toBe(570); + }); + + it("uses a live turn baseline after load, resume, or fork without in-memory counters", async () => { + fixture.getCodexAcpAgent().getSessionState(threadId).totalTokenUsage = null; + fixture.sendServerNotification(usage("historic", count(50000, 5000))); + const result = await turn("resumed", [ + usage("resumed", count(50100, 5020), count(100, 20)), + usage("resumed", count(50300, 5050), count(200, 30)), + ]); + expect(result.usage?.totalTokens).toBe(350); + }); + + it("does not charge a rate-limit-only snapshot at the beginning of the next turn", async () => { + await turn("one", [usage("one", count(100, 20))]); + const result = await turn("two", [ + usage("two", count(100, 20)), // same history, attributed to the new active turn + usage("two", count(300, 50), count(200, 30)), + ]); + expect(result.usage?.totalTokens).toBe(230); + }); + + it("keeps a missing resume baseline partial instead of charging historical last usage", async () => { + fixture.getCodexAcpAgent().getSessionState(threadId).totalTokenUsage = null; + const result = await turn("resumed", [ + usage("resumed", count(50000, 5000), count(100, 20)), + usage("resumed", count(50200, 5030), count(200, 30)), + ]); + expect(result.usage?.totalTokens).toBe(230); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + }); + + it("reports missing usage instead of borrowing the last prompt", async () => { + await turn("one", [usage("one", count(100, 20))]); + const result = await turn("two", []); + expect(result.usage).toBeNull(); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + }); + + it("retains all observed requests when cancelled", async () => { + const result = await turn("one", [ + usage("one", count(100, 20)), + usage("one", count(300, 50), count(200, 30)), + ], "interrupted"); + expect(result.stopReason).toBe("cancelled"); + expect(result.usage?.totalTokens).toBe(350); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + }); + + it("retains observed usage on a typed failed turn", async () => { + await fixture.getCodexAcpAgent().initialize({protocolVersion: 1, clientCapabilities: { + _meta: {jetbrains: {air: {version: 1, capabilities: ["sessionFailure"]}}}, + }}); + const result = await turn("one", [usage("one", count(100, 20))], "failed"); + expect(result.usage?.totalTokens).toBe(120); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + expect(result._meta?.["jetbrains"]).toMatchObject({air: {sessionFailure: {severity: "error"}}}); + }); + + it("preserves known usage and marks counter replacement incomplete", async () => { + const result = await turn("one", [ + usage("one", count(100, 20)), + usage("one", count(0, 0)), + usage("one", count(200, 30)), + ]); + expect(result.usage?.totalTokens).toBe(350); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + }); + + it("rejects synthetic context-full totals without inference counts", async () => { + const synthetic = {...count(0, 0), totalTokens: 128000}; + const result = await turn("one", [usage("one", count(100, 20)), usage("one", synthetic)]); + expect(result.usage?.totalTokens).toBe(120); + expect(result._meta?.["usageAccounting"]).toMatchObject({completeness: "partial"}); + }); +}); diff --git a/src/__tests__/CodexACPAgent/token-usage-events.test.ts b/src/__tests__/CodexACPAgent/token-usage-events.test.ts index dc49c1fc..fa96bd3d 100644 --- a/src/__tests__/CodexACPAgent/token-usage-events.test.ts +++ b/src/__tests__/CodexACPAgent/token-usage-events.test.ts @@ -61,7 +61,7 @@ describe('Token Usage Events', () => { inputTokens: 4000, cachedInputTokens: 1000, cacheWriteInputTokens: 0, - outputTokens: 900, + outputTokens: 1000, reasoningOutputTokens: 100, }, last: { @@ -69,7 +69,7 @@ describe('Token Usage Events', () => { inputTokens: 2000, cachedInputTokens: 500, cacheWriteInputTokens: 0, - outputTokens: 450, + outputTokens: 500, reasoningOutputTokens: 50, }, modelContextWindow: 128000, @@ -133,7 +133,7 @@ describe('Token Usage Events', () => { ); }); - it('should use last token usage from multiple updates', async () => { + it('should accumulate all requests from multiple updates', async () => { const notifications: ServerNotification[] = [ createTokenUsageNotification(sessionId, { total: { totalTokens: 1000, inputTokens: 800, cachedInputTokens: 0, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 0 }, @@ -146,8 +146,8 @@ describe('Token Usage Events', () => { modelContextWindow: 128000, }), createTokenUsageNotification(sessionId, { - total: { totalTokens: 3500, inputTokens: 2800, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 600, reasoningOutputTokens: 100 }, - last: { totalTokens: 1500, inputTokens: 1200, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 100 }, + total: { totalTokens: 3500, inputTokens: 2800, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 700, reasoningOutputTokens: 100 }, + last: { totalTokens: 1500, inputTokens: 1200, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 300, reasoningOutputTokens: 100 }, modelContextWindow: 128000, }), ]; @@ -202,7 +202,7 @@ describe('Token Usage Events', () => { inputTokens: 4000, cachedInputTokens: 1000, cacheWriteInputTokens: 0, - outputTokens: 900, + outputTokens: 1000, reasoningOutputTokens: 100, }, last: { @@ -210,7 +210,7 @@ describe('Token Usage Events', () => { inputTokens: 2000, cachedInputTokens: 500, cacheWriteInputTokens: 0, - outputTokens: 450, + outputTokens: 500, reasoningOutputTokens: 50, }, modelContextWindow: 128000, @@ -233,8 +233,8 @@ describe('Token Usage Events', () => { modelContextWindow: 128000, }), createTokenUsageNotification(sessionId, { - total: { totalTokens: 3500, inputTokens: 2800, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 600, reasoningOutputTokens: 100 }, - last: { totalTokens: 1500, inputTokens: 1200, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 100 }, + total: { totalTokens: 3500, inputTokens: 2800, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 700, reasoningOutputTokens: 100 }, + last: { totalTokens: 1500, inputTokens: 1200, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 300, reasoningOutputTokens: 100 }, modelContextWindow: 128000, }), ])(); @@ -245,8 +245,8 @@ describe('Token Usage Events', () => { it('should skip usage_update when model context window is unavailable', async () => { const events = await setupPromptAndReturnEvents([ createTokenUsageNotification(sessionId, { - total: { totalTokens: 5000, inputTokens: 4000, cachedInputTokens: 1000, cacheWriteInputTokens: 0, outputTokens: 900, reasoningOutputTokens: 100 }, - last: { totalTokens: 2500, inputTokens: 2000, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 450, reasoningOutputTokens: 50 }, + total: { totalTokens: 5000, inputTokens: 4000, cachedInputTokens: 1000, cacheWriteInputTokens: 0, outputTokens: 1000, reasoningOutputTokens: 100 }, + last: { totalTokens: 2500, inputTokens: 2000, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 500, reasoningOutputTokens: 50 }, modelContextWindow: null, }), ])(); diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts index f358664c..a070f720 100644 --- a/src/__tests__/acp-test-utils.ts +++ b/src/__tests__/acp-test-utils.ts @@ -400,7 +400,7 @@ export function createTestSessionState(overrides?: Partial): Sessi return { currentTurnId: null, lastTokenUsage: null, - totalTokenUsage: null, + totalTokenUsage: {totalTokens: 0, inputTokens: 0, cachedInputTokens: 0, outputTokens: 0, reasoningOutputTokens: 0}, modelContextWindow: null, rateLimits: null, account: null, diff --git a/src/subagents/CodexSubagentSubscriptions.ts b/src/subagents/CodexSubagentSubscriptions.ts index 27998e08..8bb49cbc 100644 --- a/src/subagents/CodexSubagentSubscriptions.ts +++ b/src/subagents/CodexSubagentSubscriptions.ts @@ -46,8 +46,13 @@ export class CodexSubagentSubscriptions { this.registerInteractiveHandlers(session, subscription.rootSessionId); } + /** Child thread ids discovered under a root session; cleared with it. */ + childSessionIds(rootSessionId: string): string[] { + return [...this.sessions.get(rootSessionId)?.children ?? []]; + } + clear(rootSessionId: string): void { - for (const childSessionId of this.sessions.get(rootSessionId)?.children ?? []) { + for (const childSessionId of this.childSessionIds(rootSessionId)) { this.client.clearThreadHandlers(childSessionId); } this.sessions.delete(rootSessionId);