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 index 3008f3d4d5..7c3883986a 100644 --- 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 @@ -1,5 +1,5 @@ import { createServer } from "node:http"; -import { mkdir, mkdtemp, rm } from "node:fs/promises"; +import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises"; import path from "node:path"; import pino from "pino"; import { expect, test } from "vitest"; @@ -13,12 +13,18 @@ test.each([ { placement: "delta", cacheTokens: 0 }, { placement: "start", cacheTokens: 1_000 }, { placement: "delta", cacheTokens: 1_000 }, + { placement: "start", cacheTokens: 0, model: "glm-5.2[1m]" }, + { placement: "start", cacheTokens: 0, model: "sonnet", resolvedModel: "glm-5.2[1m]" }, ])( - "counts $placement input and $cacheTokens cached tokens through Claude Code", - async ({ placement, cacheTokens }) => { + "counts $placement input and $cacheTokens cached tokens through Claude Code ($model)", + async ({ placement, cacheTokens, model = "claude-sonnet-4-6", resolvedModel = model }) => { const cwd = await mkdtemp(path.resolve(".claude-context-usage-")); const configDir = path.join(cwd, "config"); await mkdir(configDir); + await writeFile( + path.join(configDir, "settings.json"), + JSON.stringify({ env: { ANTHROPIC_DEFAULT_SONNET_MODEL: resolvedModel } }), + ); let requestCount = 0; const server = createServer(async (request, response) => { for await (const chunk of request) { @@ -45,7 +51,7 @@ test.each([ id: `msg_context_usage_${requestCount}`, type: "message", role: "assistant", - model: "claude-sonnet-4-6", + model: resolvedModel, content: [], stop_reason: null, stop_sequence: null, @@ -86,11 +92,9 @@ test.each([ }, }, }); - const session = await client.createSession( - { provider: "claude", cwd, model: "claude-sonnet-4-6" }, - undefined, - { persistSession: false }, - ); + const session = await client.createSession({ provider: "claude", cwd, model }, undefined, { + persistSession: false, + }); try { for (let turn = 0; turn < 2; turn++) { const events = await collectSessionTurnEvents(session, "say hello", { @@ -111,7 +115,10 @@ test.each([ expect(events).toContainEqual( expect.objectContaining({ type: "usage_updated", - usage: expect.objectContaining({ contextWindowUsedTokens: expectedTokens }), + usage: expect.objectContaining({ + contextWindowUsedTokens: expectedTokens, + contextWindowMaxTokens: resolvedModel.endsWith("[1m]") ? 1_000_000 : 200_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 new file mode 100644 index 0000000000..1f0c5bd36c --- /dev/null +++ b/packages/server/src/server/agent/providers/claude/agent.context-window.test.ts @@ -0,0 +1,112 @@ +import { expect, test } from "vitest"; +import { createTestLogger } from "../../../../test-utils/test-logger.js"; +import { collectSessionTurnEvents } from "../test-utils/session-stream-adapter.js"; +import { ClaudeAgentClient } from "./agent.js"; +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 { + return ({ prompt }) => { + let model = runtimeModel; + const messages = (async function* () { + if (typeof prompt === "string") throw new Error("Expected streaming input"); + for await (const _message of prompt) { + yield { + type: "system", + subtype: "init", + session_id: "context-session", + permissionMode: "default", + model, + }; + yield { + type: "stream_event", + session_id: "context-session", + event: { type: "message_start", message: { usage: { input_tokens: 249_950 } } }, + }; + yield { + type: "stream_event", + session_id: "context-session", + event: { type: "message_delta", usage: { output_tokens: 50 } }, + }; + yield { + type: "result", + subtype: "success", + session_id: "context-session", + duration_ms: 1, + duration_api_ms: 1, + is_error: false, + num_turns: 1, + result: "done", + stop_reason: null, + total_cost_usd: 0, + usage: { input_tokens: 249_950, output_tokens: 50, cache_read_input_tokens: 0 }, + modelUsage: { [model]: { contextWindow: reportedLimit } }, + permission_denials: [], + uuid: "context-result", + }; + } + })(); + return Object.assign(messages, { + close: () => { + void messages.return(); + }, + interrupt: async () => undefined, + setModel: async (value: string) => { + model = value; + }, + setPermissionMode: async () => undefined, + supportedModels: async () => [], + supportedCommands: async () => [], + }) as ReturnType; + }; +} + +test.each([ + { selected: "glm-5.2[1m]", runtime: "glm-5.2[1m]", expected: 1_000_000 }, + { selected: "sonnet", runtime: "glm-5.2[1m]", expected: 1_000_000 }, + { selected: "glm-5.2", runtime: "glm-5.2", expected: undefined }, +])( + "streams the context limit for $selected resolved as $runtime", + async ({ selected, runtime, expected }) => { + const client = new ClaudeAgentClient({ + logger: createTestLogger(), + queryFactory: contextQuery(runtime, 800_000), + resolveBinary: async () => process.execPath, + }); + const session = await client.createSession({ + provider: "claude", + cwd: process.cwd(), + model: selected, + }); + try { + const first = await collectSessionTurnEvents(session, "first turn"); + const live = first.filter((event) => event.type === "usage_updated"); + expect(live.map((event) => event.usage.contextWindowMaxTokens)).toEqual([expected, expected]); + expect(live.map((event) => event.usage.contextWindowUsedTokens)).toEqual([249_950, 250_000]); + expect(first).toContainEqual( + expect.objectContaining({ + type: "turn_completed", + usage: expect.objectContaining({ contextWindowMaxTokens: 800_000 }), + }), + ); + + const second = await collectSessionTurnEvents(session, "second turn"); + expect( + second + .filter((event) => event.type === "usage_updated") + .map((event) => event.usage.contextWindowMaxTokens), + ).toEqual([800_000, 800_000]); + + await session.setModel?.("another-custom-model"); + const switched = await collectSessionTurnEvents(session, "after model switch"); + expect( + switched + .filter((event) => event.type === "usage_updated") + .map((event) => event.usage.contextWindowMaxTokens), + ).toEqual([undefined, undefined]); + } finally { + await session.close(); + } + }, +); 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 5e26e6a55a..d07b0572b7 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -2805,6 +2805,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 150, }, }), @@ -2840,7 +2841,13 @@ describe("ClaudeAgentSession context window usage", () => { test("selected 1M Claude models seed active context window usage from the catalog", async () => { const session = await createSessionForTurns( - [[createInitMessage(), createMessageStartEvent(), createSuccessResult()]], + [ + [ + { ...createInitMessage(), model: "claude-sonnet-5[1m]" }, + createMessageStartEvent(), + createSuccessResult(), + ], + ], { model: "claude-sonnet-5[1m]" }, ); @@ -2880,6 +2887,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 175, }, }), @@ -2922,6 +2930,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 20_300, }, }), @@ -2963,6 +2972,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 1_069, }, }), @@ -3026,6 +3036,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 55, }, }), @@ -3035,6 +3046,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 62, }, }), @@ -3072,6 +3084,7 @@ describe("ClaudeAgentSession context window usage", () => { type: "usage_updated", provider: "claude", usage: { + contextWindowMaxTokens: 200_000, contextWindowUsedTokens: 704, }, }), diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index 51885f74e8..9936969e7d 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -31,7 +31,7 @@ import { readTaskNotificationToolUseIdFromHistoryRecord, } from "./task-notification-tool-call.js"; import { - findClaudeModel, + getClaudeContextWindowMaxTokens, getClaudeModelsWithSettings, normalizeClaudeRuntimeModelId, resolveConfiguredClaudeModel, @@ -2240,7 +2240,7 @@ class ClaudeAgentSession implements AgentSession { this.resolveBinary = options.resolveBinary; this.rewindSdk = options.rewindSdk ?? realClaudeRewindSdk; this.contextUsage = new ClaudeContextUsageState( - findClaudeModel(this.config.model)?.contextWindowMaxTokens, + getClaudeContextWindowMaxTokens(this.config.model), ); const handle = options.handle; @@ -2572,7 +2572,7 @@ class ClaudeAgentSession implements AgentSession { await this.applyFastModeFeature(false, activeQuery); } this.contextUsage.setInitialContextWindowMaxTokens( - findClaudeModel(this.config.model)?.contextWindowMaxTokens, + getClaudeContextWindowMaxTokens(this.config.model), ); this.lastOptionsModel = normalizedModelId ?? this.lastOptionsModel; this.lastRuntimeModel = null; @@ -4821,6 +4821,12 @@ class ClaudeAgentSession implements AgentSession { this.observePermissionMode(message.permissionMode); this.persistence = null; if (message.model) { + const runtimeContextWindowMaxTokens = getClaudeContextWindowMaxTokens(message.model); + // Seed usage before the first result, including settings-resolved aliases. + // A repeated init for the same model must retain the SDK's reported limit. + if (message.model !== this.lastRuntimeModel && runtimeContextWindowMaxTokens !== undefined) { + this.contextUsage.setInitialContextWindowMaxTokens(runtimeContextWindowMaxTokens); + } const normalizedRuntimeModel = normalizeClaudeRuntimeModelId(message.model); this.logger.debug( { runtimeModel: message.model, normalizedRuntimeModel }, diff --git a/packages/server/src/server/agent/providers/claude/models.test.ts b/packages/server/src/server/agent/providers/claude/models.test.ts index 1ef8177bbf..16d13b9d68 100644 --- a/packages/server/src/server/agent/providers/claude/models.test.ts +++ b/packages/server/src/server/agent/providers/claude/models.test.ts @@ -15,7 +15,27 @@ import { } from "./model-manifest.js"; import { discoverClaudeModels } from "./model-discovery.js"; import type { ModelInfo } from "@anthropic-ai/claude-agent-sdk"; -import { findClaudeModel, getClaudeModels, normalizeClaudeRuntimeModelId } from "./models.js"; +import { + findClaudeModel, + getClaudeContextWindowMaxTokens, + getClaudeModels, + normalizeClaudeRuntimeModelId, +} from "./models.js"; + +describe("getClaudeContextWindowMaxTokens", () => { + it.each([ + ["glm-5.2[1m]", 1_000_000], + [" glm-5.2[1M] ", 1_000_000], + ["claude-sonnet-4-6", 200_000], + ["claude-sonnet-4-6[1m]", 1_000_000], + ["glm-5.2", undefined], + ["glm-5.2[1m]-other", undefined], + [undefined, undefined], + [null, undefined], + ])("resolves %s to %s without guessing unknown models", (model, expected) => { + expect(getClaudeContextWindowMaxTokens(model)).toBe(expected); + }); +}); const createdClaudeConfigDirs: string[] = []; diff --git a/packages/server/src/server/agent/providers/claude/models.ts b/packages/server/src/server/agent/providers/claude/models.ts index 39e0dcf1dd..2a9b801fdb 100644 --- a/packages/server/src/server/agent/providers/claude/models.ts +++ b/packages/server/src/server/agent/providers/claude/models.ts @@ -48,6 +48,15 @@ export function findClaudeModel( return getClaudeModels().find((model) => model.id === normalizedModelId); } +export function getClaudeContextWindowMaxTokens( + modelId: string | null | undefined, +): number | undefined { + if (modelId?.trim().toLowerCase().endsWith("[1m]")) { + return 1_000_000; + } + return findClaudeModel(modelId)?.contextWindowMaxTokens; +} + export async function getClaudeModelsWithSettings( logger: Logger, configDir: string,