Show live context usage during active turns

packages/server/src/server/agent/providers/claude/agent.redesign.test.ts

@@ -14,6 +14,7 @@ interface QueryMock {
   close: ReturnType<typeof vi.fn>;
   setPermissionMode: ReturnType<typeof vi.fn>;
   setModel: ReturnType<typeof vi.fn>;
+  getContextUsage: ReturnType<typeof vi.fn>;
   supportedModels: ReturnType<typeof vi.fn>;
   supportedCommands: ReturnType<typeof vi.fn>;
   rewindFiles: ReturnType<typeof vi.fn>;
@@ -53,6 +54,7 @@ function createBaseQueryMock(nextImpl: QueryMock["next"]): QueryMock {
     close: vi.fn(() => undefined),
     setPermissionMode: vi.fn(async () => undefined),
     setModel: vi.fn(async () => undefined),
+    getContextUsage: vi.fn(async () => undefined),
     supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]),
     supportedCommands: vi.fn(async () => []),
     rewindFiles: vi.fn(async () => ({ canRewind: true })),
@@ -1078,39 +1080,52 @@ test("reuses one autonomous run for unbound stream_event bursts with no foregrou
   const internal: {
     turnState: "idle" | "foreground" | "autonomous";
     nextTurnOrdinal: number;
-    routeSdkMessageFromPump: (message: Record<string, unknown>) => void;
+    routeSdkMessageFromPump: (
+      message: Record<string, unknown>,
+      activeQuery: QueryMock,
+    ) => Promise<void>;
     autonomousTurn: { id: string } | null;
   } = asInternals(session);
+  const queryMock = createBaseQueryMock(vi.fn(async () => ({ done: true, value: undefined })));
 
   internal.turnState = "idle";
-  internal.routeSdkMessageFromPump({
+  await internal.routeSdkMessageFromPump(
-    type: "stream_event",
+    {
-    event: {
+      type: "stream_event",
-      type: "content_block_delta",
+      event: {
-      delta: { type: "text_delta", text: "AUTO " },
+        type: "content_block_delta",
+        delta: { type: "text_delta", text: "AUTO " },
+      },
     },
-  });
+    queryMock,
+  );
 
   const firstRunId = internal.autonomousTurn?.id ?? null;
   expect(firstRunId).toBe("autonomous-turn-1");
   expect(internal.nextTurnOrdinal).toBe(2);
 
-  internal.routeSdkMessageFromPump({
+  await internal.routeSdkMessageFromPump(
-    type: "stream_event",
+    {
-    event: {
+      type: "stream_event",
-      type: "content_block_delta",
+      event: {
-      delta: { type: "text_delta", text: "WAKE" },
+        type: "content_block_delta",
+        delta: { type: "text_delta", text: "WAKE" },
+      },
     },
-  });
+    queryMock,
+  );
   expect(internal.autonomousTurn?.id).toBe(firstRunId);
   expect(internal.nextTurnOrdinal).toBe(2);
 
-  internal.routeSdkMessageFromPump({
+  await internal.routeSdkMessageFromPump(
-    type: "result",
+    {
-    subtype: "success",
+      type: "result",
-    usage: buildUsage(),
+      subtype: "success",
-    total_cost_usd: 0,
+      usage: buildUsage(),
-  });
+      total_cost_usd: 0,
+    },
+    queryMock,
+  );
   expect(internal.autonomousTurn).toBeNull();
 
   await session.close();

packages/server/src/server/agent/providers/claude/agent.test.ts

@@ -561,6 +561,7 @@ describe("ClaudeAgentSession features", () => {
       return: queryReturn,
       applyFlagSettings: vi.fn(async () => undefined),
       setModel: vi.fn(async () => undefined),
+      getContextUsage: vi.fn(async () => undefined),
       [Symbol.asyncIterator](): AsyncIterator<SDKMessage, void> {
         return {
           next: async () => {
@@ -1046,6 +1047,7 @@ describe("ClaudeAgentSession context window usage", () => {
 
   interface QueryFactoryForTurnsOptions {
     currentContextUsageByTurn?: Array<Record<string, unknown> | undefined>;
+    model?: string;
   }
 
   async function createSessionForTest(): Promise<TestClaudeSession> {
@@ -1069,6 +1071,7 @@ describe("ClaudeAgentSession context window usage", () => {
     return await client.createSession({
       provider: "claude",
       cwd: process.cwd(),
+      model: options?.model,
     });
   }
 
@@ -1353,6 +1356,7 @@ describe("ClaudeAgentSession context window usage", () => {
         close: () => undefined,
         setPermissionMode: async () => undefined,
         setModel: async () => undefined,
+        getContextUsage: async () => undefined,
         supportedModels: async () => [],
         supportedCommands: async () => [
           {
@@ -1745,6 +1749,30 @@ describe("ClaudeAgentSession context window usage", () => {
     }
   });
 
+  test("selected Claude models seed active context window usage with max tokens", async () => {
+    const session = await createSessionForTurns(
+      [[createInitMessage(), createMessageStartEvent(), createSuccessResult()]],
+      { model: "claude-sonnet-4-6" },
+    );
+
+    try {
+      const events = await collectStreamEvents(session);
+
+      expect(events).toContainEqual(
+        expect.objectContaining({
+          type: "usage_updated",
+          provider: "claude",
+          usage: {
+            contextWindowMaxTokens: 200_000,
+            contextWindowUsedTokens: 150,
+          },
+        }),
+      );
+    } finally {
+      await session.close();
+    }
+  });
+
   test("message_delta stream events update per-request usage", async () => {
     const session = await createSessionForTurns([
       [

packages/server/src/server/agent/providers/claude/agent.ts

@@ -1603,6 +1603,23 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
   return maxContextWindow;
 }
 
+function resolveInitialContextWindowSize(modelId: string | null | undefined): number | undefined {
+  const normalized = typeof modelId === "string" ? modelId.trim().toLowerCase() : "";
+  if (!normalized) {
+    return undefined;
+  }
+  if (normalized.includes("[1m]") || normalized.includes("context-1m")) {
+    return 1_000_000;
+  }
+  if (normalized.includes("claude-fable-5")) {
+    return 1_000_000;
+  }
+  if (/(?:^|[~/_-])(?:claude[-_ ]*)?(opus|sonnet|haiku)(?:$|[-_ ./])/.test(normalized)) {
+    return 200_000;
+  }
+  return undefined;
+}
+
 function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
   const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage);
   if (!messageUsage) {
@@ -1722,11 +1739,19 @@ class ClaudeContextUsageState {
   private streamRequestOutputTokens: number | undefined;
   private completedResultTurns = 0;
 
+  constructor(initialContextWindowMaxTokens?: number) {
+    this.contextWindowMaxTokens = initialContextWindowMaxTokens;
+  }
+
   beginTurn(): void {
     this.streamRequestInputTokens = undefined;
     this.streamRequestOutputTokens = undefined;
   }
 
+  setInitialContextWindowMaxTokens(contextWindowMaxTokens: number | undefined): void {
+    this.contextWindowMaxTokens = contextWindowMaxTokens;
+  }
+
   recordModelUsage(modelUsage: unknown): number | undefined {
     const contextWindowMaxTokens = extractContextWindowSize(modelUsage);
     if (contextWindowMaxTokens !== undefined) {
@@ -1879,7 +1904,7 @@ class ClaudeAgentSession implements AgentSession {
   private pendingInterruptAbort = false;
   private foregroundHasVisibleActivity = false;
   private activeTurnHasAssistantText = false;
-  private readonly contextUsage = new ClaudeContextUsageState();
+  private readonly contextUsage: ClaudeContextUsageState;
   private userMessageIds: string[] = [];
   private readonly emittedUserMessageIds = new Set<string>();
   private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = [];
@@ -1897,6 +1922,9 @@ class ClaudeAgentSession implements AgentSession {
     this.logger = options.logger.child({ agentId: this.agentId });
     this.queryFactory = options.queryFactory;
     this.resolveBinary = options.resolveBinary;
+    this.contextUsage = new ClaudeContextUsageState(
+      resolveInitialContextWindowSize(this.config.model),
+    );
     const handle = options.handle;
 
     if (handle) {
@@ -2140,6 +2168,9 @@ class ClaudeAgentSession implements AgentSession {
     if (!claudeModelSupportsFastMode(this.config.model) && this.config.featureValues?.fast_mode) {
       await this.applyFastModeFeature(false, activeQuery);
     }
+    this.contextUsage.setInitialContextWindowMaxTokens(
+      resolveInitialContextWindowSize(this.config.model),
+    );
     this.lastOptionsModel = normalizedModelId ?? this.lastOptionsModel;
     this.lastRuntimeModel = null;
     this.cachedRuntimeInfo = null;
@@ -3455,14 +3486,8 @@ class ClaudeAgentSession implements AgentSession {
   private async queryCurrentContextUsage(
     activeQuery: Query,
   ): Promise<ClaudeCurrentContextUsage | undefined> {
-    const query = activeQuery as Query & {
-      getContextUsage?: () => Promise<SDKControlGetContextUsageResponse>;
-    };
-    if (typeof query.getContextUsage !== "function") {
-      return undefined;
-    }
     try {
-      const usage = await withTimeout(query.getContextUsage(), 3_000, "timeout");
+      const usage = await withTimeout(activeQuery.getContextUsage(), 3_000, "timeout");
       return readCurrentContextUsage(usage);
     } catch (error) {
       this.logger.debug({ err: error }, "Claude context usage query failed");

packages/server/src/server/daemon-e2e/claude-live-usage.e2e.test.ts

@@ -0,0 +1,187 @@
+import { mkdtempSync } from "node:fs";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import pino from "pino";
+import { describe, expect, test, vi } from "vitest";
+
+import { ClaudeAgentClient } from "../agent/providers/claude/agent.js";
+import { DaemonClient } from "../test-utils/daemon-client.js";
+import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js";
+
+function deferred(): { promise: Promise<void>; resolve: () => void } {
+  let resolve!: () => void;
+  const promise = new Promise<void>((innerResolve) => {
+    resolve = innerResolve;
+  });
+  return { promise, resolve };
+}
+
+function createControlledClaudeQueryFactory(resultGate: Promise<void>) {
+  return vi.fn(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
+    const queuedMessages: Array<Record<string, unknown>> = [];
+    const waiters: Array<() => void> = [];
+    let closed = false;
+
+    function wakeNextWaiter() {
+      waiters.shift()?.();
+    }
+
+    function enqueue(message: Record<string, unknown>) {
+      queuedMessages.push(message);
+      wakeNextWaiter();
+    }
+
+    void (async () => {
+      for await (const _message of prompt) {
+        enqueue({
+          type: "system",
+          subtype: "init",
+          session_id: "session-live-usage",
+          permissionMode: "default",
+          model: "claude-sonnet-4-6",
+        });
+        enqueue({
+          type: "stream_event",
+          event: {
+            type: "message_start",
+            message: {
+              usage: {
+                input_tokens: 100,
+                cache_creation_input_tokens: 20,
+                cache_read_input_tokens: 30,
+              },
+            },
+          },
+          session_id: "session-live-usage",
+        });
+        enqueue({
+          type: "stream_event",
+          event: {
+            type: "message_delta",
+            usage: {
+              output_tokens: 25,
+            },
+          },
+          session_id: "session-live-usage",
+        });
+        await resultGate;
+        enqueue({
+          type: "result",
+          subtype: "success",
+          duration_ms: 100,
+          duration_api_ms: 75,
+          is_error: false,
+          num_turns: 1,
+          result: "done",
+          stop_reason: null,
+          total_cost_usd: 0.25,
+          usage: {
+            input_tokens: 10,
+            cache_read_input_tokens: 5,
+            output_tokens: 7,
+          },
+          modelUsage: {
+            "claude-sonnet-4-6": { contextWindow: 200_000 },
+          },
+          permission_denials: [],
+          uuid: "result-live-usage",
+          session_id: "session-live-usage",
+        });
+        break;
+      }
+      closed = true;
+      wakeNextWaiter();
+    })();
+
+    return {
+      next: vi.fn(async () => {
+        for (;;) {
+          if (queuedMessages.length > 0 || closed) {
+            break;
+          }
+          await new Promise<void>((resolve) => {
+            waiters.push(resolve);
+          });
+        }
+        if (queuedMessages.length === 0) {
+          return { done: true, value: undefined };
+        }
+        return { done: false, value: queuedMessages.shift() };
+      }),
+      interrupt: vi.fn(async () => undefined),
+      return: vi.fn(async () => {
+        closed = true;
+        wakeNextWaiter();
+        return undefined;
+      }),
+      close: vi.fn(() => {
+        closed = true;
+        wakeNextWaiter();
+      }),
+      setPermissionMode: vi.fn(async () => undefined),
+      setModel: vi.fn(async () => undefined),
+      getContextUsage: vi.fn(async () => undefined),
+      supportedModels: vi.fn(async () => []),
+      supportedCommands: vi.fn(async () => []),
+      rewindFiles: vi.fn(async () => ({ canRewind: true })),
+      [Symbol.asyncIterator]() {
+        return this;
+      },
+    };
+  });
+}
+
+describe("daemon E2E (claude live usage)", () => {
+  test("publishes renderable active context usage through agent_update", async () => {
+    const logger = pino({ level: "silent" });
+    const cwd = mkdtempSync(path.join(tmpdir(), "paseo-claude-live-usage-"));
+    const resultGate = deferred();
+    const daemon = await createTestPaseoDaemon({
+      agentClients: {
+        claude: new ClaudeAgentClient({
+          logger,
+          queryFactory: createControlledClaudeQueryFactory(resultGate.promise),
+          resolveBinary: async () => "/test/claude/bin",
+        }),
+      },
+      logger,
+    });
+    const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` });
+
+    try {
+      await client.connect();
+      await client.fetchAgents({ subscribe: { subscriptionId: "claude-live-usage" } });
+
+      const agent = await client.createAgent({
+        provider: "claude",
+        model: "claude-sonnet-4-6",
+        cwd,
+        title: "Claude live usage",
+      });
+
+      await client.sendMessage(agent.id, "Report live usage");
+
+      const activeSnapshot = await client.waitForAgentUpsert(
+        agent.id,
+        (snapshot) =>
+          snapshot.status === "running" &&
+          snapshot.lastUsage?.contextWindowMaxTokens === 200_000 &&
+          snapshot.lastUsage.contextWindowUsedTokens === 175,
+        10_000,
+      );
+
+      expect(activeSnapshot.lastUsage).toEqual({
+        contextWindowMaxTokens: 200_000,
+        contextWindowUsedTokens: 175,
+      });
+
+      resultGate.resolve();
+      const finalState = await client.waitForFinish(agent.id, 10_000);
+      expect(finalState.status).toBe("idle");
+    } finally {
+      resultGate.resolve();
+      await client.close().catch(() => undefined);
+      await daemon.close();
+    }
+  });
+});