fix: use current model's context window for usage_update size

timvisher-dd · claude · timvisher-dd · commit 0962531d354d · 2026-03-12T14:06:53.000-04:00
The original code took Math.min across all modelUsage entries, which
reported a 200k context window when subagents (Sonnet/Haiku) were
mixed with Opus 1M. Track the top-level assistant model and look up
its context window specifically.

The SDK's streaming path keys modelUsage by the requested model alias
(e.g. "claude-opus-4-6") while BetaMessage.model on assistant messages
has the resolved API response model (e.g. "claude-opus-4-6-20250514").
Fall back to prefix matching (preferring the longest/most-specific
match) when exact key lookup misses.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/src/acp-agent.ts b/src/acp-agent.ts
@@ -478,6 +478,7 @@ export class ClaudeAcpAgent implements Agent {
     };
 
     let lastAssistantTotalUsage: number | null = null;
+    let lastAssistantModel: string | null = null;
 
     const userMessage = promptToClaude(params);
 
@@ -576,10 +577,25 @@ export class ClaudeAcpAgent implements Agent {
             session.accumulatedUsage.cachedReadTokens += message.usage.cache_read_input_tokens;
             session.accumulatedUsage.cachedWriteTokens += message.usage.cache_creation_input_tokens;
 
-            // Calculate context window size from modelUsage (minimum across all models used)
-            const contextWindows = Object.values(message.modelUsage).map((m) => m.contextWindow);
-            const contextWindowSize =
-              contextWindows.length > 0 ? Math.min(...contextWindows) : 200000;
+            // Calculate context window size from the current model's usage.
+            // The modelUsage keys may use the requested model alias (e.g. "claude-opus-4-6")
+            // while message.model on assistant messages has the resolved API response model
+            // (e.g. "claude-opus-4-6-20250514"), so we fall back to prefix matching.
+            const currentModel = lastAssistantModel;
+            const matchingModelUsage = currentModel
+              ? message.modelUsage[currentModel] ??
+                Object.entries(message.modelUsage)
+                  .filter(
+                    ([key]) =>
+                      currentModel.startsWith(key) || key.startsWith(currentModel),
+                  )
+                  .sort((a, b) => b[0].length - a[0].length)[0]?.[1]
+              : undefined;
+            // Fallback to 200k: this is hit when lastAssistantModel is null (e.g. the
+            // assistant message lacked a model field) or no modelUsage key matches.
+            // 200k is a conservative default — the Anthropic API should always populate
+            // BetaMessage.model, so this path is unlikely in practice.
+            const contextWindowSize = matchingModelUsage?.contextWindow ?? 200000;
 
             // Send usage_update notification
             if (lastAssistantTotalUsage !== null) {
@@ -690,6 +706,11 @@ export class ClaudeAcpAgent implements Agent {
             }
 
             // Store latest assistant usage (excluding subagents)
+            // Sum all token types as a proxy for post-turn context occupancy:
+            // current turn's output will become next turn's input.
+            // Note: per the Anthropic API, input_tokens excludes cache tokens —
+            // cache_read and cache_creation are reported separately, so summing
+            // all four fields is not double-counting.
             if ((message.message as any).usage && message.parent_tool_use_id === null) {
               const messageWithUsage = message.message as unknown as SDKResultMessage;
               lastAssistantTotalUsage =
@@ -698,6 +719,15 @@ export class ClaudeAcpAgent implements Agent {
                 messageWithUsage.usage.cache_read_input_tokens +
                 messageWithUsage.usage.cache_creation_input_tokens;
             }
+            // Track the current top-level model for context window size lookup
+            // (exclude subagent messages to stay in sync with lastAssistantTotalUsage)
+            if (
+              message.type === "assistant" &&
+              message.parent_tool_use_id === null &&
+              message.message.model
+            ) {
+              lastAssistantModel = message.message.model;
+            }
 
             // Slash commands like /compact can generate invalid output... doesn't match
             // their own docs: https://docs.anthropic.com/en/docs/claude-code/sdk/sdk-slash-commands#%2Fcompact-compact-conversation-history
diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts
@@ -1416,3 +1416,287 @@ describe("stop reason propagation", () => {
     ).rejects.toThrow("Internal error");
   });
 });
