proxy: tee upstream body and extract Anthropic usage tokens

ProxyEvent gains two new fields:
- usage: Anthropic's input/output/cache_creation/cache_read tokens.
- firstByteMs: ms to upstream response headers (durationMs now measures
  end-to-end since the event fires after usage is extracted).

Implementation (teeForUsage):
- Tee the response body. Client gets one side immediately; we read the
  other in the background.
- SSE: scan up to 64 KiB for 'event: message_start' + the data: payload.
  Usage is always in the first event so this terminates within ms of
  upstream first byte.
- JSON: buffer up to 4 MiB and parse.
- Error responses (status >= 400) and bodyless responses skip extraction.
- Drains the tee fully in the background to avoid backpressure.

Without this, the proxy had no way to validate its own value
proposition. Now we can compute cache hit rate, real input-token
savings, etc. directly from log events.

Tests: +3 (JSON usage, SSE message_start usage, error skips extraction).

src/core/proxy.ts

@@ -8,6 +8,7 @@
  */
 
 import { transformRequest, type TransformOptions, type TransformInfo } from './transform.js';
+import type { Usage } from './types.js';
 
 export interface ProxyConfig {
   /** Anthropic API base, no trailing slash. Defaults to api.anthropic.com. */
@@ -24,11 +25,111 @@ export interface ProxyEvent {
   method: string;
   path: string;
   status: number;
+  /** Wall-clock ms from request start to event fire (≈ end of upstream response
+   *  body, since we now wait for usage extraction). For first-byte latency see
+   *  firstByteMs. */
   durationMs: number;
+  /** Wall-clock ms from request start to upstream response headers. */
+  firstByteMs?: number;
   info?: TransformInfo;
+  /** Usage block from Anthropic's response — input/output/cache tokens. */
+  usage?: Usage;
   error?: string;
 }
 
+/**
+ * Tee the response body so we can scan for the usage block (SSE: in the
+ * message_start event; non-stream: at the top of the JSON) without buffering
+ * the whole stream or blocking the client. Returns the un-touched response
+ * to forward to the client + a Promise that resolves to the parsed Usage
+ * (or undefined if we couldn't find one within the budget).
+ */
+function teeForUsage(res: Response): {
+  response: Response;
+  usagePromise: Promise<Usage | undefined>;
+} {
+  // Errors and bodyless responses: nothing to extract.
+  if (!res.body || res.status >= 400) {
+    return { response: res, usagePromise: Promise.resolve(undefined) };
+  }
+  const ct = (res.headers.get('content-type') ?? '').toLowerCase();
+  const [forClient, forUs] = res.body.tee();
+
+  const usagePromise = (async (): Promise<Usage | undefined> => {
+    const reader = forUs.getReader();
+    const decoder = new TextDecoder();
+    let buf = '';
+    const drain = async () => {
+      try {
+        while (true) {
+          const { done } = await reader.read();
+          if (done) break;
+        }
+      } catch {
+        /* ignore */
+      }
+    };
+
+    try {
+      if (ct.includes('text/event-stream')) {
+        // SSE: usage is in the FIRST event (`message_start`). Cap scan at 64
+        // KiB so we don't hold the tee buffer open for the entire stream.
+        const MAX = 65536;
+        while (buf.length < MAX) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buf += decoder.decode(value, { stream: true });
+          const idx = buf.indexOf('event: message_start');
+          if (idx >= 0) {
+            // The data: line follows. Match the first data: after that idx.
+            const m = /^data:\s*(.+)$/m.exec(buf.slice(idx));
+            if (m) {
+              try {
+                const j = JSON.parse(m[1]!);
+                void drain();
+                return j?.message?.usage as Usage | undefined;
+              } catch {
+                /* not yet a complete JSON line — keep reading */
+              }
+            }
+          }
+        }
+        void drain();
+        return undefined;
+      }
+
+      if (ct.includes('application/json')) {
+        // Non-stream: buffer fully (capped at 4 MiB).
+        const MAX = 4 * 1024 * 1024;
+        while (buf.length < MAX) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buf += decoder.decode(value, { stream: true });
+        }
+        try {
+          const j = JSON.parse(buf);
+          return j?.usage as Usage | undefined;
+        } catch {
+          return undefined;
+        }
+      }
+    } catch {
+      /* tee may be released early if the client aborts — ignore */
+    }
+    void drain();
+    return undefined;
+  })();
+
+  return {
+    response: new Response(forClient, {
+      status: res.status,
+      statusText: res.statusText,
+      headers: res.headers,
+    }),
+    usagePromise,
+  };
+}
+
 const DEFAULT_UPSTREAM = 'https://api.anthropic.com';
 
 /** Headers we strip on the way out — they're hop-by-hop or proxy-injected. */
