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 02e93388b..9c05414dc 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -1046,7 +1046,7 @@ describe("ClaudeAgentSession context window usage", () => { const logger = createTestLogger(); interface QueryFactoryForTurnsOptions { - currentContextUsageByTurn?: Array | undefined>; + getContextUsage?: ReturnType; model?: string; } @@ -1091,8 +1091,8 @@ describe("ClaudeAgentSession context window usage", () => { const queuedMessages: Array> = []; const waiters: Array<() => void> = []; let turnIndex = 0; - let contextUsageIndex = 0; const closedRef = { value: false }; + const getContextUsage = options?.getContextUsage ?? vi.fn(async () => undefined); function wakeNextWaiter() { const waiter = waiters.shift(); @@ -1140,11 +1140,7 @@ 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; - }), + getContextUsage, supportedModels: vi.fn(async () => []), supportedCommands: vi.fn(async () => []), rewindFiles: vi.fn(async () => ({ canRewind: true })), @@ -1165,26 +1161,6 @@ describe("ClaudeAgentSession context window usage", () => { }; } - 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", @@ -1282,6 +1258,21 @@ describe("ClaudeAgentSession context window usage", () => { }; } + function createCompactBoundary(overrides: Record = {}): Record { + return { + type: "system", + subtype: "compact_boundary", + compact_metadata: { + trigger: "manual", + pre_tokens: 14_990, + post_tokens: 704, + }, + uuid: "compact-boundary-1", + session_id: "session-1", + ...overrides, + }; + } + test("passes persistSession through to the Claude SDK query options", async () => { const createResultTurn = (sessionId: string) => [ { @@ -1556,7 +1547,10 @@ describe("ClaudeAgentSession context window usage", () => { } }); - test("reports Claude's current context usage after an Agent subagent runs", async () => { + test("does not probe current context usage after an Agent subagent runs", async () => { + const getContextUsage = vi.fn(async () => { + throw new Error("getContextUsage should not be called during result handling"); + }); const session = await createSessionForTurns( [ [ @@ -1575,21 +1569,20 @@ describe("ClaudeAgentSession context window usage", () => { }), ], ], - { - currentContextUsageByTurn: [createClaudeCurrentContextUsage(12_345, 200_000)], - }, + { getContextUsage }, ); try { const result = await session.run("turn"); + expect(getContextUsage).not.toHaveBeenCalled(); expect(result.usage).toEqual({ inputTokens: 9_000, cachedInputTokens: 700, outputTokens: 400, totalCostUsd: 0.25, contextWindowMaxTokens: 200_000, - contextWindowUsedTokens: 12_345, + contextWindowUsedTokens: 175, }); } finally { await session.close(); @@ -1637,6 +1630,121 @@ describe("ClaudeAgentSession context window usage", () => { } }); + test("uses parent request usage after a real subagent tool result", async () => { + const getContextUsage = vi.fn(async () => { + throw new Error("getContextUsage should not be called during result handling"); + }); + const session = await createSessionForTurns( + [ + [ + createInitMessage(), + createMessageStartEvent({ + input_tokens: 3, + cache_creation_input_tokens: 16_999, + cache_read_input_tokens: 0, + }), + createAgentToolStartEvent(), + createMessageDeltaEvent(163), + { + type: "assistant", + parent_tool_use_id: "toolu-agent-1", + message: { + role: "assistant", + content: [{ type: "text", text: "SUBAGENT_OK" }], + usage: { + input_tokens: 3, + cache_creation_input_tokens: 1_182, + cache_read_input_tokens: 0, + output_tokens: 8, + }, + }, + uuid: "subagent-assistant-1", + session_id: "session-1", + }, + { + ...createSubagentTaskNotification(), + status: "completed", + summary: "Probe subagent test", + usage: { + total_tokens: 1_193, + tool_uses: 0, + }, + }, + { + type: "user", + parent_tool_use_id: null, + message: { + role: "user", + content: [ + { + type: "tool_result", + tool_use_id: "toolu-agent-1", + content: [ + { type: "text", text: "SUBAGENT_OK" }, + { + type: "text", + text: "agentId: subagent-1\nsubagent_tokens: 1194\ntool_uses: 0", + }, + ], + }, + ], + }, + uuid: "subagent-tool-result-1", + session_id: "session-1", + }, + createMessageStartEvent({ + input_tokens: 1, + cache_creation_input_tokens: 253, + cache_read_input_tokens: 16_999, + }), + createMessageDeltaEvent(8), + createSuccessResult({ + usage: { + input_tokens: 4, + cache_creation_input_tokens: 17_252, + cache_read_input_tokens: 16_999, + output_tokens: 171, + iterations: [ + { + input_tokens: 1, + cache_creation_input_tokens: 253, + cache_read_input_tokens: 16_999, + output_tokens: 8, + }, + ], + }, + modelUsage: { + "claude-sonnet-4-6": { + inputTokens: 7, + outputTokens: 180, + cacheReadInputTokens: 16_999, + cacheCreationInputTokens: 18_434, + contextWindow: 200_000, + }, + }, + }), + ], + ], + { getContextUsage }, + ); + + try { + const result = await session.run("turn"); + + expect(getContextUsage).not.toHaveBeenCalled(); + expect(result.usage).toEqual({ + inputTokens: 4, + cachedInputTokens: 16_999, + outputTokens: 171, + totalCostUsd: 0.25, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 17_261, + }); + } finally { + await session.close(); + } + }); + test("falls back to the active result iteration when current and stream usage are unavailable", async () => { const session = await createSessionForTurns([ [ @@ -1842,6 +1950,165 @@ describe("ClaudeAgentSession context window usage", () => { } }); + test("manual compact boundary updates context usage from post tokens", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createMessageDeltaEvent(25), + createCompactBoundary(), + createSuccessResult({ + total_cost_usd: 0.04, + usage: { + input_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + output_tokens: 0, + iterations: [], + }, + }), + ], + ]); + + try { + const events = await collectStreamEvents(session, "/compact"); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowUsedTokens: 704, + }, + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + provider: "claude", + usage: { + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + totalCostUsd: 0.04, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 704, + }, + }), + ); + } finally { + await session.close(); + } + }); + + test("zero-token stream events after compact keep post-token usage", async () => { + const session = await createSessionForTurns([ + [ + createInitMessage(), + createMessageStartEvent(), + createMessageDeltaEvent(25), + createCompactBoundary(), + createMessageStartEvent({ + input_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }), + createMessageDeltaEvent(0), + createSuccessResult({ + total_cost_usd: 0.04, + usage: { + input_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + output_tokens: 0, + iterations: [], + }, + }), + ], + ]); + + try { + const events = await collectStreamEvents(session, "/compact"); + + expect( + events.filter( + (event) => event.type === "usage_updated" && event.usage.contextWindowUsedTokens === 0, + ), + ).toEqual([]); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + provider: "claude", + usage: { + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + totalCostUsd: 0.04, + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 704, + }, + }), + ); + } finally { + await session.close(); + } + }); + + test("starting a new turn clears interrupted compact usage", async () => { + const session = await createSessionForTurns([ + [ + createSuccessResult({ + total_cost_usd: 0.04, + usage: { + input_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + output_tokens: 0, + iterations: [], + }, + }), + ], + ]); + + try { + const compactEvents = (session as unknown as TestClaudeSession).translateMessageToEvents( + createCompactBoundary(), + ); + expect(compactEvents).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowUsedTokens: 704, + }, + }), + ); + + const events = await collectStreamEvents(session, "next turn"); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + provider: "claude", + usage: expect.objectContaining({ + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + totalCostUsd: 0.04, + }), + }), + ); + expect( + events.some( + (event) => + event.type === "turn_completed" && event.usage.contextWindowUsedTokens !== undefined, + ), + ).toBe(false); + } finally { + await session.close(); + } + }); + test("result.result is surfaced as an assistant message when no model output was produced", async () => { const session = await createSessionForTest(); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index a7b9bdad5..dbfeb9186 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -11,7 +11,6 @@ import { type PermissionResult, type PermissionUpdate, type Query, - type SDKControlGetContextUsageResponse, type SDKMessage, type SDKPartialAssistantMessage, type SDKResultMessage, @@ -1682,31 +1681,6 @@ function readLegacyResultUsageTokens(usage: unknown): number | undefined { 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"; } @@ -1715,6 +1689,7 @@ class ClaudeContextUsageState { private contextWindowMaxTokens: number | undefined; private streamRequestInputTokens: number | undefined; private streamRequestOutputTokens: number | undefined; + private compactedContextWindowUsedTokens: number | undefined; private completedResultTurns = 0; constructor(initialContextWindowMaxTokens?: number) { @@ -1724,6 +1699,7 @@ class ClaudeContextUsageState { beginTurn(): void { this.streamRequestInputTokens = undefined; this.streamRequestOutputTokens = undefined; + this.compactedContextWindowUsedTokens = undefined; } setInitialContextWindowMaxTokens(contextWindowMaxTokens: number | undefined): void { @@ -1738,12 +1714,6 @@ class ClaudeContextUsageState { 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) { @@ -1774,11 +1744,7 @@ class ClaudeContextUsageState { return this.createUsageUpdatedEvent(usedTokens); } - buildResultUsage( - message: SDKResultMessage, - modelUsage: unknown, - currentContextUsage: ClaudeCurrentContextUsage | undefined, - ): AgentUsage | undefined { + buildResultUsage(message: SDKResultMessage, modelUsage: unknown): AgentUsage | undefined { try { if (!message.usage) { return undefined; @@ -1791,7 +1757,6 @@ class ClaudeContextUsageState { }; const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage); - this.recordCurrentContextUsage(currentContextUsage); if (this.contextWindowMaxTokens !== undefined) { usage.contextWindowMaxTokens = this.contextWindowMaxTokens; } else if (modelContextWindowMaxTokens !== undefined) { @@ -1802,12 +1767,13 @@ class ClaudeContextUsageState { readActiveUsageTokens(message.usage) ?? (this.completedResultTurns === 0 ? readLegacyResultUsageTokens(message.usage) : undefined); const usedTokens = - currentContextUsage?.totalTokens ?? this.streamUsedTokens() ?? activeResultUsageTokens; + this.streamUsedTokens() ?? activeResultUsageTokens ?? this.compactedContextWindowUsedTokens; if (usedTokens !== undefined) { usage.contextWindowUsedTokens = usedTokens; } return usage; } finally { + this.compactedContextWindowUsedTokens = undefined; this.completedResultTurns += 1; } } @@ -1819,7 +1785,8 @@ class ClaudeContextUsageState { ) { return undefined; } - return this.streamRequestInputTokens + this.streamRequestOutputTokens; + const usedTokens = this.streamRequestInputTokens + this.streamRequestOutputTokens; + return usedTokens > 0 ? usedTokens : undefined; } private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent { @@ -1835,6 +1802,24 @@ class ClaudeContextUsageState { usage, }; } + + buildCompactionUsageEvent(postTokens: number | undefined): AgentStreamEvent { + this.streamRequestInputTokens = undefined; + this.streamRequestOutputTokens = undefined; + this.compactedContextWindowUsedTokens = postTokens; + const usage: AgentUsage = {}; + if (this.contextWindowMaxTokens !== undefined) { + usage.contextWindowMaxTokens = this.contextWindowMaxTokens; + } + if (postTokens !== undefined) { + usage.contextWindowUsedTokens = postTokens; + } + return { + type: "usage_updated", + provider: "claude", + usage, + }; + } } class ClaudeAgentSession implements AgentSession { @@ -3282,7 +3267,7 @@ class ClaudeAgentSession implements AgentSession { if (await this.handleMissingResumedConversation(message, activeQuery)) { return true; } - await this.routeSdkMessageFromPump(message, activeQuery); + await this.routeSdkMessageFromPump(message); return false; }; const drainActiveQuery = async (): Promise => { @@ -3358,7 +3343,7 @@ class ClaudeAgentSession implements AgentSession { ); } - private async routeSdkMessageFromPump(message: SDKMessage, activeQuery: Query): Promise { + private async routeSdkMessageFromPump(message: SDKMessage): Promise { if (this.shouldSuppressStaleResult(message)) { return; } @@ -3388,12 +3373,7 @@ class ClaudeAgentSession implements AgentSession { "provider.claude.parsed_event", ); - const events = await this.buildPumpedMessageEvents( - message, - activeQuery, - identifiers.messageId, - turnId, - ); + const events = await this.buildPumpedMessageEvents(message, identifiers.messageId, turnId); if (events.length === 0) { return; @@ -3430,18 +3410,12 @@ class ClaudeAgentSession implements AgentSession { 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({ @@ -3461,18 +3435,6 @@ class ClaudeAgentSession implements AgentSession { return [...messageEvents, ...assistantTimelineEvents]; } - private async queryCurrentContextUsage( - activeQuery: Query, - ): Promise { - try { - const usage = await withTimeout(activeQuery.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, @@ -3540,7 +3502,6 @@ class ClaudeAgentSession implements AgentSession { options?: { suppressAssistantText?: boolean; suppressReasoning?: boolean; - currentContextUsage?: ClaudeCurrentContextUsage; }, ): AgentStreamEvent[] { const parentToolUseId = @@ -3591,9 +3552,7 @@ class ClaudeAgentSession implements AgentSession { this.appendStreamEventEvents(message, events, options); break; case "result": - this.appendResultEvents(message, events, { - currentContextUsage: options?.currentContextUsage, - }); + this.appendResultEvents(message, events); break; default: break; @@ -3667,6 +3626,7 @@ class ClaudeAgentSession implements AgentSession { }, provider: "claude", }); + events.push(this.contextUsage.buildCompactionUsageEvent(compactMetadata?.postTokens)); return; } if (message.subtype === "task_notification") { @@ -3794,9 +3754,8 @@ class ClaudeAgentSession implements AgentSession { private appendResultEvents( message: Extract, events: AgentStreamEvent[], - options?: { currentContextUsage?: ClaudeCurrentContextUsage }, ): void { - const usage = this.convertUsage(message, message.modelUsage, options?.currentContextUsage); + const usage = this.convertUsage(message, message.modelUsage); 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 @@ -3963,12 +3922,8 @@ class ClaudeAgentSession implements AgentSession { return null; } - private convertUsage( - message: SDKResultMessage, - modelUsage?: unknown, - currentContextUsage?: ClaudeCurrentContextUsage, - ): AgentUsage | undefined { - return this.contextUsage.buildResultUsage(message, modelUsage, currentContextUsage); + private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined { + return this.contextUsage.buildResultUsage(message, modelUsage); } private handlePermissionRequest: CanUseTool = async ( @@ -4869,7 +4824,9 @@ function hasToolLikeBlock(block?: ClaudeContentChunk | null): boolean { return type.includes("tool"); } -function readCompactionMetadata(source: unknown): { trigger?: string; preTokens?: number } | null { +function readCompactionMetadata( + source: unknown, +): { trigger?: string; preTokens?: number; postTokens?: number } | null { const sourceRecord = toObjectRecord(source); if (!sourceRecord) { return null; @@ -4887,7 +4844,9 @@ function readCompactionMetadata(source: unknown): { trigger?: string; preTokens? const trigger = typeof metadata.trigger === "string" ? metadata.trigger : undefined; const preTokensRaw = metadata.preTokens ?? metadata.pre_tokens; const preTokens = typeof preTokensRaw === "number" ? preTokensRaw : undefined; - return { trigger, preTokens }; + const postTokensRaw = metadata.postTokens ?? metadata.post_tokens; + const postTokens = typeof postTokensRaw === "number" ? postTokensRaw : undefined; + return { trigger, preTokens, postTokens }; } return null; }