diff --git a/src/CodexAcpClient.ts b/src/CodexAcpClient.ts index 0c7af207..309e27a3 100644 --- a/src/CodexAcpClient.ts +++ b/src/CodexAcpClient.ts @@ -36,6 +36,7 @@ import type { ThreadGoal, ThreadGoalStatus, ThreadSourceKind, + ThreadTokenUsage, TurnCompletedNotification, TurnSteerResponse, UserInput, @@ -656,6 +657,10 @@ export class CodexAcpClient { }); } + getThreadTokenUsage(sessionId: string): ThreadTokenUsage | null { + return this.codexClient.getThreadTokenUsage(sessionId); + } + async waitForSessionNotifications(sessionId: string): Promise { while (true) { const queue = this.sessionNotificationQueues.get(sessionId); diff --git a/src/CodexAcpServer.ts b/src/CodexAcpServer.ts index f1fb9cec..5ed0f3b9 100644 --- a/src/CodexAcpServer.ts +++ b/src/CodexAcpServer.ts @@ -35,7 +35,7 @@ import { REASONING_EFFORT_CONFIG_ID, } from "./ModelConfigOption"; import type {TokenCount} from "./TokenCount"; -import {toPromptUsage} from "./TokenCount"; +import {subtractTokenCounts, toPromptUsage, toTokenCount} from "./TokenCount"; import {CodexCommands} from "./CodexCommands"; import {SteeringQueue} from "./SteeringQueue"; import type {QuotaMeta} from "./QuotaMeta"; @@ -1878,6 +1878,12 @@ export class CodexAcpServer { prompt: params.prompt, }); const sessionState = this.getSessionState(params.sessionId); + const latestThreadTokenUsage = this.codexAcpClient.getThreadTokenUsage(params.sessionId); + if (latestThreadTokenUsage != null) { + sessionState.totalTokenUsage = toTokenCount(latestThreadTokenUsage.total); + sessionState.modelContextWindow = latestThreadTokenUsage.modelContextWindow; + } + const promptStartTokenUsage = sessionState.totalTokenUsage; sessionState.currentTurnId = null; sessionState.lastTokenUsage = null; const activePrompt = this.trackActivePrompt(params.sessionId); @@ -1915,7 +1921,7 @@ export class CodexAcpServer { elicitationHandler); if (activePrompt.signal.aborted) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } const commandPromise = this.availableCommands.tryHandleCommand(params.prompt, sessionState, { @@ -1957,28 +1963,32 @@ export class CodexAcpServer { this.cancelBeforeTurnStarted(activePrompt), ]); if (commandResult === null) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } if (commandResult.handled) { logger.log("Prompt handled by a command"); await this.codexAcpClient.waitForSessionNotifications(params.sessionId); if (commandResult.turnCompleted?.turn.status === "interrupted") { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } const error = eventHandler.getFailure(); if (error) { // noinspection ExceptionCaughtLocallyJS throw error; } + const promptTokenUsage = subtractTokenCounts( + sessionState.totalTokenUsage, + promptStartTokenUsage, + ); return { stopReason: "end_turn", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), - _meta: this.buildQuotaMeta(sessionState), + usage: this.buildPromptUsage(promptTokenUsage), + _meta: this.buildQuotaMeta(sessionState, promptTokenUsage), }; } if (this.sessionIsClosing(params.sessionId)) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } const modelId = ModelId.fromString(sessionState.currentModelId); @@ -2036,14 +2046,14 @@ export class CodexAcpServer { ]); if (turnCompleted === null) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } await this.codexAcpClient.waitForSessionNotifications(params.sessionId); if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } const error = eventHandler.getFailure(); @@ -2065,7 +2075,7 @@ export class CodexAcpServer { activePrompt.signal, ); if (this.promptShouldStop(params.sessionId, activePrompt)) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } if (approved && !this.promptShouldStop(params.sessionId, activePrompt)) { await this.applyCollaborationModeChange(sessionState, DEFAULT_COLLABORATION_MODE); @@ -2113,13 +2123,13 @@ export class CodexAcpServer { ]); if (turnCompleted === null) { - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } await this.codexAcpClient.waitForSessionNotifications(params.sessionId); if (turnCompleted.turn.status === "interrupted") { await eventHandler.flushPendingPlanUpdates(); - return this.cancelledPromptResponse(sessionState); + return this.cancelledPromptResponse(sessionState, promptStartTokenUsage); } const implementationError = eventHandler.getFailure(); @@ -2134,10 +2144,14 @@ export class CodexAcpServer { this.createPromptFallbackTitle(params.prompt), ); + const promptTokenUsage = subtractTokenCounts( + sessionState.totalTokenUsage, + promptStartTokenUsage, + ); return { stopReason: "end_turn", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), - _meta: this.buildQuotaMeta(sessionState), + usage: this.buildPromptUsage(promptTokenUsage), + _meta: this.buildQuotaMeta(sessionState, promptTokenUsage), }; } catch (err) { logger.error(`Prompt for session ${params.sessionId} failed`, err); @@ -2215,38 +2229,46 @@ export class CodexAcpServer { } } - private cancelledPromptResponse(sessionState: SessionState): acp.PromptResponse { + private cancelledPromptResponse( + sessionState: SessionState, + promptStartTokenUsage: TokenCount | null, + ): acp.PromptResponse { + const promptTokenUsage = subtractTokenCounts( + sessionState.totalTokenUsage, + promptStartTokenUsage, + ); return { stopReason: "cancelled", - usage: this.buildPromptUsage(sessionState.lastTokenUsage), - _meta: this.buildQuotaMeta(sessionState), + usage: this.buildPromptUsage(promptTokenUsage), + _meta: this.buildQuotaMeta(sessionState, promptTokenUsage), }; } - private buildQuotaMeta(sessionState: SessionState): { quota: QuotaMeta } { - const lastTokenUsage = sessionState.lastTokenUsage; - + private buildQuotaMeta( + sessionState: SessionState, + promptTokenUsage: TokenCount | null, + ): { quota: QuotaMeta } { // Remove the "[reasoning-level]" suffix from currentModelId if present const modelName = sessionState.currentModelId.replace(/\[.*?]$/, ''); // FIXME: currently all tokens are reported for the current model - const modelUsage = (lastTokenUsage != null) - ? [{ model: modelName, token_count: lastTokenUsage }] + const modelUsage = (promptTokenUsage != null) + ? [{ model: modelName, token_count: promptTokenUsage }] : []; return { quota: { - token_count: sessionState.lastTokenUsage, + token_count: promptTokenUsage, model_usage: modelUsage } }; } - private buildPromptUsage(lastTokenUsage: TokenCount | null): acp.Usage | null { - if (lastTokenUsage == null) { + private buildPromptUsage(promptTokenUsage: TokenCount | null): acp.Usage | null { + if (promptTokenUsage == null) { return null; } - return toPromptUsage(lastTokenUsage); + return toPromptUsage(promptTokenUsage); } private async runWithProcessCheck(operation: () => Promise): Promise { diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts index eb26c83f..b92a8446 100644 --- a/src/CodexAppServerClient.ts +++ b/src/CodexAppServerClient.ts @@ -52,6 +52,7 @@ import type { ThreadSettings, ThreadStartParams, ThreadStartResponse, + ThreadTokenUsage, ThreadUnsubscribeParams, ThreadUnsubscribeResponse, ToolRequestUserInputParams, @@ -145,6 +146,7 @@ export class CodexAppServerClient { private readonly threadGoalUpdateCaptures = new Map void>>(); private readonly threadGoalClearedCaptures = new Map void>>(); private readonly threadSettings = new Map(); + private readonly threadTokenUsage = new Map(); private readonly staleTurnIds = new Map>(); constructor(connection: MessageConnection) { @@ -186,6 +188,9 @@ export class CodexAppServerClient { if (this.handleStaleTurnNotification(serverNotification, routing)) { return; } + if (serverNotification.method === "thread/tokenUsage/updated") { + this.threadTokenUsage.set(serverNotification.params.threadId, serverNotification.params.tokenUsage); + } this.notify(serverNotification); for (const callback of this.codexEventHandlers) { callback({ eventType: "notification", ...serverNotification }); @@ -260,6 +265,7 @@ export class CodexAppServerClient { this.notificationHandlers.delete(threadId); this.approvalHandlers.delete(threadId); this.elicitationHandlers.delete(threadId); + this.threadTokenUsage.delete(threadId); } async initialize(params: InitializeParams): Promise { @@ -532,6 +538,10 @@ export class CodexAppServerClient { return this.threadSettings.get(threadId); } + getThreadTokenUsage(threadId: string): ThreadTokenUsage | null { + return this.threadTokenUsage.get(threadId) ?? null; + } + async threadSettingsUpdate(params: ExperimentalThreadSettingsUpdateParams): Promise { await this.connection.sendRequest("thread/settings/update", params); } diff --git a/src/TokenCount.ts b/src/TokenCount.ts index 1f7cf0b5..e3f42e9c 100644 --- a/src/TokenCount.ts +++ b/src/TokenCount.ts @@ -19,6 +19,46 @@ export interface TokenCount { reasoningOutputTokens: number; } +/** + * Returns the category-by-category usage added to a cumulative token count. + * Each category is clamped independently because an upstream counter may be + * reset or corrected between snapshots. + */ +export function subtractTokenCounts( + current: TokenCount | null, + previous: TokenCount | null, +): TokenCount | null { + if (current == null) { + return null; + } + + const difference = (currentValue: number, previousValue: number): number => + Math.max(0, currentValue - previousValue); + const baseline = previous ?? { + totalTokens: 0, + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + reasoningOutputTokens: 0, + }; + + const inputTokens = difference(current.inputTokens, baseline.inputTokens); + const cachedInputTokens = difference(current.cachedInputTokens, baseline.cachedInputTokens); + const outputTokens = difference(current.outputTokens, baseline.outputTokens); + const reasoningOutputTokens = Math.min( + outputTokens, + difference(current.reasoningOutputTokens, baseline.reasoningOutputTokens), + ); + + return { + totalTokens: inputTokens + cachedInputTokens + outputTokens, + inputTokens, + cachedInputTokens, + outputTokens, + reasoningOutputTokens, + }; +} + /** * Maps Codex's TokenUsageBreakdown to our TokenCount interface. * This explicit mapping ensures compile-time errors if Codex changes their types. diff --git a/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json b/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json index 426011ab..ecd060ef 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-cancelled.json @@ -1,29 +1,29 @@ { "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 } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json b/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json index 60ba2e2e..f6e92140 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-end-turn.json @@ -1,30 +1,30 @@ { "stopReason": "end_turn", "usage": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedReadTokens": 500, - "outputTokens": 450, - "thoughtTokens": 50 + "totalTokens": 5000, + "inputTokens": 3100, + "cachedReadTokens": 1000, + "outputTokens": 900, + "thoughtTokens": 100 }, "_meta": { "quota": { "token_count": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedInputTokens": 500, - "outputTokens": 450, - "reasoningOutputTokens": 50 + "totalTokens": 5000, + "inputTokens": 3100, + "cachedInputTokens": 1000, + "outputTokens": 900, + "reasoningOutputTokens": 100 }, "model_usage": [ { "model": "model-id", "token_count": { - "totalTokens": 2500, - "inputTokens": 1500, - "cachedInputTokens": 500, - "outputTokens": 450, - "reasoningOutputTokens": 50 + "totalTokens": 5000, + "inputTokens": 3100, + "cachedInputTokens": 1000, + "outputTokens": 900, + "reasoningOutputTokens": 100 } } ] diff --git a/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json b/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json index aebc1783..439c5fca 100644 --- a/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json +++ b/src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json @@ -1,29 +1,29 @@ { "stopReason": "end_turn", "usage": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2400, "cachedReadTokens": 500, - "outputTokens": 200, + "outputTokens": 600, "thoughtTokens": 100 }, "_meta": { "quota": { "token_count": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2400, "cachedInputTokens": 500, - "outputTokens": 200, + "outputTokens": 600, "reasoningOutputTokens": 100 }, "model_usage": [ { "model": "model-id", "token_count": { - "totalTokens": 1500, - "inputTokens": 700, + "totalTokens": 3500, + "inputTokens": 2400, "cachedInputTokens": 500, - "outputTokens": 200, + "outputTokens": 600, "reasoningOutputTokens": 100 } } diff --git a/src/__tests__/CodexACPAgent/data/token-usage-replayed-baseline.json b/src/__tests__/CodexACPAgent/data/token-usage-replayed-baseline.json new file mode 100644 index 00000000..7095b847 --- /dev/null +++ b/src/__tests__/CodexACPAgent/data/token-usage-replayed-baseline.json @@ -0,0 +1,33 @@ +{ + "stopReason": "end_turn", + "usage": { + "totalTokens": 1500, + "inputTokens": 900, + "cachedReadTokens": 300, + "outputTokens": 300, + "thoughtTokens": 50 + }, + "_meta": { + "quota": { + "token_count": { + "totalTokens": 1500, + "inputTokens": 900, + "cachedInputTokens": 300, + "outputTokens": 300, + "reasoningOutputTokens": 50 + }, + "model_usage": [ + { + "model": "model-id", + "token_count": { + "totalTokens": 1500, + "inputTokens": 900, + "cachedInputTokens": 300, + "outputTokens": 300, + "reasoningOutputTokens": 50 + } + } + ] + } + } +} diff --git a/src/__tests__/CodexACPAgent/data/token-usage-second-prompt.json b/src/__tests__/CodexACPAgent/data/token-usage-second-prompt.json new file mode 100644 index 00000000..ee120a5a --- /dev/null +++ b/src/__tests__/CodexACPAgent/data/token-usage-second-prompt.json @@ -0,0 +1,33 @@ +{ + "stopReason": "end_turn", + "usage": { + "totalTokens": 2900, + "inputTokens": 1400, + "cachedReadTokens": 900, + "outputTokens": 600, + "thoughtTokens": 150 + }, + "_meta": { + "quota": { + "token_count": { + "totalTokens": 2900, + "inputTokens": 1400, + "cachedInputTokens": 900, + "outputTokens": 600, + "reasoningOutputTokens": 150 + }, + "model_usage": [ + { + "model": "model-id", + "token_count": { + "totalTokens": 2900, + "inputTokens": 1400, + "cachedInputTokens": 900, + "outputTokens": 600, + "reasoningOutputTokens": 150 + } + } + ] + } + } +} diff --git a/src/__tests__/CodexACPAgent/token-usage-events.test.ts b/src/__tests__/CodexACPAgent/token-usage-events.test.ts index dc49c1fc..c7f3d473 100644 --- a/src/__tests__/CodexACPAgent/token-usage-events.test.ts +++ b/src/__tests__/CodexACPAgent/token-usage-events.test.ts @@ -30,8 +30,13 @@ describe('Token Usage Events', () => { vi.clearAllMocks(); }); describe('PromptResponse usage', () => { - function setupPromptWithTokenUsage(notifications: ServerNotification[], turnStatus: string = "completed") { + function setupPromptsWithTokenUsage( + notificationsByPrompt: ServerNotification[][], + turnStatuses: string[] = [], + sessionState = createTestSessionState({ sessionId }), + ) { const codexAcpAgent = mockFixture.getCodexAcpAgent(); + let promptIndex = 0; mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue({ turn: { id: "turn-id", items: [], status: "inProgress", error: null } @@ -40,25 +45,35 @@ describe('Token Usage Events', () => { // awaitTurnCompleted sends notifications before resolving mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockImplementation(async () => { // Send notifications during turn (after handler is registered) - for (const notification of notifications) { + const currentPromptIndex = promptIndex++; + for (const notification of notificationsByPrompt[currentPromptIndex] ?? []) { mockFixture.sendServerNotification(notification); } return { threadId: sessionId, - turn: { id: "turn-id", items: [], status: turnStatus, error: null } + turn: { + id: "turn-id", + items: [], + status: turnStatuses[currentPromptIndex] ?? "completed", + error: null, + } }; }); - vi.spyOn(codexAcpAgent, 'getSessionState').mockReturnValue(createTestSessionState({ sessionId })); + vi.spyOn(codexAcpAgent, 'getSessionState').mockReturnValue(sessionState); return codexAcpAgent; } + function setupPromptWithTokenUsage(notifications: ServerNotification[], turnStatus: string = "completed") { + return setupPromptsWithTokenUsage([notifications], [turnStatus]); + } + it('should include token_count in PromptResponse on end_turn', async () => { const tokenUsageNotification = createTokenUsageNotification(sessionId, { total: { totalTokens: 5000, - inputTokens: 4000, + inputTokens: 4100, cachedInputTokens: 1000, cacheWriteInputTokens: 0, outputTokens: 900, @@ -133,7 +148,31 @@ describe('Token Usage Events', () => { ); }); - it('should use last token usage from multiple updates', async () => { + it('should subtract cumulative usage replayed before prompt subscription', async () => { + mockFixture.sendServerNotification(createTokenUsageNotification(sessionId, { + total: { totalTokens: 2000, inputTokens: 1600, cachedInputTokens: 400, cacheWriteInputTokens: 0, outputTokens: 400, reasoningOutputTokens: 50 }, + last: { totalTokens: 1000, inputTokens: 800, cachedInputTokens: 200, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 25 }, + modelContextWindow: 128000, + })); + const codexAcpAgent = setupPromptWithTokenUsage([ + createTokenUsageNotification(sessionId, { + total: { totalTokens: 3500, inputTokens: 2800, cachedInputTokens: 700, cacheWriteInputTokens: 0, outputTokens: 700, reasoningOutputTokens: 100 }, + last: { totalTokens: 1500, inputTokens: 1200, cachedInputTokens: 300, cacheWriteInputTokens: 0, outputTokens: 300, reasoningOutputTokens: 50 }, + modelContextWindow: 128000, + }), + ]); + + const response = await codexAcpAgent.prompt({ + sessionId, + prompt: [{ type: 'text', text: 'test prompt' }], + }); + + await expect(`${JSON.stringify(response, null, 2)}\n`).toMatchFileSnapshot( + 'data/token-usage-replayed-baseline.json' + ); + }); + + it('should report cumulative usage from multiple updates within one prompt', async () => { const notifications: ServerNotification[] = [ createTokenUsageNotification(sessionId, { total: { totalTokens: 1000, inputTokens: 800, cachedInputTokens: 0, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 0 }, @@ -146,8 +185,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: 2900, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 600, reasoningOutputTokens: 100 }, + last: { totalTokens: 1500, inputTokens: 1300, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 100 }, modelContextWindow: 128000, }), ]; @@ -163,6 +202,43 @@ describe('Token Usage Events', () => { 'data/token-usage-multiple-updates.json' ); }); + + it('should report only the cumulative usage added by a second prompt', async () => { + const codexAcpAgent = setupPromptsWithTokenUsage([ + [ + createTokenUsageNotification(sessionId, { + total: { totalTokens: 3500, inputTokens: 2900, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 600, reasoningOutputTokens: 100 }, + last: { totalTokens: 1500, inputTokens: 1300, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 200, reasoningOutputTokens: 100 }, + modelContextWindow: 128000, + }), + ], + [ + createTokenUsageNotification(sessionId, { + total: { totalTokens: 4800, inputTokens: 3900, cachedInputTokens: 900, cacheWriteInputTokens: 0, outputTokens: 900, reasoningOutputTokens: 180 }, + last: { totalTokens: 1300, inputTokens: 1000, cachedInputTokens: 400, cacheWriteInputTokens: 0, outputTokens: 300, reasoningOutputTokens: 80 }, + modelContextWindow: 128000, + }), + createTokenUsageNotification(sessionId, { + total: { totalTokens: 6400, inputTokens: 5200, cachedInputTokens: 1400, cacheWriteInputTokens: 0, outputTokens: 1200, reasoningOutputTokens: 250 }, + last: { totalTokens: 1600, inputTokens: 1300, cachedInputTokens: 500, cacheWriteInputTokens: 0, outputTokens: 300, reasoningOutputTokens: 70 }, + modelContextWindow: 128000, + }), + ], + ]); + + await codexAcpAgent.prompt({ + sessionId, + prompt: [{ type: 'text', text: 'first prompt' }], + }); + const response = await codexAcpAgent.prompt({ + sessionId, + prompt: [{ type: 'text', text: 'second prompt' }], + }); + + await expect(`${JSON.stringify(response, null, 2)}\n`).toMatchFileSnapshot( + 'data/token-usage-second-prompt.json' + ); + }); }); describe('session/update usage_update', () => { diff --git a/src/__tests__/TokenCount.test.ts b/src/__tests__/TokenCount.test.ts new file mode 100644 index 00000000..dd742c67 --- /dev/null +++ b/src/__tests__/TokenCount.test.ts @@ -0,0 +1,29 @@ +import {describe, expect, it} from "vitest"; +import {subtractTokenCounts} from "../TokenCount"; + +describe("subtractTokenCounts", () => { + it("clamps cumulative counter decreases and keeps the total consistent", () => { + expect(subtractTokenCounts( + { + totalTokens: 900, + inputTokens: 650, + cachedInputTokens: 100, + outputTokens: 150, + reasoningOutputTokens: 20, + }, + { + totalTokens: 1000, + inputTokens: 600, + cachedInputTokens: 200, + outputTokens: 200, + reasoningOutputTokens: 40, + }, + )).toEqual({ + totalTokens: 50, + inputTokens: 50, + cachedInputTokens: 0, + outputTokens: 0, + reasoningOutputTokens: 0, + }); + }); +});