Fix subagent context usage reporting

Task-level usage reports processed work, not the parent session's active context. Prefer current context data and parent stream usage before falling back to result totals.

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

@@ -13,11 +13,11 @@ import {
   normalizeClaudeAskUserQuestionUpdatedInput,
   toClaudeSdkMcpConfig,
 } from "./agent.js";
-import type { AgentTimelineItem, AgentUsage, AgentStreamEvent } from "../../agent-sdk-types.js";
+import { streamSession } from "../test-utils/session-stream-adapter.js";
+import type { AgentSession, AgentTimelineItem, AgentStreamEvent } from "../../agent-sdk-types.js";
 
 interface TestClaudeSession {
   translateMessageToEvents(message: SDKMessage): AgentStreamEvent[];
-  convertUsage(message: SDKMessage): AgentUsage | undefined;
 }
 
 afterEach(() => {
@@ -1044,6 +1044,10 @@ describe("ClaudeAgentClient.listImportableSessions", () => {
 describe("ClaudeAgentSession context window usage", () => {
   const logger = createTestLogger();
 
+  interface QueryFactoryForTurnsOptions {
+    currentContextUsageByTurn?: Array<Record<string, unknown> | undefined>;
+  }
+
   async function createSessionForTest(): Promise<TestClaudeSession> {
     const client = new ClaudeAgentClient({ logger, resolveBinary: async () => "/test/claude/bin" });
     const session = await client.createSession({
@@ -1053,11 +1057,38 @@ describe("ClaudeAgentSession context window usage", () => {
     return session as unknown as TestClaudeSession;
   }
 
-  function createQueryFactoryForTurns(turns: Array<Array<Record<string, unknown>>>) {
+  async function createSessionForTurns(
+    turns: Array<Array<Record<string, unknown>>>,
+    options?: QueryFactoryForTurnsOptions,
+  ): Promise<AgentSession> {
+    const client = new ClaudeAgentClient({
+      logger,
+      queryFactory: createQueryFactoryForTurns(turns, options),
+      resolveBinary: async () => "/test/claude/bin",
+    });
+    return await client.createSession({
+      provider: "claude",
+      cwd: process.cwd(),
+    });
+  }
+
+  async function collectStreamEvents(session: AgentSession, prompt = "turn") {
+    const events: AgentStreamEvent[] = [];
+    for await (const event of streamSession(session, prompt)) {
+      events.push(event);
+    }
+    return events;
+  }
+
+  function createQueryFactoryForTurns(
+    turns: Array<Array<Record<string, unknown>>>,
+    options?: QueryFactoryForTurnsOptions,
+  ) {
     return vi.fn(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
       const queuedMessages: Array<Record<string, unknown>> = [];
       const waiters: Array<() => void> = [];
       let turnIndex = 0;
+      let contextUsageIndex = 0;
       const closedRef = { value: false };
 
       function wakeNextWaiter() {
@@ -1106,6 +1137,11 @@ describe("ClaudeAgentSession context window usage", () => {
         }),
         setPermissionMode: vi.fn(async () => undefined),
         setModel: vi.fn(async () => undefined),
+        getContextUsage: vi.fn(async () => {
+          const usage = options?.currentContextUsageByTurn?.[contextUsageIndex];
+          contextUsageIndex += 1;
+          return usage;
+        }),
         supportedModels: vi.fn(async () => []),
         supportedCommands: vi.fn(async () => []),
         rewindFiles: vi.fn(async () => ({ canRewind: true })),
@@ -1116,6 +1152,133 @@ describe("ClaudeAgentSession context window usage", () => {
     });
   }
 
+  function createInitMessage(sessionId = "session-1"): Record<string, unknown> {
+    return {
+      type: "system",
+      subtype: "init",
+      session_id: sessionId,
+      permissionMode: "default",
+      model: "claude-sonnet-4-6",
+    };
+  }
+
+  function createClaudeCurrentContextUsage(
+    totalTokens: number,
+    maxTokens: number,
+  ): Record<string, unknown> {
+    return {
+      categories: [],
+      totalTokens,
+      maxTokens,
+      rawMaxTokens: maxTokens,
+      percentage: totalTokens / maxTokens,
+      gridRows: [],
+      model: "claude-sonnet-4-6",
+      memoryFiles: [],
+      mcpTools: [],
+      agents: [],
+      isAutoCompactEnabled: true,
+      apiUsage: null,
+    };
+  }
+
+  function createSuccessResult(overrides: Record<string, unknown> = {}): Record<string, unknown> {
+    return {
+      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-1",
+      session_id: "session-1",
+      ...overrides,
+    };
+  }
+
+  function createMessageStartEvent(
+    usage: Record<string, unknown> = {
+      input_tokens: 100,
+      cache_creation_input_tokens: 20,
+      cache_read_input_tokens: 30,
+    },
+  ): Record<string, unknown> {
+    return {
+      type: "stream_event",
+      event: {
+        type: "message_start",
+        message: { usage },
+      },
+      session_id: "session-1",
+    };
+  }
+
+  function createMessageDeltaEvent(outputTokens: number): Record<string, unknown> {
+    return {
+      type: "stream_event",
+      event: {
+        type: "message_delta",
+        usage: {
+          output_tokens: outputTokens,
+        },
+      },
+      session_id: "session-1",
+    };
+  }
+
+  function createAgentToolStartEvent(): Record<string, unknown> {
+    return {
+      type: "stream_event",
+      event: {
+        type: "content_block_start",
+        index: 0,
+        content_block: {
+          type: "tool_use",
+          id: "toolu-agent-1",
+          name: "Agent",
+          input: {
+            description: "Check something in a subagent",
+            prompt: "Return a short answer",
+          },
+        },
+      },
+      session_id: "session-1",
+    };
+  }
+
+  function createSubagentTaskNotification(): Record<string, unknown> {
+    return {
+      type: "system",
+      subtype: "task_notification",
+      uuid: "task-note-1",
+      task_id: "task-1",
+      tool_use_id: "toolu-agent-1",
+      status: "running",
+      summary: "Subagent is working",
+      usage: {
+        total_tokens: 18_876,
+        tool_uses: 1,
+        duration_ms: 50,
+        input_tokens: 12_000,
+        cache_read_input_tokens: 6_000,
+        output_tokens: 876,
+      },
+      session_id: "session-1",
+    };
+  }
+
   test("passes persistSession through to the Claude SDK query options", async () => {
     const createResultTurn = (sessionId: string) => [
       {
@@ -1389,357 +1552,152 @@ describe("ClaudeAgentSession context window usage", () => {
     }
   });
 
-  test("convertUsage includes contextWindowMaxTokens and derives used tokens from result usage as initial fallback", async () => {
+  test("reports Claude's current context usage after an Agent subagent runs", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns(
-
+      [
-    const usage = session.convertUsage(
+        [
+          createInitMessage(),
+          createMessageStartEvent(),
+          createAgentToolStartEvent(),
+          createSubagentTaskNotification(),
+          createMessageDeltaEvent(25),
+          createSuccessResult({
+            usage: {
+              input_tokens: 9_000,
+              cache_creation_input_tokens: 300,
+              cache_read_input_tokens: 700,
+              output_tokens: 400,
+            },
+          }),
+        ],
+      ],
       {
-        type: "result",
+        currentContextUsageByTurn: [createClaudeCurrentContextUsage(12_345, 200_000)],
-        subtype: "success",
-        usage: {
-          input_tokens: 10,
-          cache_read_input_tokens: 5,
-          output_tokens: 7,
-        },
-        total_cost_usd: 0.12,
-      },
-      {
-        "claude-sonnet-4-6": { contextWindow: 200_000 },
-        "claude-opus-4-6": { contextWindow: 1_000_000 },
       },
     );
 
-    expect(usage).toEqual({
+    try {
-      inputTokens: 10,
+      const result = await session.run("turn");
-      cachedInputTokens: 5,
+
-      outputTokens: 7,
+      expect(result.usage).toEqual({
-      totalCostUsd: 0.12,
+        inputTokens: 9_000,
-      contextWindowMaxTokens: 1_000_000,
+        cachedInputTokens: 700,
-      contextWindowUsedTokens: 22,
+        outputTokens: 400,
-    });
+        totalCostUsd: 0.25,
+        contextWindowMaxTokens: 200_000,
+        contextWindowUsedTokens: 12_345,
+      });
+    } finally {
+      await session.close();
+    }
   });
 
-  test("contextWindowUsedTokens falls back to result usage when no task_progress was received", async () => {
+  test("does not report task notification tokens as parent context usage", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageStartEvent(),
+        createAgentToolStartEvent(),
+        createSubagentTaskNotification(),
+        {
+          type: "system",
+          subtype: "task_progress",
+          task_id: "task-1",
+          description: "Subagent progress",
+          usage: {
+            total_tokens: 9_999,
+            tool_uses: 1,
+            duration_ms: 50,
+          },
+          uuid: "task-progress-1",
+          session_id: "session-1",
+        },
+        createMessageDeltaEvent(25),
+        createSuccessResult(),
+      ],
+    ]);
 
-    const usage = session.convertUsage({
+    try {
-      type: "result",
+      const result = await session.run("turn");
-      subtype: "success",
-      usage: {
-        input_tokens: 10,
-        cache_creation_input_tokens: 3,
-        cache_read_input_tokens: 5,
-        output_tokens: 7,
-      },
-      total_cost_usd: 0.12,
-    });
 
-    expect(usage).toEqual({
+      expect(result.usage).toEqual({
-      inputTokens: 10,
-      cachedInputTokens: 5,
-      outputTokens: 7,
-      totalCostUsd: 0.12,
-      contextWindowUsedTokens: 25,
-    });
-  });
-
-  test("contextWindowUsedTokens is populated from task_progress usage data", async () => {
-    const session = await createSessionForTest();
-
-    session.translateMessageToEvents({
-      type: "system",
-      subtype: "task_progress",
-      task_id: "task-1",
-      description: "Processing",
-      usage: {
-        total_tokens: 999,
-        tool_uses: 1,
-        duration_ms: 50,
-        input_tokens: 345,
-        cache_read_input_tokens: 55,
-      },
-      uuid: "task-progress-1",
-      session_id: "session-1",
-    });
-
-    const events = session.translateMessageToEvents({
-      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-1",
-      session_id: "session-1",
-    });
-
-    expect(events).toContainEqual({
-      type: "turn_completed",
-      provider: "claude",
-      usage: {
         inputTokens: 10,
         cachedInputTokens: 5,
         outputTokens: 7,
         totalCostUsd: 0.25,
         contextWindowMaxTokens: 200_000,
-        contextWindowUsedTokens: 999,
-      },
-    });
-  });
-
-  test("task_progress emits a usage_updated event", async () => {
-    const session = await createSessionForTest();
-
-    const events = session.translateMessageToEvents({
-      type: "system",
-      subtype: "task_progress",
-      task_id: "task-1",
-      description: "Processing",
-      usage: {
-        total_tokens: 999,
-        tool_uses: 1,
-        duration_ms: 50,
-      },
-      uuid: "task-progress-1",
-      session_id: "session-1",
-    });
-
-    expect(events).toContainEqual({
-      type: "usage_updated",
-      provider: "claude",
-      usage: {
-        contextWindowUsedTokens: 999,
-      },
-    });
-  });
-
-  test("task_notification emits a usage_updated event", async () => {
-    const session = await createSessionForTest();
-
-    const events = session.translateMessageToEvents({
-      type: "system",
-      subtype: "task_notification",
-      uuid: "task-note-1",
-      task_id: "task-1",
-      status: "running",
-      summary: "Background task still running",
-      usage: {
-        total_tokens: 777,
-        tool_uses: 1,
-        duration_ms: 50,
-      },
-      session_id: "session-1",
-    } as unknown as SDKMessage);
-
-    expect(events).toContainEqual({
-      type: "usage_updated",
-      provider: "claude",
-      usage: {
-        contextWindowUsedTokens: 777,
-      },
-    });
-  });
-
-  test("message_start stream events emit usage_updated with per-request usage", async () => {
-    const session = await createSessionForTest();
-
-    const events = session.translateMessageToEvents({
-      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-1",
-    } as unknown as SDKMessage);
-
-    expect(events).toContainEqual({
-      type: "usage_updated",
-      provider: "claude",
-      usage: {
-        contextWindowUsedTokens: 150,
-      },
-    });
-  });
-
-  test("message_delta stream events update per-request usage", async () => {
-    const session = await createSessionForTest();
-
-    session.translateMessageToEvents({
-      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-1",
-    } as unknown as SDKMessage);
-
-    const events = session.translateMessageToEvents({
-      type: "stream_event",
-      event: {
-        type: "message_delta",
-        usage: {
-          output_tokens: 25,
-        },
-      },
-      session_id: "session-1",
-    } as unknown as SDKMessage);
-
-    expect(events).toContainEqual({
-      type: "usage_updated",
-      provider: "claude",
-      usage: {
         contextWindowUsedTokens: 175,
-      },
+      });
-    });
+    } finally {
+      await session.close();
+    }
   });
 
-  test("task_progress usage takes priority over derived result usage", async () => {
+  test("falls back to the active result iteration when current and stream usage are unavailable", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns([
-
-    session.translateMessageToEvents({
-      type: "system",
-      subtype: "task_progress",
-      task_id: "task-1",
-      description: "Processing",
-      usage: {
-        total_tokens: 999,
-        tool_uses: 1,
-        duration_ms: 50,
-        input_tokens: 345,
-        cache_read_input_tokens: 55,
-      },
-      uuid: "task-progress-1",
-      session_id: "session-1",
-    });
-
-    const usage = session.convertUsage({
-      type: "result",
-      subtype: "success",
-      usage: {
-        input_tokens: 10,
-        cache_creation_input_tokens: 3,
-        cache_read_input_tokens: 5,
-        output_tokens: 7,
-      },
-      total_cost_usd: 0.12,
-    });
-
-    expect(usage).toEqual({
-      inputTokens: 10,
-      cachedInputTokens: 5,
-      outputTokens: 7,
-      totalCostUsd: 0.12,
-      contextWindowUsedTokens: 999,
-    });
-  });
-
-  test("contextWindowUsedTokens persists across turns from last task_progress", async () => {
-    const queryFactory = createQueryFactoryForTurns([
       [
-        {
+        createInitMessage(),
-          type: "system",
+        createSuccessResult({
-          subtype: "init",
-          session_id: "session-1",
-          permissionMode: "default",
-          model: "claude-sonnet-4-6",
-        },
-        {
-          type: "system",
-          subtype: "task_progress",
-          task_id: "task-1",
-          description: "Processing",
           usage: {
-            total_tokens: 999,
+            input_tokens: 5_000,
-            tool_uses: 1,
+            cache_creation_input_tokens: 500,
-            duration_ms: 50,
+            cache_read_input_tokens: 600,
-            input_tokens: 345,
+            output_tokens: 700,
-            cache_read_input_tokens: 55,
+            iterations: [
+              {
+                input_tokens: 100,
+                cache_creation_input_tokens: 10,
+                cache_read_input_tokens: 20,
+                output_tokens: 30,
+              },
+              {
+                input_tokens: 2,
+                cache_creation_input_tokens: 10,
+                cache_read_input_tokens: 100,
+                output_tokens: 5,
+              },
+            ],
           },
-          uuid: "task-progress-1",
+        }),
-          session_id: "session-1",
+      ],
-        },
+    ]);
-        {
+
-          type: "result",
+    try {
-          subtype: "success",
+      const result = await session.run("turn");
-          duration_ms: 100,
+
-          duration_api_ms: 75,
+      expect(result.usage).toEqual({
-          is_error: false,
+        inputTokens: 5_000,
-          num_turns: 1,
+        cachedInputTokens: 600,
-          result: "done",
+        outputTokens: 700,
-          stop_reason: null,
+        totalCostUsd: 0.25,
-          total_cost_usd: 0.25,
+        contextWindowMaxTokens: 200_000,
-          usage: {
+        contextWindowUsedTokens: 117,
-            input_tokens: 10,
+      });
-            cache_read_input_tokens: 5,
+    } finally {
-            output_tokens: 7,
+      await session.close();
-          },
+    }
-          modelUsage: {
+  });
-            "claude-sonnet-4-6": { contextWindow: 200_000 },
+
-          },
+  test("does not use aggregate result totals after the first result turn", async () => {
-          permission_denials: [],
+    const session = await createSessionForTurns([
-          uuid: "result-1",
+      [
-          session_id: "session-1",
+        createInitMessage(),
-        },
+        createMessageStartEvent(),
+        createMessageDeltaEvent(25),
+        createSuccessResult(),
       ],
       [
-        {
+        createSuccessResult({
-          type: "result",
-          subtype: "success",
-          duration_ms: 110,
-          duration_api_ms: 80,
-          is_error: false,
-          num_turns: 1,
-          result: "still done",
-          stop_reason: null,
           total_cost_usd: 0.1,
           usage: {
-            input_tokens: 11,
+            input_tokens: 1_000,
-            cache_creation_input_tokens: 3,
+            cache_read_input_tokens: 200,
-            cache_read_input_tokens: 6,
+            output_tokens: 300,
-            output_tokens: 8,
-          },
-          modelUsage: {
-            "claude-sonnet-4-6": { contextWindow: 200_000 },
           },
-          permission_denials: [],
           uuid: "result-2",
-          session_id: "session-1",
+        }),
-        },
       ],
     ]);
-    const client = new ClaudeAgentClient({
-      logger,
-      queryFactory,
-      resolveBinary: async () => "/test/claude/bin",
-    });
-    const session = await client.createSession({
-      provider: "claude",
-      cwd: process.cwd(),
-    });
 
     try {
       const firstTurn = await session.run("turn 1");
@@ -1751,175 +1709,109 @@ describe("ClaudeAgentSession context window usage", () => {
         outputTokens: 7,
         totalCostUsd: 0.25,
         contextWindowMaxTokens: 200_000,
-        contextWindowUsedTokens: 999,
+        contextWindowUsedTokens: 175,
       });
-      // Turn 2 has no task_progress, so contextWindowUsedTokens retains the
-      // last known value from turn 1 rather than deriving from accumulated
-      // result.usage (which would be incorrect — those are session-level totals).
       expect(secondTurn.usage).toEqual({
-        inputTokens: 11,
+        inputTokens: 1_000,
-        cachedInputTokens: 6,
+        cachedInputTokens: 200,
-        outputTokens: 8,
+        outputTokens: 300,
         totalCostUsd: 0.1,
         contextWindowMaxTokens: 200_000,
-        contextWindowUsedTokens: 999,
       });
     } finally {
       await session.close();
     }
   });
 
-  test("convertUsage derives used tokens from result usage as fallback when task_progress is missing", async () => {
+  test("message_start stream events emit usage_updated with per-request usage", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns([
+      [createInitMessage(), createMessageStartEvent(), createSuccessResult()],
+    ]);
 
-    const usage = session.convertUsage({
+    try {
-      type: "result",
+      const events = await collectStreamEvents(session);
-      subtype: "success",
-      usage: {
-        input_tokens: 10,
-        cache_read_input_tokens: 5,
-        output_tokens: 7,
-      },
-      total_cost_usd: 0.12,
-    });
 
-    expect(usage).toEqual({
+      expect(events).toContainEqual(
-      inputTokens: 10,
+        expect.objectContaining({
-      cachedInputTokens: 5,
+          type: "usage_updated",
-      outputTokens: 7,
+          provider: "claude",
-      totalCostUsd: 0.12,
+          usage: {
-      contextWindowUsedTokens: 22,
+            contextWindowUsedTokens: 150,
-    });
+          },
+        }),
+      );
+    } finally {
+      await session.close();
+    }
   });
 
-  test("convertUsage uses per-request stream usage when no task_progress is available", async () => {
+  test("message_delta stream events update per-request usage", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageStartEvent(),
+        createMessageDeltaEvent(25),
+        createSuccessResult(),
+      ],
+    ]);
 
-    session.translateMessageToEvents({
+    try {
-      type: "stream_event",
+      const events = await collectStreamEvents(session);
-      event: {
+
-        type: "message_start",
+      expect(events).toContainEqual(
-        message: {
+        expect.objectContaining({
+          type: "usage_updated",
+          provider: "claude",
           usage: {
-            input_tokens: 100,
+            contextWindowUsedTokens: 175,
-            cache_creation_input_tokens: 20,
-            cache_read_input_tokens: 30,
           },
-        },
+        }),
-      },
+      );
-      session_id: "session-1",
+    } finally {
-    } as unknown as SDKMessage);
+      await session.close();
-    session.translateMessageToEvents({
+    }
-      type: "stream_event",
-      event: {
-        type: "message_delta",
-        usage: {
-          output_tokens: 25,
-        },
-      },
-      session_id: "session-1",
-    } as unknown as SDKMessage);
-
-    const usage = session.convertUsage({
-      type: "result",
-      subtype: "success",
-      usage: {
-        input_tokens: 10,
-        cache_read_input_tokens: 5,
-        output_tokens: 7,
-      },
-      total_cost_usd: 0.12,
-    });
-
-    expect(usage).toEqual({
-      inputTokens: 10,
-      cachedInputTokens: 5,
-      outputTokens: 7,
-      totalCostUsd: 0.12,
-      contextWindowUsedTokens: 175,
-    });
   });
 
   test("per-request stream usage is not cumulative across API calls in a turn", async () => {
-    const session = await createSessionForTest();
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageStartEvent(),
+        createMessageDeltaEvent(25),
+        createMessageStartEvent({
+          input_tokens: 40,
+          cache_creation_input_tokens: 5,
+          cache_read_input_tokens: 10,
+        }),
+        createMessageDeltaEvent(7),
+        createSuccessResult(),
+      ],
+    ]);
 
-    session.translateMessageToEvents({
+    try {
-      type: "stream_event",
+      const events = await collectStreamEvents(session);
-      event: {
+
-        type: "message_start",
+      expect(events).toContainEqual(
-        message: {
+        expect.objectContaining({
+          type: "usage_updated",
+          provider: "claude",
           usage: {
-            input_tokens: 100,
+            contextWindowUsedTokens: 55,
-            cache_creation_input_tokens: 20,
-            cache_read_input_tokens: 30,
           },
-        },
+        }),
-      },
+      );
-      session_id: "session-1",
+      expect(events).toContainEqual(
-    } as unknown as SDKMessage);
+        expect.objectContaining({
-    session.translateMessageToEvents({
+          type: "usage_updated",
-      type: "stream_event",
+          provider: "claude",
-      event: {
-        type: "message_delta",
-        usage: {
-          output_tokens: 25,
-        },
-      },
-      session_id: "session-1",
-    } as unknown as SDKMessage);
-
-    const secondStartEvents = session.translateMessageToEvents({
-      type: "stream_event",
-      event: {
-        type: "message_start",
-        message: {
           usage: {
-            input_tokens: 40,
+            contextWindowUsedTokens: 62,
-            cache_creation_input_tokens: 5,
-            cache_read_input_tokens: 10,
           },
-        },
+        }),
-      },
+      );
-      session_id: "session-1",
+    } finally {
-    } as unknown as SDKMessage);
+      await session.close();
-
+    }
-    expect(secondStartEvents).toContainEqual({
-      type: "usage_updated",
-      provider: "claude",
-      usage: {
-        contextWindowUsedTokens: 55,
-      },
-    });
-
-    session.translateMessageToEvents({
-      type: "stream_event",
-      event: {
-        type: "message_delta",
-        usage: {
-          output_tokens: 7,
-        },
-      },
-      session_id: "session-1",
-    } as unknown as SDKMessage);
-
-    const usage = session.convertUsage({
-      type: "result",
-      subtype: "success",
-      usage: {
-        input_tokens: 10,
-        cache_read_input_tokens: 5,
-        output_tokens: 7,
-      },
-      total_cost_usd: 0.12,
-    });
-
-    expect(usage).toEqual({
-      inputTokens: 10,
-      cachedInputTokens: 5,
-      outputTokens: 7,
-      totalCostUsd: 0.12,
-      contextWindowUsedTokens: 62,
-    });
   });
 
   test("result.result is surfaced as an assistant message when no model output was produced", async () => {

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

@@ -11,9 +11,9 @@ import {
   type PermissionResult,
   type PermissionUpdate,
   type Query,
+  type SDKControlGetContextUsageResponse,
   type SDKMessage,
   type SDKPartialAssistantMessage,
-  type SDKTaskProgressMessage,
   type SDKResultMessage,
   type SDKSystemMessage,
   type SDKUserMessage,
@@ -1603,27 +1603,6 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
   return maxContextWindow;
 }
 
-function readUsageTotalTokens(usage: unknown): number | undefined {
-  if (!usage || typeof usage !== "object") {
-    return undefined;
-  }
-  const totalTokens = (usage as { total_tokens?: unknown }).total_tokens;
-  if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) {
-    return undefined;
-  }
-  return totalTokens;
-}
-
-function readContextWindowUsedTokensFromTaskProgress(
-  message: SDKTaskProgressMessage,
-): number | undefined {
-  return readUsageTotalTokens(message.usage);
-}
-
-function readUsageFromTaskNotification(message: { usage?: unknown }): number | undefined {
-  return readUsageTotalTokens(message.usage);
-}
-
 function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
   const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage);
   if (!messageUsage) {
@@ -1658,6 +1637,203 @@ function readStreamRequestOutputTokens(event: Record<string, unknown>): number |
   return outputTokens;
 }
 
+function readLastUsageIteration(usage: unknown): Record<string, unknown> | undefined {
+  const iterations = toObjectRecord(usage)?.iterations;
+  if (!Array.isArray(iterations)) {
+    return undefined;
+  }
+  for (let index = iterations.length - 1; index >= 0; index -= 1) {
+    const candidate = toObjectRecord(iterations[index]);
+    if (candidate) {
+      return candidate;
+    }
+  }
+  return undefined;
+}
+
+function readUsageTokenTotal(usage: Record<string, unknown>): number | undefined {
+  const usageWithCacheCreation = usage as typeof usage & {
+    cache_creation_input_tokens?: unknown;
+  };
+  const inputTokens =
+    typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens)
+      ? usage.input_tokens
+      : 0;
+  const cacheCreationInputTokens =
+    typeof usageWithCacheCreation.cache_creation_input_tokens === "number" &&
+    Number.isFinite(usageWithCacheCreation.cache_creation_input_tokens)
+      ? usageWithCacheCreation.cache_creation_input_tokens
+      : 0;
+  const cacheReadInputTokens =
+    typeof usage.cache_read_input_tokens === "number" &&
+    Number.isFinite(usage.cache_read_input_tokens)
+      ? usage.cache_read_input_tokens
+      : 0;
+  const outputTokens =
+    typeof usage.output_tokens === "number" && Number.isFinite(usage.output_tokens)
+      ? usage.output_tokens
+      : 0;
+  const total = inputTokens + cacheCreationInputTokens + cacheReadInputTokens + outputTokens;
+  return total > 0 ? total : undefined;
+}
+
+function readActiveUsageTokens(usage: unknown): number | undefined {
+  const activeUsage = readLastUsageIteration(usage);
+  return activeUsage ? readUsageTokenTotal(activeUsage) : undefined;
+}
+
+function readLegacyResultUsageTokens(usage: unknown): number | undefined {
+  const usageRecord = toObjectRecord(usage);
+  return usageRecord ? readUsageTokenTotal(usageRecord) : undefined;
+}
+
+interface ClaudeCurrentContextUsage {
+  totalTokens: number;
+  maxTokens?: number;
+}
+
+function readCurrentContextUsage(
+  value: SDKControlGetContextUsageResponse | unknown,
+): ClaudeCurrentContextUsage | undefined {
+  const record = toObjectRecord(value);
+  if (!record) {
+    return undefined;
+  }
+  const totalTokens = record.totalTokens;
+  if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) {
+    return undefined;
+  }
+  const maxTokens = record.maxTokens;
+  return {
+    totalTokens,
+    ...(typeof maxTokens === "number" && Number.isFinite(maxTokens) && maxTokens > 0
+      ? { maxTokens }
+      : {}),
+  };
+}
+
+function isClaudeSubagentToolName(name: string | undefined): boolean {
+  return name === "Task" || name === "Agent";
+}
+
+class ClaudeContextUsageState {
+  private contextWindowMaxTokens: number | undefined;
+  private streamRequestInputTokens: number | undefined;
+  private streamRequestOutputTokens: number | undefined;
+  private completedResultTurns = 0;
+
+  beginTurn(): void {
+    this.streamRequestInputTokens = undefined;
+    this.streamRequestOutputTokens = undefined;
+  }
+
+  recordModelUsage(modelUsage: unknown): number | undefined {
+    const contextWindowMaxTokens = extractContextWindowSize(modelUsage);
+    if (contextWindowMaxTokens !== undefined) {
+      this.contextWindowMaxTokens = contextWindowMaxTokens;
+    }
+    return this.contextWindowMaxTokens;
+  }
+
+  recordCurrentContextUsage(usage: ClaudeCurrentContextUsage | undefined): void {
+    if (usage?.maxTokens !== undefined) {
+      this.contextWindowMaxTokens = usage.maxTokens;
+    }
+  }
+
+  buildStreamUsageEvent(event: unknown): AgentStreamEvent | null {
+    const streamEvent = toObjectRecord(event);
+    if (!streamEvent) {
+      return null;
+    }
+    const eventType = readTrimmedString(streamEvent.type);
+    if (eventType === "message_start") {
+      const inputTokens = readStreamRequestInputTokens(streamEvent);
+      if (typeof inputTokens !== "number") {
+        return null;
+      }
+      this.streamRequestInputTokens = inputTokens;
+      this.streamRequestOutputTokens = 0;
+    } else if (eventType === "message_delta") {
+      const outputTokens = readStreamRequestOutputTokens(streamEvent);
+      if (typeof outputTokens !== "number") {
+        return null;
+      }
+      this.streamRequestOutputTokens = outputTokens;
+    } else {
+      return null;
+    }
+
+    const usedTokens = this.streamUsedTokens();
+    if (usedTokens === undefined) {
+      return null;
+    }
+    return this.createUsageUpdatedEvent(usedTokens);
+  }
+
+  buildResultUsage(
+    message: SDKResultMessage,
+    modelUsage: unknown,
+    currentContextUsage: ClaudeCurrentContextUsage | undefined,
+  ): AgentUsage | undefined {
+    try {
+      if (!message.usage) {
+        return undefined;
+      }
+      const usage: AgentUsage = {
+        inputTokens: message.usage.input_tokens,
+        cachedInputTokens: message.usage.cache_read_input_tokens,
+        outputTokens: message.usage.output_tokens,
+        totalCostUsd: message.total_cost_usd,
+      };
+
+      const modelContextWindowMaxTokens = this.recordModelUsage(modelUsage ?? message.modelUsage);
+      this.recordCurrentContextUsage(currentContextUsage);
+      if (this.contextWindowMaxTokens !== undefined) {
+        usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
+      } else if (modelContextWindowMaxTokens !== undefined) {
+        usage.contextWindowMaxTokens = modelContextWindowMaxTokens;
+      }
+
+      const activeResultUsageTokens =
+        readActiveUsageTokens(message.usage) ??
+        (this.completedResultTurns === 0 ? readLegacyResultUsageTokens(message.usage) : undefined);
+      const usedTokens =
+        currentContextUsage?.totalTokens ?? this.streamUsedTokens() ?? activeResultUsageTokens;
+      if (usedTokens !== undefined) {
+        usage.contextWindowUsedTokens = usedTokens;
+      }
+      return usage;
+    } finally {
+      this.completedResultTurns += 1;
+    }
+  }
+
+  private streamUsedTokens(): number | undefined {
+    if (
+      typeof this.streamRequestInputTokens !== "number" ||
+      typeof this.streamRequestOutputTokens !== "number"
+    ) {
+      return undefined;
+    }
+    return this.streamRequestInputTokens + this.streamRequestOutputTokens;
+  }
+
+  private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
+    const usage: AgentUsage = {
+      contextWindowUsedTokens,
+    };
+    if (this.contextWindowMaxTokens !== undefined) {
+      usage.contextWindowMaxTokens = this.contextWindowMaxTokens;
+    }
+    return {
+      type: "usage_updated",
+      provider: "claude",
+      usage,
+    };
+  }
+}
+
 class ClaudeAgentSession implements AgentSession {
   readonly provider = "claude" as const;
   readonly capabilities = CLAUDE_CAPABILITIES;
@@ -1703,10 +1879,7 @@ class ClaudeAgentSession implements AgentSession {
   private pendingInterruptAbort = false;
   private foregroundHasVisibleActivity = false;
   private activeTurnHasAssistantText = false;
-  private lastContextWindowUsedTokens: number | undefined;
+  private readonly contextUsage = new ClaudeContextUsageState();
-  private lastContextWindowMaxTokens: number | undefined;
-  private lastStreamRequestInputTokens: number | undefined;
-  private lastStreamRequestOutputTokens: number | undefined;
   private userMessageIds: string[] = [];
   private readonly emittedUserMessageIds = new Set<string>();
   private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = [];
@@ -1840,6 +2013,7 @@ class ClaudeAgentSession implements AgentSession {
     this.activeForegroundTurnId = turnId;
     this.foregroundHasVisibleActivity = false;
     this.activeTurnHasAssistantText = false;
+    this.contextUsage.beginTurn();
     this.transitionTurnState("foreground", "foreground turn started");
     this.clearRecentStderr();
 
@@ -3005,6 +3179,7 @@ class ClaudeAgentSession implements AgentSession {
       id: this.createTurnId("autonomous"),
     };
     this.activeTurnHasAssistantText = false;
+    this.contextUsage.beginTurn();
     this.notifySubscribers({ type: "turn_started", provider: "claude" });
     this.syncTurnState("autonomous turn started");
   }
@@ -3098,7 +3273,7 @@ class ClaudeAgentSession implements AgentSession {
       if (await this.handleMissingResumedConversation(message, activeQuery)) {
         return true;
       }
-      this.routeSdkMessageFromPump(message);
+      await this.routeSdkMessageFromPump(message, activeQuery);
       return false;
     };
     const drainActiveQuery = async (): Promise<boolean> => {
@@ -3174,7 +3349,7 @@ class ClaudeAgentSession implements AgentSession {
     );
   }
 
-  private routeSdkMessageFromPump(message: SDKMessage): void {
+  private async routeSdkMessageFromPump(message: SDKMessage, activeQuery: Query): Promise<void> {
     if (this.shouldSuppressStaleResult(message)) {
       return;
     }
@@ -3204,26 +3379,12 @@ class ClaudeAgentSession implements AgentSession {
       "provider.claude.parsed_event",
     );
 
-    const messageEvents = this.translateMessageToEvents(message, {
+    const events = await this.buildPumpedMessageEvents(
-      suppressAssistantText: true,
+      message,
-      suppressReasoning: true,
+      activeQuery,
-    });
+      identifiers.messageId,
-    const assistantTimelineEvents = this.timelineAssembler
+      turnId,
-      .consume({
+    );
-        message,
-        runId: turnId,
-        messageIdHint: identifiers.messageId,
-      })
-      .map(
-        (item) =>
-          ({
-            type: "timeline",
-            item,
-            provider: "claude",
-          }) satisfies AgentStreamEvent,
-      );
-
-    const events = [...messageEvents, ...assistantTimelineEvents];
 
     if (events.length === 0) {
       return;
@@ -3258,6 +3419,57 @@ class ClaudeAgentSession implements AgentSession {
     this.dispatchEvents(events);
   }
 
+  private async buildPumpedMessageEvents(
+    message: SDKMessage,
+    activeQuery: Query,
+    messageIdHint: string | null,
+    turnId: string | null,
+  ): Promise<AgentStreamEvent[]> {
+    const currentContextUsage =
+      message.type === "result" && message.subtype === "success"
+        ? await this.queryCurrentContextUsage(activeQuery)
+        : undefined;
+    const messageEvents = this.translateMessageToEvents(message, {
+      suppressAssistantText: true,
+      suppressReasoning: true,
+      currentContextUsage,
+    });
+    const assistantTimelineEvents = this.timelineAssembler
+      .consume({
+        message,
+        runId: turnId,
+        messageIdHint,
+      })
+      .map(
+        (item) =>
+          ({
+            type: "timeline",
+            item,
+            provider: "claude",
+          }) satisfies AgentStreamEvent,
+      );
+
+    return [...messageEvents, ...assistantTimelineEvents];
+  }
+
+  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");
+      return readCurrentContextUsage(usage);
+    } catch (error) {
+      this.logger.debug({ err: error }, "Claude context usage query failed");
+      return undefined;
+    }
+  }
+
   private async handleMissingResumedConversation(
     message: SDKMessage,
     activeQuery: Query,
@@ -3325,6 +3537,7 @@ class ClaudeAgentSession implements AgentSession {
     options?: {
       suppressAssistantText?: boolean;
       suppressReasoning?: boolean;
+      currentContextUsage?: ClaudeCurrentContextUsage;
     },
   ): AgentStreamEvent[] {
     const parentToolUseId =
@@ -3375,7 +3588,9 @@ class ClaudeAgentSession implements AgentSession {
         this.appendStreamEventEvents(message, events, options);
         break;
       case "result":
-        this.appendResultEvents(message, events);
+        this.appendResultEvents(message, events, {
+          currentContextUsage: options?.currentContextUsage,
+        });
         break;
       default:
         break;
@@ -3456,11 +3671,7 @@ class ClaudeAgentSession implements AgentSession {
       return;
     }
     if (message.subtype === "task_progress") {
-      this.lastContextWindowUsedTokens =
+      return;
-        readContextWindowUsedTokensFromTaskProgress(message) ?? this.lastContextWindowUsedTokens;
-      if (typeof this.lastContextWindowUsedTokens === "number") {
-        events.push(this.createUsageUpdatedEvent(this.lastContextWindowUsedTokens));
-      }
     }
   }
 
@@ -3470,12 +3681,12 @@ class ClaudeAgentSession implements AgentSession {
   ): void {
     // TODO: subagent timelines are best-effort. Subagent task_notifications
     // arrive without parent_tool_use_id but with tool_use_id pointing at the
-    // parent's Task call, so they slip past the sidechain router and pollute
+    // the parent's subagent tool call, so they slip past the sidechain router and pollute
     // the parent timeline. Drop them here; eventually thread them into the
-    // parent Task tool call's sub_agent log instead.
+    // parent tool call's sub_agent log instead.
     const taskUseId = message.tool_use_id;
     const cachedTool = taskUseId ? this.toolUseCache.get(taskUseId) : undefined;
-    if (cachedTool?.name === "Task") {
+    if (isClaudeSubagentToolName(cachedTool?.name)) {
       return;
     }
     const taskNotificationItem = mapTaskNotificationSystemRecordToToolCall(message);
@@ -3486,11 +3697,6 @@ class ClaudeAgentSession implements AgentSession {
         provider: "claude",
       });
     }
-    const usage = readUsageFromTaskNotification(message);
-    if (typeof usage === "number") {
-      this.lastContextWindowUsedTokens = usage;
-      events.push(this.createUsageUpdatedEvent(usage));
-    }
   }
 
   private appendUserMessageEvents(
@@ -3569,7 +3775,7 @@ class ClaudeAgentSession implements AgentSession {
     events: AgentStreamEvent[],
     options: { suppressAssistantText?: boolean; suppressReasoning?: boolean } | undefined,
   ): void {
-    const usageUpdatedEvent = this.trackStreamEventUsage(message.event);
+    const usageUpdatedEvent = this.contextUsage.buildStreamUsageEvent(message.event);
     if (usageUpdatedEvent) {
       events.push(usageUpdatedEvent);
     }
@@ -3585,8 +3791,9 @@ class ClaudeAgentSession implements AgentSession {
   private appendResultEvents(
     message: Extract<SDKMessage, { type: "result" }>,
     events: AgentStreamEvent[],
+    options?: { currentContextUsage?: ClaudeCurrentContextUsage },
   ): void {
-    const usage = this.convertUsage(message, message.modelUsage);
+    const usage = this.convertUsage(message, message.modelUsage, options?.currentContextUsage);
     if (message.subtype === "success") {
       // Built-in slash commands (e.g. /voice, /usage, "Unknown command: …")
       // run client-side in the Claude CLI with no model turn — output_tokens
@@ -3753,99 +3960,12 @@ class ClaudeAgentSession implements AgentSession {
     return null;
   }
 
-  private convertUsage(message: SDKResultMessage, modelUsage?: unknown): AgentUsage | undefined {
+  private convertUsage(
-    if (!message.usage) {
+    message: SDKResultMessage,
-      return undefined;
+    modelUsage?: unknown,
-    }
+    currentContextUsage?: ClaudeCurrentContextUsage,
-    const usage: AgentUsage = {
+  ): AgentUsage | undefined {
-      inputTokens: message.usage.input_tokens,
+    return this.contextUsage.buildResultUsage(message, modelUsage, currentContextUsage);
-      cachedInputTokens: message.usage.cache_read_input_tokens,
-      outputTokens: message.usage.output_tokens,
-      totalCostUsd: message.total_cost_usd,
-    };
-    const contextWindowMaxTokens = extractContextWindowSize(modelUsage ?? message.modelUsage);
-    if (contextWindowMaxTokens !== undefined) {
-      this.lastContextWindowMaxTokens = contextWindowMaxTokens;
-      usage.contextWindowMaxTokens = contextWindowMaxTokens;
-    } else if (this.lastContextWindowMaxTokens !== undefined) {
-      usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
-    }
-    if (typeof this.lastContextWindowUsedTokens === "number") {
-      // task_progress.total_tokens is the accurate context window fill level.
-      // Prefer it over result.usage which contains accumulated session totals.
-      usage.contextWindowUsedTokens = this.lastContextWindowUsedTokens;
-    } else if (
-      typeof this.lastStreamRequestInputTokens === "number" &&
-      typeof this.lastStreamRequestOutputTokens === "number"
-    ) {
-      usage.contextWindowUsedTokens =
-        this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens;
-    } else if (message.usage) {
-      // Fallback: derive from result.usage when no task_progress has been
-      // received yet. These values are accumulated across all API calls, but
-      // for the first turn they equal the per-call values so the estimate is
-      // reasonable. Once a task_progress arrives it takes over permanently.
-      const usageWithCacheCreation = message.usage as typeof message.usage & {
-        cache_creation_input_tokens?: number;
-      };
-      const derived =
-        (message.usage.input_tokens ?? 0) +
-        (usageWithCacheCreation.cache_creation_input_tokens ?? 0) +
-        (message.usage.cache_read_input_tokens ?? 0) +
-        (message.usage.output_tokens ?? 0);
-      if (Number.isFinite(derived) && derived > 0) {
-        usage.contextWindowUsedTokens = derived;
-      }
-    }
-    return usage;
-  }
-
-  private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
-    const usage: AgentUsage = {
-      contextWindowUsedTokens,
-    };
-    if (this.lastContextWindowMaxTokens !== undefined) {
-      usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
-    }
-    return {
-      type: "usage_updated",
-      provider: "claude",
-      usage,
-    };
-  }
-
-  private trackStreamEventUsage(event: unknown): AgentStreamEvent | null {
-    const streamEvent = toObjectRecord(event);
-    if (!streamEvent) {
-      return null;
-    }
-    const eventType = readTrimmedString(streamEvent.type);
-    if (eventType === "message_start") {
-      const inputTokens = readStreamRequestInputTokens(streamEvent);
-      if (typeof inputTokens !== "number") {
-        return null;
-      }
-      this.lastStreamRequestInputTokens = inputTokens;
-      this.lastStreamRequestOutputTokens = 0;
-    } else if (eventType === "message_delta") {
-      const outputTokens = readStreamRequestOutputTokens(streamEvent);
-      if (typeof outputTokens !== "number") {
-        return null;
-      }
-      this.lastStreamRequestOutputTokens = outputTokens;
-    } else {
-      return null;
-    }
-
-    if (
-      typeof this.lastStreamRequestInputTokens !== "number" ||
-      typeof this.lastStreamRequestOutputTokens !== "number"
-    ) {
-      return null;
-    }
-    return this.createUsageUpdatedEvent(
-      this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens,
-    );
   }
 
   private handlePermissionRequest: CanUseTool = async (