diff --git a/packages/server/src/server/agent/providers/claude/agent.context-switch.real.e2e.test.ts b/packages/server/src/server/agent/providers/claude/agent.context-switch.real.e2e.test.ts new file mode 100644 index 0000000000..f05a6ee271 --- /dev/null +++ b/packages/server/src/server/agent/providers/claude/agent.context-switch.real.e2e.test.ts @@ -0,0 +1,42 @@ +import pino from "pino"; +import { expect, test } from "vitest"; +import { ClaudeAgentClient } from "./agent.js"; +import { collectSessionTurnEvents } from "../test-utils/session-stream-adapter.js"; + +// Explicit real-provider test: run separately from the default suite. +test("uses the active model context window after switching from a 1M model", async () => { + const client = new ClaudeAgentClient({ logger: pino({ level: "silent" }) }); + const session = await client.createSession( + { provider: "claude", cwd: process.cwd(), model: "claude-haiku-4-5" }, + undefined, + { persistSession: false }, + ); + try { + for (const [model, maxTokens] of [ + ["claude-haiku-4-5", 200_000], + ["claude-opus-5-5", 1_000_000], + ["claude-haiku-4-5", 200_000], + ] as const) { + await session.setModel(model); + const events = await collectSessionTurnEvents( + session, + "Reply with exactly: OK. Do not use tools.", + { signal: AbortSignal.timeout(90_000) }, + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + usage: expect.objectContaining({ contextWindowMaxTokens: maxTokens }), + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ contextWindowMaxTokens: maxTokens }), + }), + ); + } + } finally { + await session.close(); + } +}, 300_000); 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 1f0c5bd36c..d14b27c714 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,9 +6,14 @@ 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): ClaudeQueryFactory { +function contextQuery( + 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) { @@ -41,7 +46,11 @@ function contextQuery(runtimeModel: string, reportedLimit: number): ClaudeQueryF stop_reason: null, total_cost_usd: 0, usage: { input_tokens: 249_950, output_tokens: 50, cache_read_input_tokens: 0 }, - modelUsage: { [model]: { contextWindow: reportedLimit } }, + modelUsage: Object.fromEntries( + Object.entries(modelWindowsByTurn?.[turn++] ?? { [model]: reportedLimit }).map( + ([id, contextWindow]) => [id, { contextWindow }], + ), + ), permission_denials: [], uuid: "context-result", }; @@ -110,3 +119,85 @@ test.each([ } }, ); + +test("completed turns use the active model's window after switching both ways", async () => { + const haiku = "claude-haiku-4-5"; + const opus = "claude-opus-5-5"; + 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 }, + ]), + }); + const session = await client.createSession({ + provider: "claude", + cwd: process.cwd(), + model: haiku, + }); + try { + for (const [model, limit] of [ + [haiku, 200_000], + [opus, 1_000_000], + [haiku, 200_000], + ] as const) { + await session.setModel(model); + const events = await collectSessionTurnEvents(session, "OK"); + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + usage: expect.objectContaining({ contextWindowMaxTokens: limit }), + }), + ); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ contextWindowMaxTokens: limit }), + }), + ); + } + } finally { + await session.close(); + } +}); + +test.each([ + { + name: "runtime-reported custom limit", + model: "custom-model", + windows: { "custom-model": 400_000, "prior-model": 1_000_000 }, + expected: 400_000, + }, + { + name: "missing active model entry", + model: "claude-haiku-4-5", + windows: { "prior-model": 1_000_000 }, + expected: 200_000, + }, + { + name: "invalid active model limit", + model: "claude-haiku-4-5", + windows: { "claude-haiku-4-5": -1, "prior-model": 1_000_000 }, + expected: 200_000, + }, +])("preserves $name without using another model's window", async ({ model, windows, expected }) => { + const client = new ClaudeAgentClient({ + logger: createTestLogger(), + resolveBinary: async () => process.execPath, + queryFactory: contextQuery(model, expected, [windows]), + }); + const session = await client.createSession({ provider: "claude", cwd: process.cwd(), model }); + try { + const events = await collectSessionTurnEvents(session, "OK"); + expect(events).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ contextWindowMaxTokens: expected }), + }), + ); + } finally { + await session.close(); + } +}); diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 9936969e7d..f3f7aaa5a5 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1829,30 +1829,22 @@ async function resolveClaudeAuth( } } -function extractContextWindowSize(modelUsage: unknown): number | undefined { +function extractContextWindowSize( + modelUsage: unknown, + modelId: string | null | undefined, +): number | undefined { const usageRecord = toObjectRecord(modelUsage); if (!usageRecord) { return undefined; } - let maxContextWindow: number | undefined; - for (const value of Object.values(usageRecord)) { - const valueRecord = toObjectRecord(value); - if (!valueRecord) { - continue; - } - const contextWindow = valueRecord.contextWindow; - if ( - typeof contextWindow !== "number" || - !Number.isFinite(contextWindow) || - contextWindow <= 0 - ) { - continue; - } - maxContextWindow = Math.max(maxContextWindow ?? 0, contextWindow); - } - - return maxContextWindow; + // modelUsage accumulates every model used in the session, including earlier turns + // and subagents. Only the active model describes this turn's context capacity. + const modelUsageEntry = modelId ? usageRecord[modelId] : undefined; + const contextWindow = toObjectRecord(modelUsageEntry)?.contextWindow; + return typeof contextWindow === "number" && Number.isFinite(contextWindow) && contextWindow > 0 + ? contextWindow + : undefined; } interface StreamRequestInputUsage { @@ -2006,8 +1998,8 @@ class ClaudeContextUsageState { this.contextWindowMaxTokens = contextWindowMaxTokens; } - recordModelUsage(modelUsage: unknown): number | undefined { - const contextWindowMaxTokens = extractContextWindowSize(modelUsage); + recordModelUsage(modelUsage: unknown, modelId: string | null | undefined): number | undefined { + const contextWindowMaxTokens = extractContextWindowSize(modelUsage, modelId); if (contextWindowMaxTokens !== undefined) { this.contextWindowMaxTokens = contextWindowMaxTokens; } @@ -2050,7 +2042,11 @@ class ClaudeContextUsageState { return this.createUsageUpdatedEvent(usedTokens); } - buildResultUsage(message: SDKResultMessage, modelUsage: unknown): AgentUsage | undefined { + buildResultUsage( + message: SDKResultMessage, + modelUsage: unknown, + modelId: string | null | undefined, + ): AgentUsage | undefined { try { if (!message.usage) { return undefined; @@ -2062,7 +2058,10 @@ class ClaudeContextUsageState { totalCostUsd: message.total_cost_usd, }; - const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage); + const modelContextWindowMaxTokens = this.recordModelUsage( + modelUsage ?? message.modelUsage, + modelId, + ); if (this.contextWindowMaxTokens !== undefined) { usage.contextWindowMaxTokens = this.contextWindowMaxTokens; } else if (modelContextWindowMaxTokens !== undefined) { @@ -4877,7 +4876,11 @@ class ClaudeAgentSession implements AgentSession { } private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined { - return this.contextUsage.buildResultUsage(message, modelUsage); + return this.contextUsage.buildResultUsage( + message, + modelUsage, + this.lastRuntimeModel ?? this.config.model, + ); } private handlePermissionRequest: CanUseTool = async (