From 5214612a599547b79170c5b245833400cd40af8c Mon Sep 17 00:00:00 2001
From: Alex Castillo <alex.castillog33@gmail.com>
Date: Wed, 21 Oct 2026 08:30:07 -0400
Subject: [PATCH] fix(ai): use cumulative input and cache tokens from anthropic
 message_delta when streaming

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 .../anthropic-stream-cumulative-usage.md      |  5 ++
 packages/ai/src/anthropic/index.ts            | 11 +++
 packages/ai/tests/anthropic.test.ts           | 71 +++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 .changeset/anthropic-stream-cumulative-usage.md

diff --git a/.changeset/anthropic-stream-cumulative-usage.md b/.changeset/anthropic-stream-cumulative-usage.md
new file mode 100644
index 0000000000..5d22a3be0a
--- /dev/null
+++ b/.changeset/anthropic-stream-cumulative-usage.md
@@ -0,0 +1,5 @@
+---
+'@posthog/ai': patch
+---
+
+fix: Use the cumulative input and cache token counts from Anthropic's final `message_delta` when streaming with server tools.
diff --git a/packages/ai/src/anthropic/index.ts b/packages/ai/src/anthropic/index.ts
index 65860aa2f9..23cf45ec03 100644
--- a/packages/ai/src/anthropic/index.ts
+++ b/packages/ai/src/anthropic/index.ts
@@ -222,6 +222,17 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
                     ...Object.fromEntries(Object.entries(chunk.usage).filter(([, value]) => value != null)),
                   }
                   usage.outputTokens = chunk.usage.output_tokens ?? 0
+                  // With server tools the final delta reports cumulative input and cache tokens that
+                  // exceed message_start's, so keep the larger value.
+                  usage.inputTokens = Math.max(usage.inputTokens ?? 0, chunk.usage.input_tokens ?? 0)
+                  usage.cacheCreationInputTokens = Math.max(
+                    usage.cacheCreationInputTokens ?? 0,
+                    chunk.usage.cache_creation_input_tokens ?? 0
+                  )
+                  usage.cacheReadInputTokens = Math.max(
+                    usage.cacheReadInputTokens ?? 0,
+                    chunk.usage.cache_read_input_tokens ?? 0
+                  )
                   // Update web search count if present in delta
                   if (chunk.usage.server_tool_use?.web_search_requests !== undefined) {
                     usage.webSearchCount = chunk.usage.server_tool_use.web_search_requests
diff --git a/packages/ai/tests/anthropic.test.ts b/packages/ai/tests/anthropic.test.ts
index 0ff398ffa3..1aca9ab453 100644
--- a/packages/ai/tests/anthropic.test.ts
+++ b/packages/ai/tests/anthropic.test.ts
@@ -1463,3 +1463,74 @@ describe('PostHogAnthropic - streaming error safety', () => {
     expect(unhandledRejections).toEqual([])
   })
 })
+
+describe('PostHogAnthropic - streaming usage', () => {
+  let usageMockPostHogClient: PostHog
+  let usageClient: PostHogAnthropic
+
+  beforeEach(() => {
+    vi.clearAllMocks()
+    usageMockPostHogClient = new (PostHog as any)()
+    usageClient = new PostHogAnthropic({
+      apiKey: 'test-api-key',
+      posthog: usageMockPostHogClient as any,
+    })
+  })
+
+  const captureStreamUsage = async (deltaUsage: Record<string, number>): Promise<Record<string, unknown>> => {
+    const events = [
+      {
+        type: 'message_start',
+        message: {
+          usage: { input_tokens: 50, cache_creation_input_tokens: 10, cache_read_input_tokens: 20, output_tokens: 1 },
+        },
+      },
+      { type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: deltaUsage },
+      { type: 'message_stop' },
+    ]
+    const MessagesMock = AnthropicOriginal.Messages as vi.MockedClass<typeof AnthropicOriginal.Messages>
+    ;(MessagesMock.prototype.create as unknown as vi.Mock) = vi.fn().mockResolvedValue({
+      async *[Symbol.asyncIterator]() {
+        yield* events
+      },
+    })
+
+    const stream = await usageClient.messages.create({
+      model: 'claude-3-5-sonnet-20240620',
+      max_tokens: 100,
+      messages: [{ role: 'user', content: 'hi' }],
+      stream: true,
+      posthogDistinctId: 'anthropic-usage-user',
+    } as any)
+    for await (const _chunk of stream as unknown as AsyncIterable<unknown>) {
+      // consume
+    }
+    await new Promise(process.nextTick)
+
+    expect(usageMockPostHogClient.capture).toHaveBeenCalledTimes(1)
+    return (usageMockPostHogClient.capture as vi.Mock).mock.calls[0][0].properties
+  }
+
+  test('uses the cumulative input and cache tokens from message_delta when they grow (server tools)', async () => {
+    const properties = await captureStreamUsage({
+      input_tokens: 120,
+      cache_creation_input_tokens: 30,
+      cache_read_input_tokens: 70,
+      output_tokens: 25,
+    })
+
+    expect(properties['$ai_input_tokens']).toBe(120)
+    expect(properties['$ai_cache_creation_input_tokens']).toBe(30)
+    expect(properties['$ai_cache_read_input_tokens']).toBe(70)
+    expect(properties['$ai_output_tokens']).toBe(25)
+  })
+
+  test('keeps the message_start input and cache tokens when message_delta omits them or reports less', async () => {
+    const properties = await captureStreamUsage({ input_tokens: 5, output_tokens: 25 })
+
+    expect(properties['$ai_input_tokens']).toBe(50)
+    expect(properties['$ai_cache_creation_input_tokens']).toBe(10)
+    expect(properties['$ai_cache_read_input_tokens']).toBe(20)
+    expect(properties['$ai_output_tokens']).toBe(25)
+  })
+})