diff --git a/packages/server/src/server/agent/providers/claude/agent.context-usage.local.e2e.test.ts b/packages/server/src/server/agent/providers/claude/agent.context-usage.local.e2e.test.ts new file mode 100644 index 0000000000..3008f3d4d5 --- /dev/null +++ b/packages/server/src/server/agent/providers/claude/agent.context-usage.local.e2e.test.ts @@ -0,0 +1,133 @@ +import { createServer } from "node:http"; +import { mkdir, mkdtemp, rm } from "node:fs/promises"; +import path from "node:path"; +import pino from "pino"; +import { expect, test } from "vitest"; +import { collectSessionTurnEvents } from "../test-utils/session-stream-adapter.js"; +import { ClaudeAgentClient } from "./agent.js"; + +// Exercise the real Claude Code process against a local Messages API fixture. +// The input counter's location is the only difference between the two backends. +test.each([ + { placement: "start", cacheTokens: 0 }, + { placement: "delta", cacheTokens: 0 }, + { placement: "start", cacheTokens: 1_000 }, + { placement: "delta", cacheTokens: 1_000 }, +])( + "counts $placement input and $cacheTokens cached tokens through Claude Code", + async ({ placement, cacheTokens }) => { + const cwd = await mkdtemp(path.resolve(".claude-context-usage-")); + const configDir = path.join(cwd, "config"); + await mkdir(configDir); + let requestCount = 0; + const server = createServer(async (request, response) => { + for await (const chunk of request) { + void chunk; + } + if (!request.url?.startsWith("/v1/messages")) { + response.writeHead(200, { "content-type": "application/json" }); + response.end("{}"); + return; + } + if (request.url.includes("count_tokens")) { + response.writeHead(200, { "content-type": "application/json" }); + response.end(JSON.stringify({ input_tokens: 20_236 })); + return; + } + const inputTokens = 20_236 + requestCount++; + response.writeHead(200, { "content-type": "text/event-stream" }); + const emit = (data: { type: string; [key: string]: unknown }) => { + response.write(`event: ${data.type}\ndata: ${JSON.stringify(data)}\n\n`); + }; + emit({ + type: "message_start", + message: { + id: `msg_context_usage_${requestCount}`, + type: "message", + role: "assistant", + model: "claude-sonnet-4-6", + content: [], + stop_reason: null, + stop_sequence: null, + usage: { + input_tokens: placement === "start" ? inputTokens : 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: cacheTokens, + }, + }, + }); + emit({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }); + emit({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hello" } }); + emit({ type: "content_block_stop", index: 0 }); + emit({ + type: "message_delta", + delta: { stop_reason: "end_turn", stop_sequence: null }, + usage: { + output_tokens: 64, + ...(placement === "delta" ? { input_tokens: inputTokens } : {}), + }, + }); + emit({ type: "message_stop" }); + response.end(); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("Fixture has no TCP address"); + const client = new ClaudeAgentClient({ + logger: pino({ level: "silent" }), + runtimeSettings: { + env: { + ANTHROPIC_BASE_URL: `http://127.0.0.1:${address.port}`, + ANTHROPIC_AUTH_TOKEN: "fixture-token", + ANTHROPIC_API_KEY: "", + CLAUDE_CONFIG_DIR: configDir, + CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: "1", + }, + }, + }); + const session = await client.createSession( + { provider: "claude", cwd, model: "claude-sonnet-4-6" }, + undefined, + { persistSession: false }, + ); + try { + for (let turn = 0; turn < 2; turn++) { + const events = await collectSessionTurnEvents(session, "say hello", { + signal: AbortSignal.timeout(20_000), + }); + const expectedTokens = 20_236 + turn + cacheTokens + 64; + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ + inputTokens: 20_236 + turn, + outputTokens: 64, + cachedInputTokens: cacheTokens, + contextWindowUsedTokens: expectedTokens, + }), + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + usage: expect.objectContaining({ contextWindowUsedTokens: expectedTokens }), + }), + ); + } + expect(requestCount).toBe(2); + } finally { + try { + await session.close(); + } finally { + try { + server.closeAllConnections(); + await new Promise((resolve) => server.close(() => resolve())); + } finally { + await rm(cwd, { recursive: true, force: true }); + } + } + } + }, + 60_000, +); diff --git a/packages/server/src/server/agent/providers/claude/agent.test.ts b/packages/server/src/server/agent/providers/claude/agent.test.ts index 5a5e1132a1..5e26e6a55a 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -2135,13 +2135,15 @@ describe("ClaudeAgentSession context window usage", () => { } function createMessageDeltaEvent(outputTokens: number): Record { + return createMessageDeltaUsageEvent({ output_tokens: outputTokens }); + } + + function createMessageDeltaUsageEvent(usage: Record): Record { return { type: "stream_event", event: { type: "message_delta", - usage: { - output_tokens: outputTokens, - }, + usage, }, session_id: "session-1", }; @@ -2887,6 +2889,119 @@ describe("ClaudeAgentSession context window usage", () => { } }); + test("message_delta input usage replaces an empty message_start count", async () => { + // Translating gateways (LiteLLM in front of an OpenAI-compatible backend) do not know the + // prompt size when message_start is emitted, so they report input_tokens: 0 there and send + // the cumulative counts on message_delta, as the Anthropic Messages API allows. + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent({ + input_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }), + createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }), + createSuccessResult({ + usage: { + input_tokens: 20_236, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + output_tokens: 64, + iterations: [], + }, + }), + ], + ]); + + try { + const events = await collectStreamEvents(session); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowUsedTokens: 20_300, + }, + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ + inputTokens: 20_236, + outputTokens: 64, + contextWindowUsedTokens: 20_300, + }), + }), + ); + } finally { + await session.close(); + } + }); + + test("message_delta without cache counters keeps the message_start cache usage", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent({ + input_tokens: 5, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 1_000, + }), + createMessageDeltaUsageEvent({ input_tokens: 5, output_tokens: 64 }), + createSuccessResult(), + ], + ]); + + try { + const events = await collectStreamEvents(session); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowUsedTokens: 1_069, + }, + }), + ); + } finally { + await session.close(); + } + }); + + test("message_delta before this request's message_start does not seed usage", async () => { + // A canceled request can leave its terminal message_delta queued behind the next turn. + // Nothing is known about the new request until its message_start, so that straggler + // must not produce a usage_updated event or leak into the next request's counters. + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }), + createMessageStartEvent({ + input_tokens: 40, + cache_creation_input_tokens: 5, + cache_read_input_tokens: 10, + }), + createMessageDeltaEvent(7), + createSuccessResult(), + ], + ]); + + try { + const events = await collectStreamEvents(session); + + const usedTokens = events + .filter((event) => event.type === "usage_updated") + .map((event) => event.usage.contextWindowUsedTokens); + expect(usedTokens).toEqual([55, 62]); + } finally { + await session.close(); + } + }); + test("per-request stream usage is not cumulative across API calls in a turn", async () => { const session = await createSessionForTurns([ [ diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 6c73f1f643..51885f74e8 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1855,30 +1855,59 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined { return maxContextWindow; } -function readStreamRequestInputTokens(event: Record): number | undefined { - const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage); - if (!messageUsage) { +interface StreamRequestInputUsage { + inputTokens: number; + cacheCreationInputTokens: number; + cacheReadInputTokens: number; +} + +function readUsageCount(usage: Record, key: string): number | undefined { + const value = usage[key]; + if (typeof value !== "number" || !Number.isFinite(value) || value < 0) { return undefined; } - const usage = messageUsage; - const inputTokens = - typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens) - ? usage.input_tokens - : undefined; - const cacheCreationInputTokens = - typeof usage.cache_creation_input_tokens === "number" && - Number.isFinite(usage.cache_creation_input_tokens) - ? usage.cache_creation_input_tokens - : 0; - const cacheReadInputTokens = - typeof usage.cache_read_input_tokens === "number" && - Number.isFinite(usage.cache_read_input_tokens) - ? usage.cache_read_input_tokens - : 0; - if (typeof inputTokens !== "number" || inputTokens < 0) { + return value; +} + +function readStreamRequestInputUsage( + event: Record, +): StreamRequestInputUsage | undefined { + const usage = toObjectRecord(toObjectRecord(event.message)?.usage); + if (!usage) { return undefined; } - return inputTokens + cacheCreationInputTokens + cacheReadInputTokens; + const inputTokens = readUsageCount(usage, "input_tokens"); + if (inputTokens === undefined) { + return undefined; + } + return { + inputTokens, + cacheCreationInputTokens: readUsageCount(usage, "cache_creation_input_tokens") ?? 0, + cacheReadInputTokens: readUsageCount(usage, "cache_read_input_tokens") ?? 0, + }; +} + +/** + * The input counters on `message_delta.usage` are cumulative whole-request totals and are + * optional. Mirror the Anthropic SDK's MessageStream: a counter present on the delta replaces + * the `message_start` value, a missing counter keeps it. Translating gateways report + * `input_tokens: 0` on `message_start` and only know the real count by `message_delta`. + */ +function mergeStreamDeltaInputUsage( + current: StreamRequestInputUsage, + event: Record, +): StreamRequestInputUsage { + const usage = toObjectRecord(event.usage); + if (!usage) { + return current; + } + return { + inputTokens: readUsageCount(usage, "input_tokens") ?? current.inputTokens, + cacheCreationInputTokens: + readUsageCount(usage, "cache_creation_input_tokens") ?? current.cacheCreationInputTokens, + cacheReadInputTokens: + readUsageCount(usage, "cache_read_input_tokens") ?? current.cacheReadInputTokens, + }; } function readStreamRequestOutputTokens(event: Record): number | undefined { @@ -1958,7 +1987,7 @@ function readClaudeParentToolUseId(message: SDKMessage): string | null { class ClaudeContextUsageState { private contextWindowMaxTokens: number | undefined; - private streamRequestInputTokens: number | undefined; + private streamRequestInputUsage: StreamRequestInputUsage | undefined; private streamRequestOutputTokens: number | undefined; private compactedContextWindowUsedTokens: number | undefined; private completedResultTurns = 0; @@ -1968,7 +1997,7 @@ class ClaudeContextUsageState { } beginTurn(): void { - this.streamRequestInputTokens = undefined; + this.streamRequestInputUsage = undefined; this.streamRequestOutputTokens = undefined; this.compactedContextWindowUsedTokens = undefined; } @@ -1992,11 +2021,11 @@ class ClaudeContextUsageState { } const eventType = readTrimmedString(streamEvent.type); if (eventType === "message_start") { - const inputTokens = readStreamRequestInputTokens(streamEvent); - if (typeof inputTokens !== "number") { + const inputUsage = readStreamRequestInputUsage(streamEvent); + if (!inputUsage) { return null; } - this.streamRequestInputTokens = inputTokens; + this.streamRequestInputUsage = inputUsage; this.streamRequestOutputTokens = 0; } else if (eventType === "message_delta") { const outputTokens = readStreamRequestOutputTokens(streamEvent); @@ -2004,6 +2033,12 @@ class ClaudeContextUsageState { return null; } this.streamRequestOutputTokens = outputTokens; + if (this.streamRequestInputUsage) { + this.streamRequestInputUsage = mergeStreamDeltaInputUsage( + this.streamRequestInputUsage, + streamEvent, + ); + } } else { return null; } @@ -2050,13 +2085,15 @@ class ClaudeContextUsageState { } private streamUsedTokens(): number | undefined { - if ( - typeof this.streamRequestInputTokens !== "number" || - typeof this.streamRequestOutputTokens !== "number" - ) { + const inputUsage = this.streamRequestInputUsage; + if (!inputUsage || typeof this.streamRequestOutputTokens !== "number") { return undefined; } - const usedTokens = this.streamRequestInputTokens + this.streamRequestOutputTokens; + const usedTokens = + inputUsage.inputTokens + + inputUsage.cacheCreationInputTokens + + inputUsage.cacheReadInputTokens + + this.streamRequestOutputTokens; return usedTokens > 0 ? usedTokens : undefined; } @@ -2075,7 +2112,7 @@ class ClaudeContextUsageState { } buildCompactionUsageEvent(postTokens: number | undefined): AgentStreamEvent { - this.streamRequestInputTokens = undefined; + this.streamRequestInputUsage = undefined; this.streamRequestOutputTokens = undefined; this.compactedContextWindowUsedTokens = postTokens; const usage: AgentUsage = {};