fix(claude): count input reported late in the context meter (#6703)

* fix(claude): read cumulative input usage from message_delta The context meter took input_tokens only from message_start. Backends behind a translating gateway report 0 there and the real cumulative counts on message_delta, so contextWindowUsedTokens collapsed to output_tokens. Keep the message_start counters per field and let a counter present on message_delta replace it, the same rule the Anthropic SDK MessageStream uses. A missing counter keeps the message_start value, so a delta without cache fields does not drop cache_read_input_tokens. Closes #4294 * test(claude): ignore a message_delta that precedes its message_start A canceled request can leave its terminal message_delta queued ahead of the next turn. Pin that such a straggler neither emits usage_updated nor leaks into the next request's counters. * test(claude): take the message_delta usage object as a single parameter Keep createMessageDeltaEvent(outputTokens) for existing call sites and add createMessageDeltaUsageEvent(usage) for the tests that set input counters on the delta, instead of a second positional parameter. * test(claude): verify late input usage through the real provider process Compare start-side and delta-side counters through two real provider turns against a local Messages API fixture, including omitted cache counters. The delta cases fail on the baseline with output-only context totals. Co-authored-by: kywoo26 <kywoo26@naver.com> * test(claude): always close the context usage fixture --------- Co-authored-by: kywoo26 <kywoo26@naver.com> Co-authored-by: paseo-bot[bot] <266920839+paseo-bot[bot]@users.noreply.github.com>
2026-10-14 13:22:37 +00:00 · 2026-10-12 12:59:46 +02:00
parent 99e4b53aa0
commit a2b040869d
3 changed files with 319 additions and 34 deletions
@@ -0,0 +1,133 @@
+import { createServer } from "node:http";
+import { mkdir, mkdtemp, rm } from "node:fs/promises";
+import path from "node:path";
+import pino from "pino";
+import { expect, test } from "vitest";
+import { collectSessionTurnEvents } from "../test-utils/session-stream-adapter.js";
+import { ClaudeAgentClient } from "./agent.js";
+
+// Exercise the real Claude Code process against a local Messages API fixture.
+// The input counter's location is the only difference between the two backends.
+test.each([
+  { placement: "start", cacheTokens: 0 },
+  { placement: "delta", cacheTokens: 0 },
+  { placement: "start", cacheTokens: 1_000 },
+  { placement: "delta", cacheTokens: 1_000 },
+])(
+  "counts $placement input and $cacheTokens cached tokens through Claude Code",
+  async ({ placement, cacheTokens }) => {
+    const cwd = await mkdtemp(path.resolve(".claude-context-usage-"));
+    const configDir = path.join(cwd, "config");
+    await mkdir(configDir);
+    let requestCount = 0;
+    const server = createServer(async (request, response) => {
+      for await (const chunk of request) {
+        void chunk;
+      }
+      if (!request.url?.startsWith("/v1/messages")) {
+        response.writeHead(200, { "content-type": "application/json" });
+        response.end("{}");
+        return;
+      }
+      if (request.url.includes("count_tokens")) {
+        response.writeHead(200, { "content-type": "application/json" });
+        response.end(JSON.stringify({ input_tokens: 20_236 }));
+        return;
+      }
+      const inputTokens = 20_236 + requestCount++;
+      response.writeHead(200, { "content-type": "text/event-stream" });
+      const emit = (data: { type: string; [key: string]: unknown }) => {
+        response.write(`event: ${data.type}\ndata: ${JSON.stringify(data)}\n\n`);
+      };
+      emit({
+        type: "message_start",
+        message: {
+          id: `msg_context_usage_${requestCount}`,
+          type: "message",
+          role: "assistant",
+          model: "claude-sonnet-4-6",
+          content: [],
+          stop_reason: null,
+          stop_sequence: null,
+          usage: {
+            input_tokens: placement === "start" ? inputTokens : 0,
+            output_tokens: 0,
+            cache_creation_input_tokens: 0,
+            cache_read_input_tokens: cacheTokens,
+          },
+        },
+      });
+      emit({ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } });
+      emit({ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hello" } });
+      emit({ type: "content_block_stop", index: 0 });
+      emit({
+        type: "message_delta",
+        delta: { stop_reason: "end_turn", stop_sequence: null },
+        usage: {
+          output_tokens: 64,
+          ...(placement === "delta" ? { input_tokens: inputTokens } : {}),
+        },
+      });
+      emit({ type: "message_stop" });
+      response.end();
+    });
+    await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+    const address = server.address();
+    if (!address || typeof address === "string") throw new Error("Fixture has no TCP address");
+    const client = new ClaudeAgentClient({
+      logger: pino({ level: "silent" }),
+      runtimeSettings: {
+        env: {
+          ANTHROPIC_BASE_URL: `http://127.0.0.1:${address.port}`,
+          ANTHROPIC_AUTH_TOKEN: "fixture-token",
+          ANTHROPIC_API_KEY: "",
+          CLAUDE_CONFIG_DIR: configDir,
+          CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC: "1",
+        },
+      },
+    });
+    const session = await client.createSession(
+      { provider: "claude", cwd, model: "claude-sonnet-4-6" },
+      undefined,
+      { persistSession: false },
+    );
+    try {
+      for (let turn = 0; turn < 2; turn++) {
+        const events = await collectSessionTurnEvents(session, "say hello", {
+          signal: AbortSignal.timeout(20_000),
+        });
+        const expectedTokens = 20_236 + turn + cacheTokens + 64;
+        expect(events).toContainEqual(
+          expect.objectContaining({
+            type: "turn_completed",
+            usage: expect.objectContaining({
+              inputTokens: 20_236 + turn,
+              outputTokens: 64,
+              cachedInputTokens: cacheTokens,
+              contextWindowUsedTokens: expectedTokens,
+            }),
+          }),
+        );
+        expect(events).toContainEqual(
+          expect.objectContaining({
+            type: "usage_updated",
+            usage: expect.objectContaining({ contextWindowUsedTokens: expectedTokens }),
+          }),
+        );
+      }
+      expect(requestCount).toBe(2);
+    } finally {
+      try {
+        await session.close();
+      } finally {
+        try {
+          server.closeAllConnections();
+          await new Promise<void>((resolve) => server.close(() => resolve()));
+        } finally {
+          await rm(cwd, { recursive: true, force: true });
+        }
+      }
+    }
+  },
+  60_000,
+);
@@ -2135,13 +2135,15 @@ describe("ClaudeAgentSession context window usage", () => {
  }

  function createMessageDeltaEvent(outputTokens: number): Record<string, unknown> {
+    return createMessageDeltaUsageEvent({ output_tokens: outputTokens });
+  }
+
+  function createMessageDeltaUsageEvent(usage: Record<string, unknown>): Record<string, unknown> {
    return {
      type: "stream_event",
      event: {
        type: "message_delta",
-        usage: {
-          output_tokens: outputTokens,
-        },
+        usage,
      },
      session_id: "session-1",
    };
@@ -2887,6 +2889,119 @@ describe("ClaudeAgentSession context window usage", () => {
    }
  });

+  test("message_delta input usage replaces an empty message_start count", async () => {
+    // Translating gateways (LiteLLM in front of an OpenAI-compatible backend) do not know the
+    // prompt size when message_start is emitted, so they report input_tokens: 0 there and send
+    // the cumulative counts on message_delta, as the Anthropic Messages API allows.
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageStartEvent({
+          input_tokens: 0,
+          cache_creation_input_tokens: 0,
+          cache_read_input_tokens: 0,
+        }),
+        createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }),
+        createSuccessResult({
+          usage: {
+            input_tokens: 20_236,
+            cache_creation_input_tokens: 0,
+            cache_read_input_tokens: 0,
+            output_tokens: 64,
+            iterations: [],
+          },
+        }),
+      ],
+    ]);
+
+    try {
+      const events = await collectStreamEvents(session);
+
+      expect(events).toContainEqual(
+        expect.objectContaining({
+          type: "usage_updated",
+          provider: "claude",
+          usage: {
+            contextWindowUsedTokens: 20_300,
+          },
+        }),
+      );
+      expect(events).toContainEqual(
+        expect.objectContaining({
+          type: "turn_completed",
+          usage: expect.objectContaining({
+            inputTokens: 20_236,
+            outputTokens: 64,
+            contextWindowUsedTokens: 20_300,
+          }),
+        }),
+      );
+    } finally {
+      await session.close();
+    }
+  });
+
+  test("message_delta without cache counters keeps the message_start cache usage", async () => {
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageStartEvent({
+          input_tokens: 5,
+          cache_creation_input_tokens: 0,
+          cache_read_input_tokens: 1_000,
+        }),
+        createMessageDeltaUsageEvent({ input_tokens: 5, output_tokens: 64 }),
+        createSuccessResult(),
+      ],
+    ]);
+
+    try {
+      const events = await collectStreamEvents(session);
+
+      expect(events).toContainEqual(
+        expect.objectContaining({
+          type: "usage_updated",
+          provider: "claude",
+          usage: {
+            contextWindowUsedTokens: 1_069,
+          },
+        }),
+      );
+    } finally {
+      await session.close();
+    }
+  });
+
+  test("message_delta before this request's message_start does not seed usage", async () => {
+    // A canceled request can leave its terminal message_delta queued behind the next turn.
+    // Nothing is known about the new request until its message_start, so that straggler
+    // must not produce a usage_updated event or leak into the next request's counters.
+    const session = await createSessionForTurns([
+      [
+        createInitMessage(),
+        createMessageDeltaUsageEvent({ input_tokens: 20_236, output_tokens: 64 }),
+        createMessageStartEvent({
+          input_tokens: 40,
+          cache_creation_input_tokens: 5,
+          cache_read_input_tokens: 10,
+        }),
+        createMessageDeltaEvent(7),
+        createSuccessResult(),
+      ],
+    ]);
+
+    try {
+      const events = await collectStreamEvents(session);
+
+      const usedTokens = events
+        .filter((event) => event.type === "usage_updated")
+        .map((event) => event.usage.contextWindowUsedTokens);
+      expect(usedTokens).toEqual([55, 62]);
+    } finally {
+      await session.close();
+    }
+  });
+
  test("per-request stream usage is not cumulative across API calls in a turn", async () => {
    const session = await createSessionForTurns([
      [
@@ -1855,30 +1855,59 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
  return maxContextWindow;
 }

-function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
-  const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage);
-  if (!messageUsage) {
+interface StreamRequestInputUsage {
+  inputTokens: number;
+  cacheCreationInputTokens: number;
+  cacheReadInputTokens: number;
+}
+
+function readUsageCount(usage: Record<string, unknown>, key: string): number | undefined {
+  const value = usage[key];
+  if (typeof value !== "number" || !Number.isFinite(value) || value < 0) {
    return undefined;
  }
-  const usage = messageUsage;
-  const inputTokens =
-    typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens)
-      ? usage.input_tokens
-      : undefined;
-  const cacheCreationInputTokens =
-    typeof usage.cache_creation_input_tokens === "number" &&
-    Number.isFinite(usage.cache_creation_input_tokens)
-      ? usage.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;
-  if (typeof inputTokens !== "number" || inputTokens < 0) {
+  return value;
+}
+
+function readStreamRequestInputUsage(
+  event: Record<string, unknown>,
+): StreamRequestInputUsage | undefined {
+  const usage = toObjectRecord(toObjectRecord(event.message)?.usage);
+  if (!usage) {
    return undefined;
  }
-  return inputTokens + cacheCreationInputTokens + cacheReadInputTokens;
+  const inputTokens = readUsageCount(usage, "input_tokens");
+  if (inputTokens === undefined) {
+    return undefined;
+  }
+  return {
+    inputTokens,
+    cacheCreationInputTokens: readUsageCount(usage, "cache_creation_input_tokens") ?? 0,
+    cacheReadInputTokens: readUsageCount(usage, "cache_read_input_tokens") ?? 0,
+  };
+}
+
+/**
+ * The input counters on `message_delta.usage` are cumulative whole-request totals and are
+ * optional. Mirror the Anthropic SDK's MessageStream: a counter present on the delta replaces
+ * the `message_start` value, a missing counter keeps it. Translating gateways report
+ * `input_tokens: 0` on `message_start` and only know the real count by `message_delta`.
+ */
+function mergeStreamDeltaInputUsage(
+  current: StreamRequestInputUsage,
+  event: Record<string, unknown>,
+): StreamRequestInputUsage {
+  const usage = toObjectRecord(event.usage);
+  if (!usage) {
+    return current;
+  }
+  return {
+    inputTokens: readUsageCount(usage, "input_tokens") ?? current.inputTokens,
+    cacheCreationInputTokens:
+      readUsageCount(usage, "cache_creation_input_tokens") ?? current.cacheCreationInputTokens,
+    cacheReadInputTokens:
+      readUsageCount(usage, "cache_read_input_tokens") ?? current.cacheReadInputTokens,
+  };
 }

 function readStreamRequestOutputTokens(event: Record<string, unknown>): number | undefined {
@@ -1958,7 +1987,7 @@ function readClaudeParentToolUseId(message: SDKMessage): string | null {

 class ClaudeContextUsageState {
  private contextWindowMaxTokens: number | undefined;
-  private streamRequestInputTokens: number | undefined;
+  private streamRequestInputUsage: StreamRequestInputUsage | undefined;
  private streamRequestOutputTokens: number | undefined;
  private compactedContextWindowUsedTokens: number | undefined;
  private completedResultTurns = 0;
@@ -1968,7 +1997,7 @@ class ClaudeContextUsageState {
  }

  beginTurn(): void {
-    this.streamRequestInputTokens = undefined;
+    this.streamRequestInputUsage = undefined;
    this.streamRequestOutputTokens = undefined;
    this.compactedContextWindowUsedTokens = undefined;
  }
@@ -1992,11 +2021,11 @@ class ClaudeContextUsageState {
    }
    const eventType = readTrimmedString(streamEvent.type);
    if (eventType === "message_start") {
-      const inputTokens = readStreamRequestInputTokens(streamEvent);
-      if (typeof inputTokens !== "number") {
+      const inputUsage = readStreamRequestInputUsage(streamEvent);
+      if (!inputUsage) {
        return null;
      }
-      this.streamRequestInputTokens = inputTokens;
+      this.streamRequestInputUsage = inputUsage;
      this.streamRequestOutputTokens = 0;
    } else if (eventType === "message_delta") {
      const outputTokens = readStreamRequestOutputTokens(streamEvent);
@@ -2004,6 +2033,12 @@ class ClaudeContextUsageState {
        return null;
      }
      this.streamRequestOutputTokens = outputTokens;
+      if (this.streamRequestInputUsage) {
+        this.streamRequestInputUsage = mergeStreamDeltaInputUsage(
+          this.streamRequestInputUsage,
+          streamEvent,
+        );
+      }
    } else {
      return null;
    }
@@ -2050,13 +2085,15 @@ class ClaudeContextUsageState {
  }

  private streamUsedTokens(): number | undefined {
-    if (
-      typeof this.streamRequestInputTokens !== "number" ||
-      typeof this.streamRequestOutputTokens !== "number"
-    ) {
+    const inputUsage = this.streamRequestInputUsage;
+    if (!inputUsage || typeof this.streamRequestOutputTokens !== "number") {
      return undefined;
    }
-    const usedTokens = this.streamRequestInputTokens + this.streamRequestOutputTokens;
+    const usedTokens =
+      inputUsage.inputTokens +
+      inputUsage.cacheCreationInputTokens +
+      inputUsage.cacheReadInputTokens +
+      this.streamRequestOutputTokens;
    return usedTokens > 0 ? usedTokens : undefined;
  }

@@ -2075,7 +2112,7 @@ class ClaudeContextUsageState {
  }

  buildCompactionUsageEvent(postTokens: number | undefined): AgentStreamEvent {
-    this.streamRequestInputTokens = undefined;
+    this.streamRequestInputUsage = undefined;
    this.streamRequestOutputTokens = undefined;
    this.compactedContextWindowUsedTokens = postTokens;
    const usage: AgentUsage = {};