From 3ba2d367a282fe3053a73804a416a435cb01ee08 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Mon, 13 Jul 2026 13:35:03 +0200 Subject: [PATCH] fix: Seed context window from SDK usage report (#868) Closes #596 --- src/acp-agent.ts | 118 +++++++++++++++-------- src/tests/acp-agent.test.ts | 5 +- src/tests/create-session-options.test.ts | 39 ++++++++ src/tests/session-config-options.test.ts | 69 +++++++++++++ 4 files changed, 189 insertions(+), 42 deletions(-) diff --git a/src/acp-agent.ts b/src/acp-agent.ts index 6b48f7ac..f107af74 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -336,11 +336,12 @@ type Session = { * cancel. */ forceCancelTimer?: ReturnType; emitRawSDKMessages: boolean | SDKMessageFilter[]; - /** Context window size of the last top-level assistant model, carried across + /** Context window size of the session's current model, carried across * prompts so mid-stream usage_update notifications report a correct `size` - * before the turn's first result message arrives. Defaults to - * DEFAULT_CONTEXT_WINDOW, refreshed from each result's modelUsage, and - * invalidated when the user switches the session's model. */ + * before the turn's first result message arrives. Seeded from the SDK's + * getContextUsage report at session creation (DEFAULT_CONTEXT_WINDOW when + * that and the text heuristic both fail), refreshed the same way on model + * switches, and confirmed by each result's modelUsage. */ contextWindowSize: number; /** Accumulated task list for the session, keyed by task ID. Task IDs are * per-session, so this state must not be shared across sessions. */ @@ -1743,9 +1744,9 @@ export class ClaudeAcpAgent { // dropped dramatically) and replaced within seconds by the next // result message. // - // `size` keeps coming from session.contextWindowSize (learned - // from modelUsage / the model heuristic) — getContextUsage's - // window field under-reports extended 1M windows. + // `size` keeps coming from session.contextWindowSize — + // compaction frees occupancy, it doesn't change the model's + // window. // // The "Compacting completed." text is emitted from the `status` // handler (keyed on `compact_result`), not here, so the failure @@ -2372,11 +2373,10 @@ export class ClaudeAcpAgent { const model = message.event.message.model; if (model && model !== "") { lastAssistantModel = model; - // Only upgrade from the default — once a `result` has given - // us an authoritative window, trust it over the heuristic. - // Model switches invalidate the cached window via - // `syncSessionConfigState`, which resets us back to the - // default so this branch runs again for the new model. + // Only upgrade from the default — once the SDK has given us + // an authoritative window (seeded at session creation, + // refreshed on model switches in `applyConfigOptionValue`, + // confirmed by each `result`), trust it over the heuristic. if (session.contextWindowSize === DEFAULT_CONTEXT_WINDOW) { const inferred = inferContextWindowFromModel(model); if (inferred !== null) { @@ -3693,16 +3693,22 @@ export class ClaudeAcpAgent { // carries no "1m" token. const newModelInfo = session.modelInfos.find((m) => m.value === value); if (session.models.currentModelId !== value) { - // The cached context window was learned for the previous model; reset - // to the new model's heuristic so mid-stream updates between now and - // the next `result` reflect the user's selection instead of the old - // model's window. + // The cached context window was learned for the previous model. The + // SDK is already running the new model here (user-driven switches call + // `query.setModel` before this, and the refusal-fallback sync only + // reconciles a switch the SDK already made), so ask it for the new + // window; fall back to the text heuristic so mid-stream updates + // between now and the next `result` reflect the user's selection + // instead of the old model's window. session.contextWindowSize = + (await fetchContextWindowSize(session.query, this.logger)) ?? inferContextWindowFromModel( value, + newModelInfo?.resolvedModel, newModelInfo?.displayName, newModelInfo?.description, - ) ?? DEFAULT_CONTEXT_WINDOW; + ) ?? + DEFAULT_CONTEXT_WINDOW; } session.models = { ...session.models, currentModelId: value }; @@ -4428,6 +4434,29 @@ export class ClaudeAcpAgent { effortLevel: initialEffort.currentValue as Settings["effortLevel"], }); } + // Seed the context window from the SDK's authoritative report. Text + // inference alone misses aliases that resolve to extended-context models + // with no "1m" token anywhere in their id or description (e.g. `sonnet` → + // claude-sonnet-5, natively ~1M): those streamed `usage_update.size: + // 200000` until the first result's modelUsage corrected it — again on + // every process restart or session re-creation, since the learned window + // lives only on the Session (issue #596). + // + // The inference fallback is deliberately keyed to the allowlisted entry: a + // fallback-resolved sibling's resolvedModel/displayName/description can + // describe a different context lane than the verbatim live id (e.g. an + // "opus[1m]" row matched for a bare 200k id), so on the fallback path only + // the id itself is a trustworthy window signal. + const contextWindowSize = + (await fetchContextWindowSize(q, this.logger)) ?? + inferContextWindowFromModel( + models.currentModelId, + allowlistedModelInfo?.resolvedModel, + allowlistedModelInfo?.displayName, + allowlistedModelInfo?.description, + ) ?? + DEFAULT_CONTEXT_WINDOW; + this.sessions[sessionId] = { query: q, input: input, @@ -4450,17 +4479,7 @@ export class ClaudeAcpAgent { fastModeEnabled, abortController, emitRawSDKMessages: sessionMeta?.claudeCode?.emitRawSDKMessages ?? false, - contextWindowSize: - // Deliberately keyed to the allowlisted entry: a fallback-resolved - // sibling's displayName/description can describe a different context - // lane than the verbatim live id (e.g. an "opus[1m]" row matched for - // a bare 200k id), so on the fallback path only the id itself is a - // trustworthy window signal. - inferContextWindowFromModel( - models.currentModelId, - allowlistedModelInfo?.displayName, - allowlistedModelInfo?.description, - ) ?? DEFAULT_CONTEXT_WINDOW, + contextWindowSize, taskState, toolUseCache: {}, emittedToolCalls: new Set(), @@ -6043,12 +6062,14 @@ function commonPrefixLength(a: string, b: string) { * Anthropic 1M-context variants encode "1m" as a distinct token in the SDK * model ID (e.g., "claude-opus-4-6-1m"), which `\b1m\b` catches without also * matching things like "10m" or embedded substrings. Semantic aliases like - * `default` carry no such token in the ID, but the SDK's human-facing - * `displayName`/`description` do (e.g. "Opus 4.7 (1M context)"), so callers - * pass those too — the same `\b1m\b` token appears in "1M context". The SDK's - * `ModelInfo` exposes no structured context-window field, so this text scan is - * the only pre-`result` signal available. A miss falls back to the default - * window and is corrected by `result.modelUsage` within one turn. */ + * `default` carry no such token in the ID, but their `resolvedModel` and the + * SDK's human-facing `displayName`/`description` can (e.g. + * "claude-opus-4-8[1m]", "Opus 4.7 (1M context)"), so callers pass those too. + * This text scan can't catch every model — some resolve to extended-context + * models with no "1m" anywhere (e.g. `sonnet` → claude-sonnet-5, natively + * ~1M) — which is why `fetchContextWindowSize` is preferred wherever a live + * query is available. A miss falls back to the default window and is + * corrected by `result.modelUsage` within one turn. */ function inferContextWindowFromModel(...texts: Array): number | null { if (texts.some((text) => text != null && /\b1m\b/i.test(text))) return 1_000_000; return null; @@ -6058,12 +6079,7 @@ function inferContextWindowFromModel(...texts: Array): numbe * `getContextUsage` control request. Unlike the per-message API usage numbers * (which only count message tokens), this `totalTokens` includes the system * prompt, tool schemas, MCP tools, and memory-file overhead — the real - * occupancy the user sees. Returns `null` on any control-request failure. - * - * Note: we deliberately do NOT use this response's window fields for `size`. - * They have been observed to under-report extended (1M) context windows, so - * the window keeps coming from `modelUsage` / `inferContextWindowFromModel`, - * which handle the 1M variants correctly. */ + * occupancy the user sees. Returns `null` on any control-request failure. */ async function fetchContextUsedTokens(query: Query, logger: Logger): Promise { try { const usage = await query.getContextUsage(); @@ -6074,6 +6090,28 @@ async function fetchContextUsedTokens(query: Query, logger: Logger): Promise { + try { + const usage = await query.getContextUsage(); + return usage.rawMaxTokens > 0 ? usage.rawMaxTokens : null; + } catch (error) { + logger.error("Failed to fetch context window size from SDK:", error); + return null; + } +} + /** Translate the legacy `MAX_THINKING_TOKENS` env var into the SDK's `thinking` * option. The `maxThinkingTokens` option it used to feed is deprecated and * reduced to on/off on current models, so map the value to explicit thinking diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index c39c9e16..f6de2292 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -4902,8 +4902,9 @@ describe("usage_update computation", () => { { type: "system", subtype: "compact_boundary", session_id: "test-session" }, ]); const session = agent.sessions["test-session"]; - // A 1M window learned earlier (e.g. from modelUsage) must survive compaction - // — getContextUsage's window field under-reports it, so we don't use it. + // A 1M window learned earlier (e.g. from modelUsage) must survive + // compaction — compaction frees occupancy, it doesn't change the window, + // so the handler must not overwrite it from this response. session.contextWindowSize = 1000000; (session.query as any).getContextUsage = vi .fn() diff --git a/src/tests/create-session-options.test.ts b/src/tests/create-session-options.test.ts index 18080545..ec14960c 100644 --- a/src/tests/create-session-options.test.ts +++ b/src/tests/create-session-options.test.ts @@ -7,6 +7,7 @@ import * as os from "node:os"; import * as path from "node:path"; let capturedOptions: Options | undefined; +let contextUsageResult: (() => Promise<{ rawMaxTokens: number }>) | undefined; vi.mock("@anthropic-ai/claude-agent-sdk", async () => { const actual = await vi.importActual( "@anthropic-ai/claude-agent-sdk", @@ -29,6 +30,10 @@ vi.mock("@anthropic-ai/claude-agent-sdk", async () => { setModel: async () => {}, setPermissionMode: async () => {}, supportedCommands: async () => [], + getContextUsage: () => + contextUsageResult + ? contextUsageResult() + : Promise.reject(new Error("no context usage mocked")), [Symbol.asyncIterator]: async function* () {}, }; }, @@ -58,6 +63,7 @@ describe("createSession options merging", () => { beforeEach(async () => { capturedOptions = undefined; + contextUsageResult = undefined; vi.resetModules(); const acpAgent = await import("../acp-agent.js"); @@ -614,6 +620,39 @@ describe("createSession options merging", () => { }); }); + describe("context window seeding", () => { + function sessionFor(sessionId: string) { + return (agent as unknown as { sessions: Record }).sessions[sessionId]; + } + + it("seeds contextWindowSize from the SDK's getContextUsage report", async () => { + // Aliases like `sonnet` can resolve to an extended-context model with no + // "1m" token anywhere in id/displayName/description, so the authoritative + // report must win over text inference (issue #596). + contextUsageResult = async () => ({ rawMaxTokens: 967000 }); + + const response = await agent.newSession({ cwd: process.cwd(), mcpServers: [] }); + + expect(sessionFor(response.sessionId).contextWindowSize).toBe(967000); + }); + + it("falls back to the default window when getContextUsage fails and inference misses", async () => { + // The default mock rejects getContextUsage, and the mock model + // ("claude-sonnet-4-6" / "Claude Sonnet" / "Fast") has no "1m" hint. + const response = await agent.newSession({ cwd: process.cwd(), mcpServers: [] }); + + expect(sessionFor(response.sessionId).contextWindowSize).toBe(200000); + }); + + it("ignores a nonsensical (non-positive) reported window", async () => { + contextUsageResult = async () => ({ rawMaxTokens: 0 }); + + const response = await agent.newSession({ cwd: process.cwd(), mcpServers: [] }); + + expect(sessionFor(response.sessionId).contextWindowSize).toBe(200000); + }); + }); + describe("refusal fallback dialog", () => { it("declares the dialog kind and callback when the client supports form elicitation", async () => { await agent.initialize({ diff --git a/src/tests/session-config-options.test.ts b/src/tests/session-config-options.test.ts index 962ea56c..aff7158b 100644 --- a/src/tests/session-config-options.test.ts +++ b/src/tests/session-config-options.test.ts @@ -927,6 +927,75 @@ describe("session config options", () => { }); }); + describe("context window on model change", () => { + beforeEach(() => { + populateSession(); + }); + + function getSession() { + return (agent as unknown as { sessions: Record }).sessions[SESSION_ID]; + } + + it("refreshes contextWindowSize from getContextUsage when the model changes", async () => { + const session = getSession(); + session.query.getContextUsage = vi.fn(async () => ({ rawMaxTokens: 967000 })); + + await agent.setSessionConfigOption({ + sessionId: SESSION_ID, + configId: "model", + value: "claude-sonnet-4-6", + }); + + expect(session.query.getContextUsage).toHaveBeenCalled(); + expect(session.contextWindowSize).toBe(967000); + }); + + it("falls back to resolvedModel text inference when getContextUsage fails", async () => { + const session = getSession(); + session.query.getContextUsage = vi.fn().mockRejectedValue(new Error("boom")); + session.modelInfos = session.modelInfos.map((m: ModelInfo) => + m.value === "claude-sonnet-4-6" ? { ...m, resolvedModel: "claude-sonnet-5[1m]" } : m, + ); + + await agent.setSessionConfigOption({ + sessionId: SESSION_ID, + configId: "model", + value: "claude-sonnet-4-6", + }); + + expect(session.contextWindowSize).toBe(1_000_000); + }); + + it("falls back to the default window when the SDK and inference both miss", async () => { + const session = getSession(); + session.contextWindowSize = 1_000_000; + session.query.getContextUsage = vi.fn().mockRejectedValue(new Error("boom")); + + await agent.setSessionConfigOption({ + sessionId: SESSION_ID, + configId: "model", + value: "claude-sonnet-4-6", + }); + + expect(session.contextWindowSize).toBe(200000); + }); + + it("keeps the learned window when re-asserting the current model", async () => { + const session = getSession(); + session.contextWindowSize = 1_000_000; + session.query.getContextUsage = vi.fn(async () => ({ rawMaxTokens: 200000 })); + + await agent.setSessionConfigOption({ + sessionId: SESSION_ID, + configId: "model", + value: "claude-opus-4-5", + }); + + expect(session.query.getContextUsage).not.toHaveBeenCalled(); + expect(session.contextWindowSize).toBe(1_000_000); + }); + }); + describe("auto mode availability per model", () => { /** * Augment the session populated by `populateSession()` with a Haiku entry