@@ -69,13 +170,21 @@ export function createProxy(config: ProxyConfig = {}) {
     const url = new URL(req.url);
     const path = url.pathname + url.search;
 
-    const fire = (status: number, info?: TransformInfo, error?: string): void => {
+    const fire = (
+      status: number,
+      info?: TransformInfo,
+      error?: string,
+      firstByteMs?: number,
+      usage?: Usage,
+    ): void => {
       void config.onRequest?.({
         method: req.method,
         path: url.pathname,
         status,
         durationMs: Date.now() - t0,
+        firstByteMs,
         info,
+        usage,
         error,
       });
     };
@@ -127,9 +236,20 @@ export function createProxy(config: ProxyConfig = {}) {
       });
     }
 
-    fire(upstreamRes.status, info);
+    const firstByteMs = Date.now() - t0;
 
-    return new Response(upstreamRes.body, {
+    // Tee the upstream body so we can extract Anthropic's usage block. The
+    // client gets one side immediately; we read the other in the background.
+    const { response: teed, usagePromise } = teeForUsage(upstreamRes);
+
+    // Fire the host event once usage is known (or once we've given up on
+    // finding it). Don't await — the response below is what unblocks the
+    // client; fire happens in the background.
+    void usagePromise
+      .then((usage) => fire(upstreamRes.status, info, undefined, firstByteMs, usage))
+      .catch(() => fire(upstreamRes.status, info, undefined, firstByteMs, undefined));
+
+    return new Response(teed.body, {
       status: upstreamRes.status,
       statusText: upstreamRes.statusText,
       headers: filterHeaders(upstreamRes.headers, STRIP_RES_HEADERS),

src/core/types.ts

@@ -57,6 +57,15 @@ export interface ToolDef {
 
 export type SystemField = string | Array<TextBlock | ImageBlock>;
 
+/** Anthropic's per-response token usage block. Same shape on streaming
+ *  (inside the message_start event payload) and non-streaming responses. */
+export interface Usage {
+  input_tokens?: number;
+  output_tokens?: number;
+  cache_creation_input_tokens?: number;
+  cache_read_input_tokens?: number;
+}
+
 export interface MessagesRequest {
   model: string;
   messages: Message[];

tests/proxy-usage.test.ts

@@ -0,0 +1,164 @@
+import { describe, it, expect } from 'vitest';
+import { createProxy, type ProxyEvent } from '../src/core/proxy.js';
+
+/** Tiny in-process mock upstream — accepts any request and returns whatever
+ *  the test fixture configured. Lets us assert that the proxy correctly
+ *  extracts Anthropic's usage block from both SSE and JSON responses without
+ *  touching the network. */
+function mockUpstream(handler: (req: Request) => Promise<Response> | Response) {
+  // Patch globalThis.fetch for the duration of the test.
+  const real = globalThis.fetch;
+  globalThis.fetch = ((req: Request | string | URL, init?: RequestInit) => {
+    const r = req instanceof Request ? req : new Request(String(req), init);
+    return Promise.resolve(handler(r));
+  }) as typeof fetch;
+  return () => {
+    globalThis.fetch = real;
+  };
+}
+
+const SAMPLE_REQ_BODY = JSON.stringify({
+  model: 'claude-3-5-haiku-latest',
+  messages: [{ role: 'user', content: 'hi' }],
+  system: 'short',
+});
+
+describe('proxy usage extraction', () => {
+  it('extracts usage tokens from a non-stream JSON response', async () => {
+    const restore = mockUpstream(
+      () =>
+        new Response(
+          JSON.stringify({
+            id: 'msg_1',
+            type: 'message',
+            role: 'assistant',
+            content: [{ type: 'text', text: 'hello' }],
+            usage: {
+              input_tokens: 123,
+              output_tokens: 7,
+              cache_creation_input_tokens: 0,
+              cache_read_input_tokens: 100,
+            },
+          }),
+          { status: 200, headers: { 'content-type': 'application/json' } },
+        ),
+    );
+
+    let captured: ProxyEvent | undefined;
+    const proxy = createProxy({
+      transform: {},
+      onRequest: (e) => {
+        captured = e;
+      },
+    });
+
+    const res = await proxy(
+      new Request('http://localhost/v1/messages', {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        body: SAMPLE_REQ_BODY,
+      }),
+    );
+    // Drain the client-side body so the tee is forced to finish.
+    await res.text();
+    // Give the onRequest callback a tick to fire (it's behind a void promise).
+    await new Promise((r) => setTimeout(r, 20));
+    restore();
+
+    expect(captured).toBeDefined();
+    expect(captured!.usage?.input_tokens).toBe(123);
+    expect(captured!.usage?.output_tokens).toBe(7);
+    expect(captured!.usage?.cache_read_input_tokens).toBe(100);
+    expect(captured!.firstByteMs).toBeTypeOf('number');
+  });
+
+  it('extracts usage tokens from an SSE stream (message_start event)', async () => {
+    const sseBody =
+      'event: message_start\n' +
+      'data: ' +
+      JSON.stringify({
+        type: 'message_start',
+        message: {
+          id: 'msg_2',
+          type: 'message',
+          role: 'assistant',
+          content: [],
+          usage: {
+            input_tokens: 42,
+            output_tokens: 0,
+            cache_creation_input_tokens: 5000,
+            cache_read_input_tokens: 0,
+          },
+        },
+      }) +
+      '\n\n' +
+      'event: content_block_delta\n' +
+      'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}}\n\n' +
+      'event: message_stop\n' +
+      'data: {"type":"message_stop"}\n\n';
+
+    const restore = mockUpstream(
+      () =>
+        new Response(sseBody, {
+          status: 200,
+          headers: { 'content-type': 'text/event-stream' },
+        }),
+    );
+
+    let captured: ProxyEvent | undefined;
+    const proxy = createProxy({
+      transform: {},
+      onRequest: (e) => {
+        captured = e;
+      },
+    });
+
+    const res = await proxy(
+      new Request('http://localhost/v1/messages', {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        body: SAMPLE_REQ_BODY,
+      }),
+    );
+    await res.text();
+    await new Promise((r) => setTimeout(r, 20));
+    restore();
+
+    expect(captured).toBeDefined();
+    expect(captured!.usage?.input_tokens).toBe(42);
+    expect(captured!.usage?.cache_creation_input_tokens).toBe(5000);
+  });
+
+  it('fires the event with undefined usage when the response is an error', async () => {
+    const restore = mockUpstream(
+      () =>
+        new Response(JSON.stringify({ error: { type: 'overloaded_error' } }), {
+          status: 529,
+          headers: { 'content-type': 'application/json' },
+        }),
+    );
+
+    let captured: ProxyEvent | undefined;
+    const proxy = createProxy({
+      transform: {},
+      onRequest: (e) => {
+        captured = e;
+      },
+    });
+
+    const res = await proxy(
+      new Request('http://localhost/v1/messages', {
+        method: 'POST',
+        headers: { 'content-type': 'application/json' },
+        body: SAMPLE_REQ_BODY,
+      }),
+    );
+    await res.text();
+    await new Promise((r) => setTimeout(r, 20));
+    restore();
+
+    expect(captured).toBeDefined();
+    expect(captured!.status).toBe(529);
+    expect(captured!.usage).toBeUndefined();
+  });
+});