diff --git a/src/core/proxy.ts b/src/core/proxy.ts index 3f18775..8a15121 100644 --- a/src/core/proxy.ts +++ b/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; +} { + // 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 => { + 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), diff --git a/src/core/types.ts b/src/core/types.ts index ef0d87d..d5370e0 100644 --- a/src/core/types.ts +++ b/src/core/types.ts @@ -57,6 +57,15 @@ export interface ToolDef { export type SystemField = string | Array; +/** 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[]; diff --git a/tests/proxy-usage.test.ts b/tests/proxy-usage.test.ts new file mode 100644 index 0000000..9d9428c --- /dev/null +++ b/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) { + // 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(); + }); +});