fix(claude): show custom model context limits during the first turn (#6767)

Seed active usage from the explicit 1M modifier and the model resolved by Claude Code initialization. Preserve SDK-reported limits on repeated initialization and reset them when the selected model changes.

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

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,
+            }),
           }),
         );
       }

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<ClaudeQueryFactory>;
+  };
+}
+
+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();
+    }
+  },
+);

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,
           },
         }),

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 },

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[] = [];
 

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,