diff --git a/packages/server/src/server/agent/providers/claude/agent.context-window.test.ts b/packages/server/src/server/agent/providers/claude/agent.context-window.test.ts index d14b27c714..84f31779e4 100644 --- a/packages/server/src/server/agent/providers/claude/agent.context-window.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.context-window.test.ts @@ -6,17 +6,28 @@ import type { ClaudeQueryFactory } from "./query.js"; // A query dependency that speaks the same stream interface as Claude Code. // Each request repeats init so a second turn exercises runtime-limit retention. -function contextQuery( - runtimeModel: string, - reportedLimit: number, - modelWindowsByTurn?: Array>, -): ClaudeQueryFactory { +function contextQuery({ + runtimeModel, + reportedLimit, + modelWindowsByTurn, +}: { + runtimeModel: string; + reportedLimit: number; + modelWindowsByTurn?: Array>; +}): ClaudeQueryFactory { return ({ prompt }) => { let model = runtimeModel; let turn = 0; const messages = (async function* () { if (typeof prompt === "string") throw new Error("Expected streaming input"); for await (const _message of prompt) { + const modelWindows = modelWindowsByTurn?.[turn] ?? { [model]: reportedLimit }; + turn += 1; + const usageEntries = Object.entries(modelWindows).map(([id, contextWindow]) => [ + id, + { contextWindow }, + ]); + const modelUsage = Object.fromEntries(usageEntries); yield { type: "system", subtype: "init", @@ -46,11 +57,7 @@ function contextQuery( stop_reason: null, total_cost_usd: 0, usage: { input_tokens: 249_950, output_tokens: 50, cache_read_input_tokens: 0 }, - modelUsage: Object.fromEntries( - Object.entries(modelWindowsByTurn?.[turn++] ?? { [model]: reportedLimit }).map( - ([id, contextWindow]) => [id, { contextWindow }], - ), - ), + modelUsage, permission_denials: [], uuid: "context-result", }; @@ -80,7 +87,7 @@ test.each([ async ({ selected, runtime, expected }) => { const client = new ClaudeAgentClient({ logger: createTestLogger(), - queryFactory: contextQuery(runtime, 800_000), + queryFactory: contextQuery({ runtimeModel: runtime, reportedLimit: 800_000 }), resolveBinary: async () => process.execPath, }); const session = await client.createSession({ @@ -126,11 +133,15 @@ test("completed turns use the active model's window after switching both ways", const client = new ClaudeAgentClient({ logger: createTestLogger(), resolveBinary: async () => process.execPath, - queryFactory: contextQuery(haiku, 200_000, [ - { [haiku]: 200_000 }, - { [haiku]: 200_000, [opus]: 1_000_000 }, - { [haiku]: 200_000, [opus]: 1_000_000 }, - ]), + queryFactory: contextQuery({ + runtimeModel: haiku, + reportedLimit: 200_000, + modelWindowsByTurn: [ + { [haiku]: 200_000 }, + { [haiku]: 200_000, [opus]: 1_000_000 }, + { [haiku]: 200_000, [opus]: 1_000_000 }, + ], + }), }); const session = await client.createSession({ provider: "claude", @@ -186,7 +197,11 @@ test.each([ const client = new ClaudeAgentClient({ logger: createTestLogger(), resolveBinary: async () => process.execPath, - queryFactory: contextQuery(model, expected, [windows]), + queryFactory: contextQuery({ + runtimeModel: model, + reportedLimit: expected, + modelWindowsByTurn: [windows], + }), }); const session = await client.createSession({ provider: "claude", cwd: process.cwd(), model }); try { diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index f3f7aaa5a5..a135ac7909 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -2044,7 +2044,6 @@ class ClaudeContextUsageState { buildResultUsage( message: SDKResultMessage, - modelUsage: unknown, modelId: string | null | undefined, ): AgentUsage | undefined { try { @@ -2058,10 +2057,7 @@ class ClaudeContextUsageState { totalCostUsd: message.total_cost_usd, }; - const modelContextWindowMaxTokens = this.recordModelUsage( - modelUsage ?? message.modelUsage, - modelId, - ); + const modelContextWindowMaxTokens = this.recordModelUsage(message.modelUsage, modelId); if (this.contextWindowMaxTokens !== undefined) { usage.contextWindowMaxTokens = this.contextWindowMaxTokens; } else if (modelContextWindowMaxTokens !== undefined) { @@ -4693,7 +4689,7 @@ class ClaudeAgentSession implements AgentSession { message: Extract, events: AgentStreamEvent[], ): void { - const usage = this.convertUsage(message, message.modelUsage); + const usage = this.convertUsage(message); if (message.subtype === "success") { events.push(...this.sidechainTracker.finishAll("completed")); // Built-in slash commands (e.g. /voice, /usage, "Unknown command: …") @@ -4875,12 +4871,8 @@ class ClaudeAgentSession implements AgentSession { return null; } - private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined { - return this.contextUsage.buildResultUsage( - message, - modelUsage, - this.lastRuntimeModel ?? this.config.model, - ); + private convertUsage(message: SDKResultMessage): AgentUsage | undefined { + return this.contextUsage.buildResultUsage(message, this.lastRuntimeModel ?? this.config.model); } private handlePermissionRequest: CanUseTool = async (