Keep Claude's context meter on the active model (#6814)

* fix(claude): keep context meter on the active model

Claude result modelUsage accumulates models across the session. Reading its largest window replaced Haiku capacity with a previous Opus window at turn completion. Select the runtime model entry and retain the seeded limit when that entry is absent or invalid.

* refactor(claude): name context fixture inputs and simplify result usage

---------

Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com>

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);

packages/server/src/server/agent/providers/claude/agent.context-window.test.ts

@@ -6,12 +6,28 @@ 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,
+  reportedLimit,
+  modelWindowsByTurn,
+}: {
+  runtimeModel: string;
+  reportedLimit: number;
+  modelWindowsByTurn?: Array<Record<string, number>>;
+}): 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) {
+        const modelWindows = modelWindowsByTurn?.[turn] ?? { [model]: reportedLimit };
+        turn += 1;
+        const usageEntries = Object.entries(modelWindows).map(([id, contextWindow]) => [
+          id,
+          { contextWindow },
+        ]);
+        const modelUsage = Object.fromEntries(usageEntries);
         yield {
           type: "system",
           subtype: "init",
@@ -41,7 +57,7 @@ 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,
           permission_denials: [],
           uuid: "context-result",
         };
@@ -71,7 +87,7 @@ test.each([
   async ({ selected, runtime, expected }) => {
     const client = new ClaudeAgentClient({
       logger: createTestLogger(),
-      queryFactory: contextQuery(runtime, 800_000),
+      queryFactory: contextQuery({ runtimeModel: runtime, reportedLimit: 800_000 }),
       resolveBinary: async () => process.execPath,
     });
     const session = await client.createSession({
@@ -110,3 +126,93 @@ 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({
+      runtimeModel: haiku,
+      reportedLimit: 200_000,
+      modelWindowsByTurn: [
+        { [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({
+      runtimeModel: model,
+      reportedLimit: expected,
+      modelWindowsByTurn: [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();
+  }
+});

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,10 @@ class ClaudeContextUsageState {
     return this.createUsageUpdatedEvent(usedTokens);
   }
 
-  buildResultUsage(message: SDKResultMessage, modelUsage: unknown): AgentUsage | undefined {
+  buildResultUsage(
+    message: SDKResultMessage,
+    modelId: string | null | undefined,
+  ): AgentUsage | undefined {
     try {
       if (!message.usage) {
         return undefined;
@@ -2062,7 +2057,7 @@ class ClaudeContextUsageState {
         totalCostUsd: message.total_cost_usd,
       };
 
-      const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage);
+      const modelContextWindowMaxTokens = this.recordModelUsage(message.modelUsage, modelId);
       if (this.contextWindowMaxTokens !== undefined) {
         usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
       } else if (modelContextWindowMaxTokens !== undefined) {
@@ -4696,7 +4691,7 @@ class ClaudeAgentSession implements AgentSession {
     message: Extract<SDKMessage, { type: "result" }>,
     events: AgentStreamEvent[],
   ): void {
-    const usage = this.convertUsage(message, message.modelUsage);
+    const usage = this.convertUsage(message);
     if (message.subtype === "success") {
       events.push(...this.sidechainTracker.finishAll("completed"));
       // Built-in slash commands (e.g. /voice, /usage, "Unknown command: …")
@@ -4878,8 +4873,8 @@ class ClaudeAgentSession implements AgentSession {
     return null;
   }
 
-  private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined {
-    return this.contextUsage.buildResultUsage(message, modelUsage);
+  private convertUsage(message: SDKResultMessage): AgentUsage | undefined {
+    return this.contextUsage.buildResultUsage(message, this.lastRuntimeModel ?? this.config.model);
   }
 
   private handlePermissionRequest: CanUseTool = async (