fix: emit usage updates from Claude stream events (#506)

## Summary
- emit ACP `usage_update` notifications from top-level Claude
`stream_event` message usage
- keep tracking a usage snapshot across `message_start` and
`message_delta` so partial deltas still produce a complete total
- preserve the existing `result`-path update as the final
exact/cost-bearing usage event

## Why
Claude SDK stream events already carry incremental usage on
`message_delta`, but `claude-agent-acp` currently ignores that path and
only emits `usage_update` from the final `result`. This makes usage
appear stale until the turn ends in ACP clients that expect live context
usage updates.

## Verification
- built the repo with `npm run build`
- ran a direct ACP stdio client against the patched local build
- observed `usage_update` notifications emitted before the final
result-path usage update

Patched local smoke-test output included this sequence:
- `usage_update used=20973 size=200000`
- `agent_message_chunk text=TEST_OK`
- `usage_update used=20989 size=200000`
- `usage_update used=20989 size=1000000 cost=0.1314875`

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

src/acp-agent.ts

@@ -54,7 +54,6 @@ import {
   Query,
   query,
   SDKPartialAssistantMessage,
-  SDKResultMessage,
   SDKUserMessage,
   SlashCommand,
 } from "@anthropic-ai/claude-agent-sdk";
@@ -108,6 +107,22 @@ type AccumulatedUsage = {
   cachedWriteTokens: number;
 };
 
+type UsageSnapshot = {
+  input_tokens: number;
+  output_tokens: number;
+  cache_read_input_tokens: number;
+  cache_creation_input_tokens: number;
+};
+
+const ZERO_USAGE = Object.freeze({
+  input_tokens: 0,
+  output_tokens: 0,
+  cache_read_input_tokens: 0,
+  cache_creation_input_tokens: 0,
+});
+
+const DEFAULT_CONTEXT_WINDOW = 200000;
+
 type Session = {
   query: Query;
   input: Pushable<SDKUserMessage>;
@@ -126,6 +141,12 @@ type Session = {
   nextPendingOrder: number;
   abortController: AbortController;
   emitRawSDKMessages: boolean | SDKMessageFilter[];
+  /** Context window size of the last top-level assistant 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. */
+  contextWindowSize: number;
 };
 
 /** Compute a stable fingerprint of the session-defining params so we can
@@ -545,8 +566,8 @@ export class ClaudeAcpAgent implements Agent {
     };
 
     let lastAssistantTotalUsage: number | null = null;
+    let lastAssistantUsage: UsageSnapshot | null = null;
     let lastAssistantModel: string | null = null;
-    let lastContextWindowSize: number = 200000;
 
     const userMessage = promptToClaude(params);
 
@@ -628,12 +649,13 @@ export class ClaudeAcpAgent implements Agent {
                 // "944k/1m" right after the user sees "Compacting completed",
                 // which is confusing and wrong.
                 lastAssistantTotalUsage = 0;
+                lastAssistantUsage = null;
                 await this.client.sessionUpdate({
                   sessionId: message.session_id,
                   update: {
                     sessionUpdate: "usage_update",
                     used: 0,
-                    size: lastContextWindowSize,
+                    size: session.contextWindowSize,
                   },
                 });
                 await this.client.sessionUpdate({
@@ -691,8 +713,13 @@ export class ClaudeAcpAgent implements Agent {
             const matchingModelUsage = lastAssistantModel
               ? getMatchingModelUsage(message.modelUsage, lastAssistantModel)
               : null;
-            const contextWindowSize = matchingModelUsage?.contextWindow ?? 200000;
-            lastContextWindowSize = contextWindowSize;
+            // Only overwrite when we have an authoritative value — a miss
+            // (e.g. a turn with no top-level assistant message) would
+            // otherwise discard the window learned on a prior turn and
+            // leave the next prompt's mid-stream updates reporting 200k.
+            if (matchingModelUsage) {
+              session.contextWindowSize = matchingModelUsage.contextWindow;
+            }
 
             // Send usage_update notification
             if (lastAssistantTotalUsage !== null) {
@@ -701,7 +728,7 @@ export class ClaudeAcpAgent implements Agent {
                 update: {
                   sessionUpdate: "usage_update",
                   used: lastAssistantTotalUsage,
-                  size: contextWindowSize,
+                  size: session.contextWindowSize,
                   cost: {
                     amount: message.total_cost_usd,
                     currency: "USD",
@@ -775,6 +802,57 @@ export class ClaudeAcpAgent implements Agent {
             break;
           }
           case "stream_event": {
+            if (
+              message.parent_tool_use_id === null &&
+              (message.event.type === "message_start" || message.event.type === "message_delta")
+            ) {
+              if (message.event.type === "message_start") {
+                lastAssistantUsage = snapshotFromUsage(message.event.message.usage);
+                const model = message.event.message.model;
+                if (model && model !== "<synthetic>") {
+                  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.
+                  if (session.contextWindowSize === DEFAULT_CONTEXT_WINDOW) {
+                    const inferred = inferContextWindowFromModel(model);
+                    if (inferred !== null) {
+                      session.contextWindowSize = inferred;
+                    }
+                  }
+                }
+              } else {
+                const usage = message.event.usage;
+                const prev: Readonly<UsageSnapshot> = lastAssistantUsage ?? ZERO_USAGE;
+                // Per Anthropic API, message_delta usage fields are *cumulative*;
+                // nullable fields (input_tokens and the cache fields) fall back
+                // to the prior snapshot when the server omits them from this
+                // delta. Only output_tokens is guaranteed non-null.
+                lastAssistantUsage = {
+                  input_tokens: usage.input_tokens ?? prev.input_tokens,
+                  output_tokens: usage.output_tokens,
+                  cache_read_input_tokens:
+                    usage.cache_read_input_tokens ?? prev.cache_read_input_tokens,
+                  cache_creation_input_tokens:
+                    usage.cache_creation_input_tokens ?? prev.cache_creation_input_tokens,
+                };
+              }
+
+              const nextUsage = totalTokens(lastAssistantUsage);
+              if (nextUsage !== lastAssistantTotalUsage) {
+                lastAssistantTotalUsage = nextUsage;
+                await this.client.sessionUpdate({
+                  sessionId: params.sessionId,
+                  update: {
+                    sessionUpdate: "usage_update",
+                    used: nextUsage,
+                    size: session.contextWindowSize,
+                  },
+                });
+              }
+            }
             for (const notification of streamEventToAcpNotifications(
               message,
               params.sessionId,
@@ -817,29 +895,16 @@ 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 =
-                messageWithUsage.usage.input_tokens +
-                messageWithUsage.usage.output_tokens +
-                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 &&
-              message.message.model !== "<synthetic>"
-            ) {
-              lastAssistantModel = message.message.model;
+            // Snapshot the latest top-level assistant usage and model so the
+            // next `result` can emit a usage_update tied to the right context
+            // window. Subagent messages are excluded to keep the snapshot
+            // aligned with what the user's current selection is producing.
+            if (message.type === "assistant" && message.parent_tool_use_id === null) {
+              lastAssistantUsage = snapshotFromUsage(message.message.usage);
+              lastAssistantTotalUsage = totalTokens(lastAssistantUsage);
+              if (message.message.model && message.message.model !== "<synthetic>") {
+                lastAssistantModel = message.message.model;
+              }
             }
 
             // Slash commands like /compact can generate invalid output... doesn't match
@@ -1334,6 +1399,13 @@ export class ClaudeAcpAgent implements Agent {
     if (configId === "mode") {
       session.modes = { ...session.modes, currentModeId: value };
     } else if (configId === "model") {
+      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.
+        session.contextWindowSize = inferContextWindowFromModel(value) ?? DEFAULT_CONTEXT_WINDOW;
+      }
       session.models = { ...session.models, currentModelId: value };
     }
   }
@@ -1642,6 +1714,8 @@ export class ClaudeAcpAgent implements Agent {
       nextPendingOrder: 0,
       abortController,
       emitRawSDKMessages: sessionMeta?.claudeCode?.emitRawSDKMessages ?? false,
+      contextWindowSize:
+        inferContextWindowFromModel(models.currentModelId) ?? DEFAULT_CONTEXT_WINDOW,
     };
 
     return {
@@ -1678,6 +1752,38 @@ function sessionUsage(session: Session) {
   };
 }
 
+/** Sum all four fields as a proxy for post-turn context occupancy: the current
+ *  turn's output becomes next turn's input. Per the Anthropic API, input_tokens
+ *  excludes cache tokens — cache_read and cache_creation are reported
+ *  separately — so summing all four is not double-counting. */
+function totalTokens(usage: UsageSnapshot): number {
+  return (
+    usage.input_tokens +
+    usage.output_tokens +
+    usage.cache_read_input_tokens +
+    usage.cache_creation_input_tokens
+  );
+}
+
+/** Project a nullable API usage object into our non-null snapshot shape.
+ *  Both SDK message_start and assistant message `usage` have `number | null`
+ *  cache fields; we coerce absent values to 0 so `totalTokens` never hits
+ *  NaN. Delta events have different semantics (cumulative + prev fallback)
+ *  and are handled inline. */
+function snapshotFromUsage(usage: {
+  input_tokens: number;
+  output_tokens: number;
+  cache_read_input_tokens: number | null;
+  cache_creation_input_tokens: number | null;
+}): UsageSnapshot {
+  return {
+    input_tokens: usage.input_tokens,
+    output_tokens: usage.output_tokens,
+    cache_read_input_tokens: usage.cache_read_input_tokens ?? 0,
+    cache_creation_input_tokens: usage.cache_creation_input_tokens ?? 0,
+  };
+}
+
 function createEnvForGateway(gatewayMeta?: GatewayAuthMeta) {
   if (!gatewayMeta) {
     return {};
@@ -2302,6 +2408,16 @@ function commonPrefixLength(a: string, b: string) {
   return i;
 }
 
+/** Best-effort first guess of a model's context window from its ID, used only
+ *  until a `result` message arrives with the authoritative `modelUsage` value.
+ *  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. */
+function inferContextWindowFromModel(model: string): number | null {
+  if (/\b1m\b/i.test(model)) return 1_000_000;
+  return null;
+}
+
 function getMatchingModelUsage(modelUsage: Record<string, ModelUsage>, currentModel: string) {
   let bestKey: string | null = null;
   let bestLen = 0;

src/tests/acp-agent.test.ts

@@ -1350,6 +1350,7 @@ describe("stop reason propagation", () => {
       nextPendingOrder: 0,
       abortController: new AbortController(),
       emitRawSDKMessages: false,
+      contextWindowSize: 200000,
     };
   }
 
@@ -1491,6 +1492,7 @@ describe("stop reason propagation", () => {
       pendingMessages: new Map(),
       nextPendingOrder: 0,
       emitRawSDKMessages: false,
+      contextWindowSize: 200000,
     };
 
     const response = await agent.prompt({
@@ -1566,6 +1568,7 @@ describe("session/close", () => {
       nextPendingOrder: 0,
       abortController: new AbortController(),
       emitRawSDKMessages: false,
+      contextWindowSize: 200000,
     };
     return agent.sessions[sessionId]!;
   }
@@ -1660,6 +1663,7 @@ describe("getOrCreateSession param change detection", () => {
       nextPendingOrder: 0,
       abortController: new AbortController(),
       emitRawSDKMessages: false,
+      contextWindowSize: 200000,
     };
     return agent.sessions[sessionId]!;
   }
@@ -1819,6 +1823,23 @@ describe("usage_update computation", () => {
     };
   }
 
+  function createStreamEvent(
+    eventType: "message_start" | "message_delta",
+    payload: Record<string, unknown>,
+    parentToolUseId: string | null = null,
+  ) {
+    return {
+      type: "stream_event" as const,
+      parent_tool_use_id: parentToolUseId,
+      uuid: randomUUID(),
+      session_id: "test-session",
+      event:
+        eventType === "message_start"
+          ? { type: "message_start" as const, message: payload }
+          : { type: "message_delta" as const, ...payload },
+    };
+  }
+
   function createMockAgentWithCapture() {
     const updates: any[] = [];
     const mockClient = {
@@ -1875,6 +1896,7 @@ describe("usage_update computation", () => {
       nextPendingOrder: 0,
       abortController: new AbortController(),
       emitRawSDKMessages: false,
+      contextWindowSize: 200000,
     };
   }
 
@@ -1915,6 +1937,461 @@ describe("usage_update computation", () => {
     expect(usageUpdate.update.used).toBe(1800);
   });
 
+  it("stream_event message_start emits usage_update before result", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        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,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.used).toBe(1800);
+    // First prompt of a session has no prior result to learn the window from,
+    // so the mid-stream update falls back to the default context window.
+    expect(usageUpdates[0].update.size).toBe(200000);
+    expect(usageUpdates[0].update.cost).toBeUndefined();
+    expect(usageUpdates[1].update.used).toBe(1800);
+    expect(usageUpdates[1].update.size).toBe(1000000);
+    expect(usageUpdates[1].update.cost).toBeDefined();
+  });
+
+  it("stream_event message_delta patches previous snapshot", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-20250514",
+        usage: {
+          input_tokens: 1000,
+          output_tokens: 0,
+          cache_read_input_tokens: 200,
+          cache_creation_input_tokens: 100,
+        },
+      }),
+      createStreamEvent("message_delta", {
+        usage: { output_tokens: 500 },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 1000,
+            outputTokens: 500,
+            cacheReadInputTokens: 200,
+            cacheCreationInputTokens: 100,
+            webSearchRequests: 0,
+            costUSD: 0.01,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(3);
+    expect(usageUpdates[0].update.used).toBe(1300);
+    expect(usageUpdates[0].update.cost).toBeUndefined();
+    expect(usageUpdates[1].update.used).toBe(1800);
+    expect(usageUpdates[1].update.cost).toBeUndefined();
+    expect(usageUpdates[2].update.used).toBe(1800);
+    expect(usageUpdates[2].update.cost).toBeDefined();
+  });
+
+  it("mid-stream size is inferred from a 1M model name before the first result", async () => {
+    // On the very first prompt there is no learned context window yet, so the
+    // mid-stream update would otherwise fall back to 200k. A "-1m" suffix in
+    // the SDK model ID is enough signal to emit 1_000_000 up front.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-6-1m",
+        usage: {
+          input_tokens: 2000,
+          output_tokens: 1000,
+          cache_read_input_tokens: 0,
+          cache_creation_input_tokens: 0,
+        },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-6-1m": {
+            inputTokens: 2000,
+            outputTokens: 1000,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.02,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.size).toBe(1000000);
+    expect(usageUpdates[1].update.size).toBe(1000000);
+  });
+
+  it("duplicate stream_event totals do not re-emit usage_update", async () => {
+    // A message_delta whose cumulative totals match the prior snapshot should
+    // not trigger a duplicate usage_update — only the result adds cost on top.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-20250514",
+        usage: {
+          input_tokens: 1000,
+          output_tokens: 500,
+          cache_read_input_tokens: 200,
+          cache_creation_input_tokens: 100,
+        },
+      }),
+      createStreamEvent("message_delta", {
+        usage: { output_tokens: 500 },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 1000,
+            outputTokens: 500,
+            cacheReadInputTokens: 200,
+            cacheCreationInputTokens: 100,
+            webSearchRequests: 0,
+            costUSD: 0.01,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.used).toBe(1800);
+    expect(usageUpdates[0].update.cost).toBeUndefined();
+    expect(usageUpdates[1].update.used).toBe(1800);
+    expect(usageUpdates[1].update.cost).toBeDefined();
+  });
+
+  it("mid-stream size uses the session's learned context window", async () => {
+    // Session state persists the model's context window across prompts, so a
+    // mid-stream update in a later prompt reports the real size immediately
+    // instead of snapping back to the 200k default before the result arrives.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-20250514",
+        usage: {
+          input_tokens: 2000,
+          output_tokens: 1000,
+          cache_read_input_tokens: 0,
+          cache_creation_input_tokens: 0,
+        },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 2000,
+            outputTokens: 1000,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.02,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+    // Simulate a prior prompt having learned the 1M window for this model.
+    agent.sessions["test-session"].contextWindowSize = 1000000;
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.size).toBe(1000000);
+    expect(usageUpdates[1].update.size).toBe(1000000);
+  });
+
+  it("switching to a 1M model seeds the context window from the heuristic", async () => {
+    // The heuristic runs at config-change time so mid-stream updates in the
+    // next prompt already report 1M — without waiting for message_start or
+    // the next `result` to correct us.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-6-1m",
+        usage: {
+          input_tokens: 2000,
+          output_tokens: 1000,
+          cache_read_input_tokens: 0,
+          cache_creation_input_tokens: 0,
+        },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-6-1m": {
+            inputTokens: 2000,
+            outputTokens: 1000,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.02,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+    const session = agent.sessions["test-session"];
+    expect(session.contextWindowSize).toBe(200000);
+
+    (agent as any).syncSessionConfigState(session, "model", "claude-opus-4-6-1m");
+    expect(session.contextWindowSize).toBe(1000000);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.size).toBe(1000000);
+    expect(usageUpdates[1].update.size).toBe(1000000);
+  });
+
+  it("result with no matching modelUsage preserves the learned window", async () => {
+    // A turn whose `result.modelUsage` doesn't contain the current top-level
+    // model (e.g. no top-level assistant message, or only a subagent ran) must
+    // not clobber the window learned on a prior turn — otherwise the next
+    // prompt's mid-stream updates regress to the 200k default.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-haiku-4-5-20251001": {
+            inputTokens: 100,
+            outputTokens: 50,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.001,
+            contextWindow: 200000,
+            maxOutputTokens: 8192,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+    const session = agent.sessions["test-session"];
+    session.contextWindowSize = 1000000;
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    expect(session.contextWindowSize).toBe(1000000);
+    // The emit itself falls back to session.contextWindowSize, which is
+    // unchanged from the learned value.
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    // No lastAssistantTotalUsage was set (no top-level assistant / stream
+    // event), so the result branch skips its emit entirely.
+    expect(usageUpdates).toHaveLength(0);
+  });
+
+  it("switching the session's model invalidates the learned context window", async () => {
+    // When the user switches models mid-session, the window learned for the
+    // previous model would otherwise persist into the next prompt's first
+    // mid-stream update. syncSessionConfigState should reset it so the next
+    // turn's first update falls back to the heuristic (here: 200k default).
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-sonnet-4-6",
+        usage: {
+          input_tokens: 1000,
+          output_tokens: 500,
+          cache_read_input_tokens: 0,
+          cache_creation_input_tokens: 0,
+        },
+      }),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-sonnet-4-6": {
+            inputTokens: 1000,
+            outputTokens: 500,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.01,
+            contextWindow: 200000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+    const session = agent.sessions["test-session"];
+    session.contextWindowSize = 1000000;
+    session.models = { ...session.models, currentModelId: "claude-opus-4-6-1m" };
+
+    // User flips the selector to a 200k model.
+    (agent as any).syncSessionConfigState(session, "model", "claude-sonnet-4-6");
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(2);
+    expect(usageUpdates[0].update.size).toBe(200000);
+    expect(usageUpdates[1].update.size).toBe(200000);
+  });
+
+  it("non-usage stream events do not re-emit usage_update", async () => {
+    // content_block_* and message_stop carry no usage fields; they must not
+    // trigger duplicate emits between the real message_start / message_delta
+    // / result updates.
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent("message_start", {
+        model: "claude-opus-4-20250514",
+        usage: {
+          input_tokens: 1000,
+          output_tokens: 0,
+          cache_read_input_tokens: 0,
+          cache_creation_input_tokens: 0,
+        },
+      }),
+      {
+        type: "stream_event" as const,
+        parent_tool_use_id: null,
+        uuid: randomUUID(),
+        session_id: "test-session",
+        event: { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } },
+      },
+      {
+        type: "stream_event" as const,
+        parent_tool_use_id: null,
+        uuid: randomUUID(),
+        session_id: "test-session",
+        event: { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hi" } },
+      },
+      {
+        type: "stream_event" as const,
+        parent_tool_use_id: null,
+        uuid: randomUUID(),
+        session_id: "test-session",
+        event: { type: "content_block_stop", index: 0 },
+      },
+      createStreamEvent("message_delta", {
+        usage: { output_tokens: 200 },
+      }),
+      {
+        type: "stream_event" as const,
+        parent_tool_use_id: null,
+        uuid: randomUUID(),
+        session_id: "test-session",
+        event: { type: "message_stop" },
+      },
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-opus-4-20250514": {
+            inputTokens: 1000,
+            outputTokens: 200,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.01,
+            contextWindow: 1000000,
+            maxOutputTokens: 16384,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    // Exactly three: message_start (1000), message_delta (1200), result (1200 + cost).
+    expect(usageUpdates).toHaveLength(3);
+    expect(usageUpdates[0].update.used).toBe(1000);
+    expect(usageUpdates[1].update.used).toBe(1200);
+    expect(usageUpdates[2].update.used).toBe(1200);
+    expect(usageUpdates[2].update.cost).toBeDefined();
+  });
+
+  it("subagent stream_event does not emit usage_update", async () => {
+    const { agent, updates } = createMockAgentWithCapture();
+    injectSession(agent, [
+      createStreamEvent(
+        "message_start",
+        {
+          model: "claude-haiku-4-5-20251001",
+          usage: {
+            input_tokens: 500,
+            output_tokens: 100,
+            cache_read_input_tokens: 0,
+            cache_creation_input_tokens: 0,
+          },
+        },
+        "tool_use_123",
+      ),
+      createResultMessageWithModel({
+        modelUsage: {
+          "claude-haiku-4-5-20251001": {
+            inputTokens: 500,
+            outputTokens: 100,
+            cacheReadInputTokens: 0,
+            cacheCreationInputTokens: 0,
+            webSearchRequests: 0,
+            costUSD: 0.001,
+            contextWindow: 200000,
+            maxOutputTokens: 8192,
+          },
+        },
+      }),
+      { type: "system", subtype: "session_state_changed", state: "idle" },
+    ]);
+
+    await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] });
+
+    const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update");
+    expect(usageUpdates).toHaveLength(0);
+  });
+
   it("size reflects the current model's context window, not min across all", async () => {
     const { agent, updates } = createMockAgentWithCapture();
     injectSession(agent, [
@@ -2255,6 +2732,7 @@ describe("emitRawSDKMessages", () => {
       nextPendingOrder: 0,
       abortController: new AbortController(),
       emitRawSDKMessages,
+      contextWindowSize: 200000,
     };
   }
 

src/tests/session-config-options.test.ts

@@ -93,7 +93,10 @@ describe("session config options", () => {
       cancelled: false,
       permissionMode: "default",
       settingsManager: {},
+      modes: structuredClone(MOCK_MODES),
+      models: structuredClone(MOCK_MODELS),
       configOptions: structuredClone(MOCK_CONFIG_OPTIONS),
+      contextWindowSize: 200000,
     };
   }