+
+describe("usage_update computation", () => {
+  function createAssistantMessage(overrides: {
+    model: string;
+    usage?: { input_tokens: number; output_tokens: number; cache_read_input_tokens: number; cache_creation_input_tokens: number };
+  }) {
+    return {
+      type: "assistant" as const,
+      parent_tool_use_id: null,
+      uuid: randomUUID(),
+      session_id: "test-session",
+      message: {
+        model: overrides.model,
+        content: [{ type: "text", text: "hello" }],
+        usage: overrides.usage ?? {
+          input_tokens: 100,
+          output_tokens: 50,
+          cache_read_input_tokens: 20,
+          cache_creation_input_tokens: 10,
+        },
+      },
+    };
+  }
+
+  function createResultMessageWithModel(overrides: {
+    modelUsage: Record<string, { inputTokens: number; outputTokens: number; cacheReadInputTokens: number; cacheCreationInputTokens: number; webSearchRequests: number; costUSD: number; contextWindow: number; maxOutputTokens: number }>;
+  }) {
+    return {
+      type: "result" as const,
+      subtype: "success" as const,
+      stop_reason: "end_turn",
+      is_error: false,
+      result: "",
+      errors: [],
+      duration_ms: 0,
+      duration_api_ms: 0,
+      num_turns: 1,
+      total_cost_usd: 0.01,
+      usage: {
+        input_tokens: 10,
+        output_tokens: 5,
+        cache_read_input_tokens: 0,
+        cache_creation_input_tokens: 0,
+      },
+      modelUsage: overrides.modelUsage,
+      permission_denials: [],
+      uuid: randomUUID(),
+      session_id: "test-session",
+    };
+  }
+
+  function createMockAgentWithCapture() {
+    const updates: any[] = [];
+    const mockClient = {
+      sessionUpdate: async (notification: any) => {
+        updates.push(notification);
+      },
+    } as unknown as AgentSideConnection;
+    const agent = new ClaudeAcpAgent(mockClient, { log: () => {}, error: () => {} });
+    return { agent, updates };
+  }
+
+  function injectSession(agent: ClaudeAcpAgent, messages: any[]) {
+    const gen = (function* () { yield* messages; })();
+    agent.sessions["test-session"] = {
+      query: gen as any,
+      input: new Pushable(),
+      cancelled: false,
+      cwd: "/test",
+      permissionMode: "default",
+      settingsManager: {} as any,
+      accumulatedUsage: {
+        inputTokens: 0,
+        outputTokens: 0,
+        cachedReadTokens: 0,
+        cachedWriteTokens: 0,
+      },
+      configOptions: [],
+      promptRunning: false,
+      pendingMessages: new Map(),
+      nextPendingOrder: 0,
+    };
+  }
+
+  it("used sums all token types as post-turn context occupancy proxy", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createAssistantMessage({
+        model: "claude-opus-4-20250514",
+        usage: { input_tokens: 1000, output_tokens: 500, cache_read_input_tokens: 200, cache_creation_input_tokens: 100 },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 1000, outputTokens: 500, cacheReadInputTokens: 200,
+            cacheCreationInputTokens: 100, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // used = input(1000) + output(500) + cache_read(200) + cache_creation(100) = 1800
+    expect(usageUpdate.update.used).toBe(1800);
+  });
+
+  it("size reflects the current model's context window, not min across all", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-opus-4-20250514" }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 100, outputTokens: 50, cacheReadInputTokens: 20,
+            cacheCreationInputTokens: 10, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+          "claude-sonnet-4-20250514": {
+            inputTokens: 50, outputTokens: 25, cacheReadInputTokens: 10,
+            cacheCreationInputTokens: 5, webSearchRequests: 0, costUSD: 0.005,
+            contextWindow: 200000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // size should be 1000000 (Opus), not 200000 (min of both)
+    expect(usageUpdate.update.size).toBe(1000000);
+  });
+
+  it("after model switch, size updates to the new model's window", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    // Simulate: assistant on Sonnet with both models in modelUsage
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-sonnet-4-20250514" }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 100, outputTokens: 50, cacheReadInputTokens: 20,
+            cacheCreationInputTokens: 10, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+          "claude-sonnet-4-20250514": {
+            inputTokens: 50, outputTokens: 25, cacheReadInputTokens: 10,
+            cacheCreationInputTokens: 5, webSearchRequests: 0, costUSD: 0.005,
+            contextWindow: 200000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // size should be 200000 (Sonnet - the current model)
+    expect(usageUpdate.update.size).toBe(200000);
+  });
+
+  it("after switching back to original model, size returns to original window", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    // Last assistant message is Opus again
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-sonnet-4-20250514" }),
+      createAssistantMessage({ model: "claude-opus-4-20250514" }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 200, outputTokens: 100, cacheReadInputTokens: 40,
+            cacheCreationInputTokens: 20, webSearchRequests: 0, costUSD: 0.02,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+          "claude-sonnet-4-20250514": {
+            inputTokens: 50, outputTokens: 25, cacheReadInputTokens: 10,
+            cacheCreationInputTokens: 5, webSearchRequests: 0, costUSD: 0.005,
+            contextWindow: 200000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // size should be 1000000 (Opus - switched back)
+    expect(usageUpdate.update.size).toBe(1000000);
+  });
+
+  it("subagent assistant messages do not affect size (top-level model is used)", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    // Top-level assistant on Opus, then subagent on Haiku (parent_tool_use_id set)
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-opus-4-20250514" }),
+      {
+        type: "assistant" as const,
+        parent_tool_use_id: "tool_use_123",
+        uuid: randomUUID(),
+        session_id: "test-session",
+        message: {
+          model: "claude-haiku-4-5-20251001",
+          content: [{ type: "text", text: "subagent response" }],
+          usage: { input_tokens: 50, output_tokens: 25, cache_read_input_tokens: 0, cache_creation_input_tokens: 0 },
+        },
+      },
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 100, outputTokens: 50, cacheReadInputTokens: 20,
+            cacheCreationInputTokens: 10, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+          "claude-haiku-4-5-20251001": {
+            inputTokens: 50, outputTokens: 25, cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0, webSearchRequests: 0, costUSD: 0.001,
+            contextWindow: 200000, maxOutputTokens: 8192,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // size should be 1000000 (Opus - the top-level model), NOT 200000 (Haiku subagent)
+    expect(usageUpdate.update.size).toBe(1000000);
+  });
+
+  it("prefix-matches when assistant model has date suffix but modelUsage key does not", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    // The API response has the full versioned model ID on assistant messages,
+    // but the SDK's streaming path may key modelUsage by the shorter alias.
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-opus-4-6-20250514" }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-6": {
+            inputTokens: 100, outputTokens: 50, cacheReadInputTokens: 20,
+            cacheCreationInputTokens: 10, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    // Should match via prefix: "claude-opus-4-6-20250514".startsWith("claude-opus-4-6")
+    expect(usageUpdate.update.size).toBe(1000000);
+  });
+
+  it("prefix-matches when modelUsage key has date suffix but assistant model does not", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createAssistantMessage({ model: "claude-opus-4-6" }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-6-20250514": {
+            inputTokens: 100, outputTokens: 50, cacheReadInputTokens: 20,
+            cacheCreationInputTokens: 10, webSearchRequests: 0, costUSD: 0.01,
+            contextWindow: 1000000, maxOutputTokens: 16384,
+          },
+        },
+      }),
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdate = updates.find((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdate).toBeDefined();
+    expect(usageUpdate.update.size).toBe(1000000);
+  });
+});
