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 b48493e1a..a01e12a3d 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -13,11 +13,11 @@ import { normalizeClaudeAskUserQuestionUpdatedInput, toClaudeSdkMcpConfig, } from "./agent.js"; -import type { AgentTimelineItem, AgentUsage, AgentStreamEvent } from "../../agent-sdk-types.js"; +import { streamSession } from "../test-utils/session-stream-adapter.js"; +import type { AgentSession, AgentTimelineItem, AgentStreamEvent } from "../../agent-sdk-types.js"; interface TestClaudeSession { translateMessageToEvents(message: SDKMessage): AgentStreamEvent[]; - convertUsage(message: SDKMessage): AgentUsage | undefined; } afterEach(() => { @@ -1044,6 +1044,10 @@ describe("ClaudeAgentClient.listImportableSessions", () => { describe("ClaudeAgentSession context window usage", () => { const logger = createTestLogger(); + interface QueryFactoryForTurnsOptions { + currentContextUsageByTurn?: Array | undefined>; + } + async function createSessionForTest(): Promise { const client = new ClaudeAgentClient({ logger, resolveBinary: async () => "/test/claude/bin" }); const session = await client.createSession({ @@ -1053,11 +1057,38 @@ describe("ClaudeAgentSession context window usage", () => { return session as unknown as TestClaudeSession; } - function createQueryFactoryForTurns(turns: Array>>) { + async function createSessionForTurns( + turns: Array>>, + options?: QueryFactoryForTurnsOptions, + ): Promise { + const client = new ClaudeAgentClient({ + logger, + queryFactory: createQueryFactoryForTurns(turns, options), + resolveBinary: async () => "/test/claude/bin", + }); + return await client.createSession({ + provider: "claude", + cwd: process.cwd(), + }); + } + + async function collectStreamEvents(session: AgentSession, prompt = "turn") { + const events: AgentStreamEvent[] = []; + for await (const event of streamSession(session, prompt)) { + events.push(event); + } + return events; + } + + function createQueryFactoryForTurns( + turns: Array>>, + options?: QueryFactoryForTurnsOptions, + ) { return vi.fn(({ prompt }: { prompt: AsyncIterable }) => { const queuedMessages: Array> = []; const waiters: Array<() => void> = []; let turnIndex = 0; + let contextUsageIndex = 0; const closedRef = { value: false }; function wakeNextWaiter() { @@ -1106,6 +1137,11 @@ describe("ClaudeAgentSession context window usage", () => { }), setPermissionMode: vi.fn(async () => undefined), setModel: vi.fn(async () => undefined), + getContextUsage: vi.fn(async () => { + const usage = options?.currentContextUsageByTurn?.[contextUsageIndex]; + contextUsageIndex += 1; + return usage; + }), supportedModels: vi.fn(async () => []), supportedCommands: vi.fn(async () => []), rewindFiles: vi.fn(async () => ({ canRewind: true })), @@ -1116,6 +1152,133 @@ describe("ClaudeAgentSession context window usage", () => { }); } + function createInitMessage(sessionId = "session-1"): Record { + return { + type: "system", + subtype: "init", + session_id: sessionId, + permissionMode: "default", + model: "claude-sonnet-4-6", + }; + } + + function createClaudeCurrentContextUsage( + totalTokens: number, + maxTokens: number, + ): Record { + return { + categories: [], + totalTokens, + maxTokens, + rawMaxTokens: maxTokens, + percentage: totalTokens / maxTokens, + gridRows: [], + model: "claude-sonnet-4-6", + memoryFiles: [], + mcpTools: [], + agents: [], + isAutoCompactEnabled: true, + apiUsage: null, + }; + } + + function createSuccessResult(overrides: Record = {}): Record { + return { + type: "result", + subtype: "success", + duration_ms: 100, + duration_api_ms: 75, + is_error: false, + num_turns: 1, + result: "done", + stop_reason: null, + total_cost_usd: 0.25, + usage: { + input_tokens: 10, + cache_read_input_tokens: 5, + output_tokens: 7, + }, + modelUsage: { + "claude-sonnet-4-6": { contextWindow: 200_000 }, + }, + permission_denials: [], + uuid: "result-1", + session_id: "session-1", + ...overrides, + }; + } + + function createMessageStartEvent( + usage: Record = { + input_tokens: 100, + cache_creation_input_tokens: 20, + cache_read_input_tokens: 30, + }, + ): Record { + return { + type: "stream_event", + event: { + type: "message_start", + message: { usage }, + }, + session_id: "session-1", + }; + } + + function createMessageDeltaEvent(outputTokens: number): Record { + return { + type: "stream_event", + event: { + type: "message_delta", + usage: { + output_tokens: outputTokens, + }, + }, + session_id: "session-1", + }; + } + + function createAgentToolStartEvent(): Record { + return { + type: "stream_event", + event: { + type: "content_block_start", + index: 0, + content_block: { + type: "tool_use", + id: "toolu-agent-1", + name: "Agent", + input: { + description: "Check something in a subagent", + prompt: "Return a short answer", + }, + }, + }, + session_id: "session-1", + }; + } + + function createSubagentTaskNotification(): Record { + return { + type: "system", + subtype: "task_notification", + uuid: "task-note-1", + task_id: "task-1", + tool_use_id: "toolu-agent-1", + status: "running", + summary: "Subagent is working", + usage: { + total_tokens: 18_876, + tool_uses: 1, + duration_ms: 50, + input_tokens: 12_000, + cache_read_input_tokens: 6_000, + output_tokens: 876, + }, + session_id: "session-1", + }; + } + test("passes persistSession through to the Claude SDK query options", async () => { const createResultTurn = (sessionId: string) => [ { @@ -1389,357 +1552,152 @@ describe("ClaudeAgentSession context window usage", () => { } }); - test("convertUsage includes contextWindowMaxTokens and derives used tokens from result usage as initial fallback", async () => { - const session = await createSessionForTest(); - - const usage = session.convertUsage( + test("reports Claude's current context usage after an Agent subagent runs", async () => { + const session = await createSessionForTurns( + [ + [ + createInitMessage(), + createMessageStartEvent(), + createAgentToolStartEvent(), + createSubagentTaskNotification(), + createMessageDeltaEvent(25), + createSuccessResult({ + usage: { + input_tokens: 9_000, + cache_creation_input_tokens: 300, + cache_read_input_tokens: 700, + output_tokens: 400, + }, + }), + ], + ], { - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }, - { - "claude-sonnet-4-6": { contextWindow: 200_000 }, - "claude-opus-4-6": { contextWindow: 1_000_000 }, + currentContextUsageByTurn: [createClaudeCurrentContextUsage(12_345, 200_000)], }, ); - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowMaxTokens: 1_000_000, - contextWindowUsedTokens: 22, - }); + try { + const result = await session.run("turn"); + + expect(result.usage).toEqual({ + inputTokens: 9_000, + cachedInputTokens: 700, + outputTokens: 400, + totalCostUsd: 0.25, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 12_345, + }); + } finally { + await session.close(); + } }); - test("contextWindowUsedTokens falls back to result usage when no task_progress was received", async () => { - const session = await createSessionForTest(); + test("does not report task notification tokens as parent context usage", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createAgentToolStartEvent(), + createSubagentTaskNotification(), + { + type: "system", + subtype: "task_progress", + task_id: "task-1", + description: "Subagent progress", + usage: { + total_tokens: 9_999, + tool_uses: 1, + duration_ms: 50, + }, + uuid: "task-progress-1", + session_id: "session-1", + }, + createMessageDeltaEvent(25), + createSuccessResult(), + ], + ]); - const usage = session.convertUsage({ - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_creation_input_tokens: 3, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }); + try { + const result = await session.run("turn"); - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowUsedTokens: 25, - }); - }); - - test("contextWindowUsedTokens is populated from task_progress usage data", async () => { - const session = await createSessionForTest(); - - session.translateMessageToEvents({ - type: "system", - subtype: "task_progress", - task_id: "task-1", - description: "Processing", - usage: { - total_tokens: 999, - tool_uses: 1, - duration_ms: 50, - input_tokens: 345, - cache_read_input_tokens: 55, - }, - uuid: "task-progress-1", - session_id: "session-1", - }); - - const events = session.translateMessageToEvents({ - type: "result", - subtype: "success", - duration_ms: 100, - duration_api_ms: 75, - is_error: false, - num_turns: 1, - result: "done", - stop_reason: null, - total_cost_usd: 0.25, - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - modelUsage: { - "claude-sonnet-4-6": { contextWindow: 200_000 }, - }, - permission_denials: [], - uuid: "result-1", - session_id: "session-1", - }); - - expect(events).toContainEqual({ - type: "turn_completed", - provider: "claude", - usage: { + expect(result.usage).toEqual({ inputTokens: 10, cachedInputTokens: 5, outputTokens: 7, totalCostUsd: 0.25, contextWindowMaxTokens: 200_000, - contextWindowUsedTokens: 999, - }, - }); - }); - - test("task_progress emits a usage_updated event", async () => { - const session = await createSessionForTest(); - - const events = session.translateMessageToEvents({ - type: "system", - subtype: "task_progress", - task_id: "task-1", - description: "Processing", - usage: { - total_tokens: 999, - tool_uses: 1, - duration_ms: 50, - }, - uuid: "task-progress-1", - session_id: "session-1", - }); - - expect(events).toContainEqual({ - type: "usage_updated", - provider: "claude", - usage: { - contextWindowUsedTokens: 999, - }, - }); - }); - - test("task_notification emits a usage_updated event", async () => { - const session = await createSessionForTest(); - - const events = session.translateMessageToEvents({ - type: "system", - subtype: "task_notification", - uuid: "task-note-1", - task_id: "task-1", - status: "running", - summary: "Background task still running", - usage: { - total_tokens: 777, - tool_uses: 1, - duration_ms: 50, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - expect(events).toContainEqual({ - type: "usage_updated", - provider: "claude", - usage: { - contextWindowUsedTokens: 777, - }, - }); - }); - - test("message_start stream events emit usage_updated with per-request usage", async () => { - const session = await createSessionForTest(); - - const events = session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_start", - message: { - usage: { - input_tokens: 100, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 30, - }, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - expect(events).toContainEqual({ - type: "usage_updated", - provider: "claude", - usage: { - contextWindowUsedTokens: 150, - }, - }); - }); - - test("message_delta stream events update per-request usage", async () => { - const session = await createSessionForTest(); - - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_start", - message: { - usage: { - input_tokens: 100, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 30, - }, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - const events = session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_delta", - usage: { - output_tokens: 25, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - expect(events).toContainEqual({ - type: "usage_updated", - provider: "claude", - usage: { contextWindowUsedTokens: 175, - }, - }); + }); + } finally { + await session.close(); + } }); - test("task_progress usage takes priority over derived result usage", async () => { - const session = await createSessionForTest(); - - session.translateMessageToEvents({ - type: "system", - subtype: "task_progress", - task_id: "task-1", - description: "Processing", - usage: { - total_tokens: 999, - tool_uses: 1, - duration_ms: 50, - input_tokens: 345, - cache_read_input_tokens: 55, - }, - uuid: "task-progress-1", - session_id: "session-1", - }); - - const usage = session.convertUsage({ - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_creation_input_tokens: 3, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }); - - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowUsedTokens: 999, - }); - }); - - test("contextWindowUsedTokens persists across turns from last task_progress", async () => { - const queryFactory = createQueryFactoryForTurns([ + test("falls back to the active result iteration when current and stream usage are unavailable", async () => { + const session = await createSessionForTurns([ [ - { - type: "system", - subtype: "init", - session_id: "session-1", - permissionMode: "default", - model: "claude-sonnet-4-6", - }, - { - type: "system", - subtype: "task_progress", - task_id: "task-1", - description: "Processing", + createInitMessage(), + createSuccessResult({ usage: { - total_tokens: 999, - tool_uses: 1, - duration_ms: 50, - input_tokens: 345, - cache_read_input_tokens: 55, + input_tokens: 5_000, + cache_creation_input_tokens: 500, + cache_read_input_tokens: 600, + output_tokens: 700, + iterations: [ + { + input_tokens: 100, + cache_creation_input_tokens: 10, + cache_read_input_tokens: 20, + output_tokens: 30, + }, + { + input_tokens: 2, + cache_creation_input_tokens: 10, + cache_read_input_tokens: 100, + output_tokens: 5, + }, + ], }, - uuid: "task-progress-1", - session_id: "session-1", - }, - { - type: "result", - subtype: "success", - duration_ms: 100, - duration_api_ms: 75, - is_error: false, - num_turns: 1, - result: "done", - stop_reason: null, - total_cost_usd: 0.25, - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - modelUsage: { - "claude-sonnet-4-6": { contextWindow: 200_000 }, - }, - permission_denials: [], - uuid: "result-1", - session_id: "session-1", - }, + }), + ], + ]); + + try { + const result = await session.run("turn"); + + expect(result.usage).toEqual({ + inputTokens: 5_000, + cachedInputTokens: 600, + outputTokens: 700, + totalCostUsd: 0.25, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 117, + }); + } finally { + await session.close(); + } + }); + + test("does not use aggregate result totals after the first result turn", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createMessageDeltaEvent(25), + createSuccessResult(), ], [ - { - type: "result", - subtype: "success", - duration_ms: 110, - duration_api_ms: 80, - is_error: false, - num_turns: 1, - result: "still done", - stop_reason: null, + createSuccessResult({ total_cost_usd: 0.1, usage: { - input_tokens: 11, - cache_creation_input_tokens: 3, - cache_read_input_tokens: 6, - output_tokens: 8, - }, - modelUsage: { - "claude-sonnet-4-6": { contextWindow: 200_000 }, + input_tokens: 1_000, + cache_read_input_tokens: 200, + output_tokens: 300, }, - permission_denials: [], uuid: "result-2", - session_id: "session-1", - }, + }), ], ]); - const client = new ClaudeAgentClient({ - logger, - queryFactory, - resolveBinary: async () => "/test/claude/bin", - }); - const session = await client.createSession({ - provider: "claude", - cwd: process.cwd(), - }); try { const firstTurn = await session.run("turn 1"); @@ -1751,175 +1709,109 @@ describe("ClaudeAgentSession context window usage", () => { outputTokens: 7, totalCostUsd: 0.25, contextWindowMaxTokens: 200_000, - contextWindowUsedTokens: 999, + contextWindowUsedTokens: 175, }); - // Turn 2 has no task_progress, so contextWindowUsedTokens retains the - // last known value from turn 1 rather than deriving from accumulated - // result.usage (which would be incorrect — those are session-level totals). expect(secondTurn.usage).toEqual({ - inputTokens: 11, - cachedInputTokens: 6, - outputTokens: 8, + inputTokens: 1_000, + cachedInputTokens: 200, + outputTokens: 300, totalCostUsd: 0.1, contextWindowMaxTokens: 200_000, - contextWindowUsedTokens: 999, }); } finally { await session.close(); } }); - test("convertUsage derives used tokens from result usage as fallback when task_progress is missing", async () => { - const session = await createSessionForTest(); + test("message_start stream events emit usage_updated with per-request usage", async () => { + const session = await createSessionForTurns([ + [createInitMessage(), createMessageStartEvent(), createSuccessResult()], + ]); - const usage = session.convertUsage({ - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }); + try { + const events = await collectStreamEvents(session); - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowUsedTokens: 22, - }); + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowUsedTokens: 150, + }, + }), + ); + } finally { + await session.close(); + } }); - test("convertUsage uses per-request stream usage when no task_progress is available", async () => { - const session = await createSessionForTest(); + test("message_delta stream events update per-request usage", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createMessageDeltaEvent(25), + createSuccessResult(), + ], + ]); - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_start", - message: { + try { + const events = await collectStreamEvents(session); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", usage: { - input_tokens: 100, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 30, + contextWindowUsedTokens: 175, }, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_delta", - usage: { - output_tokens: 25, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - const usage = session.convertUsage({ - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }); - - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowUsedTokens: 175, - }); + }), + ); + } finally { + await session.close(); + } }); test("per-request stream usage is not cumulative across API calls in a turn", async () => { - const session = await createSessionForTest(); + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createMessageDeltaEvent(25), + createMessageStartEvent({ + input_tokens: 40, + cache_creation_input_tokens: 5, + cache_read_input_tokens: 10, + }), + createMessageDeltaEvent(7), + createSuccessResult(), + ], + ]); - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_start", - message: { + try { + const events = await collectStreamEvents(session); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", usage: { - input_tokens: 100, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 30, + contextWindowUsedTokens: 55, }, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_delta", - usage: { - output_tokens: 25, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - const secondStartEvents = session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_start", - message: { + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", usage: { - input_tokens: 40, - cache_creation_input_tokens: 5, - cache_read_input_tokens: 10, + contextWindowUsedTokens: 62, }, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - expect(secondStartEvents).toContainEqual({ - type: "usage_updated", - provider: "claude", - usage: { - contextWindowUsedTokens: 55, - }, - }); - - session.translateMessageToEvents({ - type: "stream_event", - event: { - type: "message_delta", - usage: { - output_tokens: 7, - }, - }, - session_id: "session-1", - } as unknown as SDKMessage); - - const usage = session.convertUsage({ - type: "result", - subtype: "success", - usage: { - input_tokens: 10, - cache_read_input_tokens: 5, - output_tokens: 7, - }, - total_cost_usd: 0.12, - }); - - expect(usage).toEqual({ - inputTokens: 10, - cachedInputTokens: 5, - outputTokens: 7, - totalCostUsd: 0.12, - contextWindowUsedTokens: 62, - }); + }), + ); + } finally { + await session.close(); + } }); test("result.result is surfaced as an assistant message when no model output was produced", async () => { diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 2dc8e6745..c6fd80ffc 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -11,9 +11,9 @@ import { type PermissionResult, type PermissionUpdate, type Query, + type SDKControlGetContextUsageResponse, type SDKMessage, type SDKPartialAssistantMessage, - type SDKTaskProgressMessage, type SDKResultMessage, type SDKSystemMessage, type SDKUserMessage, @@ -1603,27 +1603,6 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined { return maxContextWindow; } -function readUsageTotalTokens(usage: unknown): number | undefined { - if (!usage || typeof usage !== "object") { - return undefined; - } - const totalTokens = (usage as { total_tokens?: unknown }).total_tokens; - if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) { - return undefined; - } - return totalTokens; -} - -function readContextWindowUsedTokensFromTaskProgress( - message: SDKTaskProgressMessage, -): number | undefined { - return readUsageTotalTokens(message.usage); -} - -function readUsageFromTaskNotification(message: { usage?: unknown }): number | undefined { - return readUsageTotalTokens(message.usage); -} - function readStreamRequestInputTokens(event: Record): number | undefined { const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage); if (!messageUsage) { @@ -1658,6 +1637,203 @@ function readStreamRequestOutputTokens(event: Record): number | return outputTokens; } +function readLastUsageIteration(usage: unknown): Record | undefined { + const iterations = toObjectRecord(usage)?.iterations; + if (!Array.isArray(iterations)) { + return undefined; + } + for (let index = iterations.length - 1; index >= 0; index -= 1) { + const candidate = toObjectRecord(iterations[index]); + if (candidate) { + return candidate; + } + } + return undefined; +} + +function readUsageTokenTotal(usage: Record): number | undefined { + const usageWithCacheCreation = usage as typeof usage & { + cache_creation_input_tokens?: unknown; + }; + const inputTokens = + typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens) + ? usage.input_tokens + : 0; + const cacheCreationInputTokens = + typeof usageWithCacheCreation.cache_creation_input_tokens === "number" && + Number.isFinite(usageWithCacheCreation.cache_creation_input_tokens) + ? usageWithCacheCreation.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; + const outputTokens = + typeof usage.output_tokens === "number" && Number.isFinite(usage.output_tokens) + ? usage.output_tokens + : 0; + const total = inputTokens + cacheCreationInputTokens + cacheReadInputTokens + outputTokens; + return total > 0 ? total : undefined; +} + +function readActiveUsageTokens(usage: unknown): number | undefined { + const activeUsage = readLastUsageIteration(usage); + return activeUsage ? readUsageTokenTotal(activeUsage) : undefined; +} + +function readLegacyResultUsageTokens(usage: unknown): number | undefined { + const usageRecord = toObjectRecord(usage); + return usageRecord ? readUsageTokenTotal(usageRecord) : undefined; +} + +interface ClaudeCurrentContextUsage { + totalTokens: number; + maxTokens?: number; +} + +function readCurrentContextUsage( + value: SDKControlGetContextUsageResponse | unknown, +): ClaudeCurrentContextUsage | undefined { + const record = toObjectRecord(value); + if (!record) { + return undefined; + } + const totalTokens = record.totalTokens; + if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) { + return undefined; + } + const maxTokens = record.maxTokens; + return { + totalTokens, + ...(typeof maxTokens === "number" && Number.isFinite(maxTokens) && maxTokens > 0 + ? { maxTokens } + : {}), + }; +} + +function isClaudeSubagentToolName(name: string | undefined): boolean { + return name === "Task" || name === "Agent"; +} + +class ClaudeContextUsageState { + private contextWindowMaxTokens: number | undefined; + private streamRequestInputTokens: number | undefined; + private streamRequestOutputTokens: number | undefined; + private completedResultTurns = 0; + + beginTurn(): void { + this.streamRequestInputTokens = undefined; + this.streamRequestOutputTokens = undefined; + } + + recordModelUsage(modelUsage: unknown): number | undefined { + const contextWindowMaxTokens = extractContextWindowSize(modelUsage); + if (contextWindowMaxTokens !== undefined) { + this.contextWindowMaxTokens = contextWindowMaxTokens; + } + return this.contextWindowMaxTokens; + } + + recordCurrentContextUsage(usage: ClaudeCurrentContextUsage | undefined): void { + if (usage?.maxTokens !== undefined) { + this.contextWindowMaxTokens = usage.maxTokens; + } + } + + buildStreamUsageEvent(event: unknown): AgentStreamEvent | null { + const streamEvent = toObjectRecord(event); + if (!streamEvent) { + return null; + } + const eventType = readTrimmedString(streamEvent.type); + if (eventType === "message_start") { + const inputTokens = readStreamRequestInputTokens(streamEvent); + if (typeof inputTokens !== "number") { + return null; + } + this.streamRequestInputTokens = inputTokens; + this.streamRequestOutputTokens = 0; + } else if (eventType === "message_delta") { + const outputTokens = readStreamRequestOutputTokens(streamEvent); + if (typeof outputTokens !== "number") { + return null; + } + this.streamRequestOutputTokens = outputTokens; + } else { + return null; + } + + const usedTokens = this.streamUsedTokens(); + if (usedTokens === undefined) { + return null; + } + return this.createUsageUpdatedEvent(usedTokens); + } + + buildResultUsage( + message: SDKResultMessage, + modelUsage: unknown, + currentContextUsage: ClaudeCurrentContextUsage | undefined, + ): AgentUsage | undefined { + try { + if (!message.usage) { + return undefined; + } + const usage: AgentUsage = { + inputTokens: message.usage.input_tokens, + cachedInputTokens: message.usage.cache_read_input_tokens, + outputTokens: message.usage.output_tokens, + totalCostUsd: message.total_cost_usd, + }; + + const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage); + this.recordCurrentContextUsage(currentContextUsage); + if (this.contextWindowMaxTokens !== undefined) { + usage.contextWindowMaxTokens = this.contextWindowMaxTokens; + } else if (modelContextWindowMaxTokens !== undefined) { + usage.contextWindowMaxTokens = modelContextWindowMaxTokens; + } + + const activeResultUsageTokens = + readActiveUsageTokens(message.usage) ?? + (this.completedResultTurns === 0 ? readLegacyResultUsageTokens(message.usage) : undefined); + const usedTokens = + currentContextUsage?.totalTokens ?? this.streamUsedTokens() ?? activeResultUsageTokens; + if (usedTokens !== undefined) { + usage.contextWindowUsedTokens = usedTokens; + } + return usage; + } finally { + this.completedResultTurns += 1; + } + } + + private streamUsedTokens(): number | undefined { + if ( + typeof this.streamRequestInputTokens !== "number" || + typeof this.streamRequestOutputTokens !== "number" + ) { + return undefined; + } + return this.streamRequestInputTokens + this.streamRequestOutputTokens; + } + + private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent { + const usage: AgentUsage = { + contextWindowUsedTokens, + }; + if (this.contextWindowMaxTokens !== undefined) { + usage.contextWindowMaxTokens = this.contextWindowMaxTokens; + } + return { + type: "usage_updated", + provider: "claude", + usage, + }; + } +} + class ClaudeAgentSession implements AgentSession { readonly provider = "claude" as const; readonly capabilities = CLAUDE_CAPABILITIES; @@ -1703,10 +1879,7 @@ class ClaudeAgentSession implements AgentSession { private pendingInterruptAbort = false; private foregroundHasVisibleActivity = false; private activeTurnHasAssistantText = false; - private lastContextWindowUsedTokens: number | undefined; - private lastContextWindowMaxTokens: number | undefined; - private lastStreamRequestInputTokens: number | undefined; - private lastStreamRequestOutputTokens: number | undefined; + private readonly contextUsage = new ClaudeContextUsageState(); private userMessageIds: string[] = []; private readonly emittedUserMessageIds = new Set(); private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = []; @@ -1840,6 +2013,7 @@ class ClaudeAgentSession implements AgentSession { this.activeForegroundTurnId = turnId; this.foregroundHasVisibleActivity = false; this.activeTurnHasAssistantText = false; + this.contextUsage.beginTurn(); this.transitionTurnState("foreground", "foreground turn started"); this.clearRecentStderr(); @@ -3005,6 +3179,7 @@ class ClaudeAgentSession implements AgentSession { id: this.createTurnId("autonomous"), }; this.activeTurnHasAssistantText = false; + this.contextUsage.beginTurn(); this.notifySubscribers({ type: "turn_started", provider: "claude" }); this.syncTurnState("autonomous turn started"); } @@ -3098,7 +3273,7 @@ class ClaudeAgentSession implements AgentSession { if (await this.handleMissingResumedConversation(message, activeQuery)) { return true; } - this.routeSdkMessageFromPump(message); + await this.routeSdkMessageFromPump(message, activeQuery); return false; }; const drainActiveQuery = async (): Promise => { @@ -3174,7 +3349,7 @@ class ClaudeAgentSession implements AgentSession { ); } - private routeSdkMessageFromPump(message: SDKMessage): void { + private async routeSdkMessageFromPump(message: SDKMessage, activeQuery: Query): Promise { if (this.shouldSuppressStaleResult(message)) { return; } @@ -3204,26 +3379,12 @@ class ClaudeAgentSession implements AgentSession { "provider.claude.parsed_event", ); - const messageEvents = this.translateMessageToEvents(message, { - suppressAssistantText: true, - suppressReasoning: true, - }); - const assistantTimelineEvents = this.timelineAssembler - .consume({ - message, - runId: turnId, - messageIdHint: identifiers.messageId, - }) - .map( - (item) => - ({ - type: "timeline", - item, - provider: "claude", - }) satisfies AgentStreamEvent, - ); - - const events = [...messageEvents, ...assistantTimelineEvents]; + const events = await this.buildPumpedMessageEvents( + message, + activeQuery, + identifiers.messageId, + turnId, + ); if (events.length === 0) { return; @@ -3258,6 +3419,57 @@ class ClaudeAgentSession implements AgentSession { this.dispatchEvents(events); } + private async buildPumpedMessageEvents( + message: SDKMessage, + activeQuery: Query, + messageIdHint: string | null, + turnId: string | null, + ): Promise { + const currentContextUsage = + message.type === "result" && message.subtype === "success" + ? await this.queryCurrentContextUsage(activeQuery) + : undefined; + const messageEvents = this.translateMessageToEvents(message, { + suppressAssistantText: true, + suppressReasoning: true, + currentContextUsage, + }); + const assistantTimelineEvents = this.timelineAssembler + .consume({ + message, + runId: turnId, + messageIdHint, + }) + .map( + (item) => + ({ + type: "timeline", + item, + provider: "claude", + }) satisfies AgentStreamEvent, + ); + + return [...messageEvents, ...assistantTimelineEvents]; + } + + private async queryCurrentContextUsage( + activeQuery: Query, + ): Promise { + const query = activeQuery as Query & { + getContextUsage?: () => Promise; + }; + if (typeof query.getContextUsage !== "function") { + return undefined; + } + try { + const usage = await withTimeout(query.getContextUsage(), 3_000, "timeout"); + return readCurrentContextUsage(usage); + } catch (error) { + this.logger.debug({ err: error }, "Claude context usage query failed"); + return undefined; + } + } + private async handleMissingResumedConversation( message: SDKMessage, activeQuery: Query, @@ -3325,6 +3537,7 @@ class ClaudeAgentSession implements AgentSession { options?: { suppressAssistantText?: boolean; suppressReasoning?: boolean; + currentContextUsage?: ClaudeCurrentContextUsage; }, ): AgentStreamEvent[] { const parentToolUseId = @@ -3375,7 +3588,9 @@ class ClaudeAgentSession implements AgentSession { this.appendStreamEventEvents(message, events, options); break; case "result": - this.appendResultEvents(message, events); + this.appendResultEvents(message, events, { + currentContextUsage: options?.currentContextUsage, + }); break; default: break; @@ -3456,11 +3671,7 @@ class ClaudeAgentSession implements AgentSession { return; } if (message.subtype === "task_progress") { - this.lastContextWindowUsedTokens = - readContextWindowUsedTokensFromTaskProgress(message) ?? this.lastContextWindowUsedTokens; - if (typeof this.lastContextWindowUsedTokens === "number") { - events.push(this.createUsageUpdatedEvent(this.lastContextWindowUsedTokens)); - } + return; } } @@ -3470,12 +3681,12 @@ class ClaudeAgentSession implements AgentSession { ): void { // TODO: subagent timelines are best-effort. Subagent task_notifications // arrive without parent_tool_use_id but with tool_use_id pointing at the - // parent's Task call, so they slip past the sidechain router and pollute + // the parent's subagent tool call, so they slip past the sidechain router and pollute // the parent timeline. Drop them here; eventually thread them into the - // parent Task tool call's sub_agent log instead. + // parent tool call's sub_agent log instead. const taskUseId = message.tool_use_id; const cachedTool = taskUseId ? this.toolUseCache.get(taskUseId) : undefined; - if (cachedTool?.name === "Task") { + if (isClaudeSubagentToolName(cachedTool?.name)) { return; } const taskNotificationItem = mapTaskNotificationSystemRecordToToolCall(message); @@ -3486,11 +3697,6 @@ class ClaudeAgentSession implements AgentSession { provider: "claude", }); } - const usage = readUsageFromTaskNotification(message); - if (typeof usage === "number") { - this.lastContextWindowUsedTokens = usage; - events.push(this.createUsageUpdatedEvent(usage)); - } } private appendUserMessageEvents( @@ -3569,7 +3775,7 @@ class ClaudeAgentSession implements AgentSession { events: AgentStreamEvent[], options: { suppressAssistantText?: boolean; suppressReasoning?: boolean } | undefined, ): void { - const usageUpdatedEvent = this.trackStreamEventUsage(message.event); + const usageUpdatedEvent = this.contextUsage.buildStreamUsageEvent(message.event); if (usageUpdatedEvent) { events.push(usageUpdatedEvent); } @@ -3585,8 +3791,9 @@ class ClaudeAgentSession implements AgentSession { private appendResultEvents( message: Extract, events: AgentStreamEvent[], + options?: { currentContextUsage?: ClaudeCurrentContextUsage }, ): void { - const usage = this.convertUsage(message, message.modelUsage); + const usage = this.convertUsage(message, message.modelUsage, options?.currentContextUsage); if (message.subtype === "success") { // Built-in slash commands (e.g. /voice, /usage, "Unknown command: …") // run client-side in the Claude CLI with no model turn — output_tokens @@ -3753,99 +3960,12 @@ class ClaudeAgentSession implements AgentSession { return null; } - private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined { - if (!message.usage) { - return undefined; - } - const usage: AgentUsage = { - inputTokens: message.usage.input_tokens, - cachedInputTokens: message.usage.cache_read_input_tokens, - outputTokens: message.usage.output_tokens, - totalCostUsd: message.total_cost_usd, - }; - const contextWindowMaxTokens = extractContextWindowSize(modelUsage ?? message.modelUsage); - if (contextWindowMaxTokens !== undefined) { - this.lastContextWindowMaxTokens = contextWindowMaxTokens; - usage.contextWindowMaxTokens = contextWindowMaxTokens; - } else if (this.lastContextWindowMaxTokens !== undefined) { - usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens; - } - if (typeof this.lastContextWindowUsedTokens === "number") { - // task_progress.total_tokens is the accurate context window fill level. - // Prefer it over result.usage which contains accumulated session totals. - usage.contextWindowUsedTokens = this.lastContextWindowUsedTokens; - } else if ( - typeof this.lastStreamRequestInputTokens === "number" && - typeof this.lastStreamRequestOutputTokens === "number" - ) { - usage.contextWindowUsedTokens = - this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens; - } else if (message.usage) { - // Fallback: derive from result.usage when no task_progress has been - // received yet. These values are accumulated across all API calls, but - // for the first turn they equal the per-call values so the estimate is - // reasonable. Once a task_progress arrives it takes over permanently. - const usageWithCacheCreation = message.usage as typeof message.usage & { - cache_creation_input_tokens?: number; - }; - const derived = - (message.usage.input_tokens ?? 0) + - (usageWithCacheCreation.cache_creation_input_tokens ?? 0) + - (message.usage.cache_read_input_tokens ?? 0) + - (message.usage.output_tokens ?? 0); - if (Number.isFinite(derived) && derived > 0) { - usage.contextWindowUsedTokens = derived; - } - } - return usage; - } - - private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent { - const usage: AgentUsage = { - contextWindowUsedTokens, - }; - if (this.lastContextWindowMaxTokens !== undefined) { - usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens; - } - return { - type: "usage_updated", - provider: "claude", - usage, - }; - } - - private trackStreamEventUsage(event: unknown): AgentStreamEvent | null { - const streamEvent = toObjectRecord(event); - if (!streamEvent) { - return null; - } - const eventType = readTrimmedString(streamEvent.type); - if (eventType === "message_start") { - const inputTokens = readStreamRequestInputTokens(streamEvent); - if (typeof inputTokens !== "number") { - return null; - } - this.lastStreamRequestInputTokens = inputTokens; - this.lastStreamRequestOutputTokens = 0; - } else if (eventType === "message_delta") { - const outputTokens = readStreamRequestOutputTokens(streamEvent); - if (typeof outputTokens !== "number") { - return null; - } - this.lastStreamRequestOutputTokens = outputTokens; - } else { - return null; - } - - if ( - typeof this.lastStreamRequestInputTokens !== "number" || - typeof this.lastStreamRequestOutputTokens !== "number" - ) { - return null; - } - return this.createUsageUpdatedEvent( - this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens, - ); + private convertUsage( + message: SDKResultMessage, + modelUsage?: unknown, + currentContextUsage?: ClaudeCurrentContextUsage, + ): AgentUsage | undefined { + return this.contextUsage.buildResultUsage(message, modelUsage, currentContextUsage); } private handlePermissionRequest: CanUseTool = async (