From ac2410b0ed89dbe888b107f5a85f9e910403deed Mon Sep 17 00:00:00 2001 From: Rio Yu <52408936+rioyu123@users.noreply.github.com> Date: Sun, 30 Aug 2026 00:10:32 +0800 Subject: [PATCH] fix(token-meter): stop consuming abandoned streams --- .../tests/token-meter-anthropic.test.ts | 66 ++++++ .../tests/token-meter-openai.test.ts | 94 ++++++++ llm/token-meter/token-meter.ts | 200 ++++++++++++------ 3 files changed, 293 insertions(+), 67 deletions(-) diff --git a/llm/token-meter/tests/token-meter-anthropic.test.ts b/llm/token-meter/tests/token-meter-anthropic.test.ts index a52ae4ba..4387ded2 100644 --- a/llm/token-meter/tests/token-meter-anthropic.test.ts +++ b/llm/token-meter/tests/token-meter-anthropic.test.ts @@ -3,6 +3,7 @@ */ import Stripe from 'stripe'; +import {Stream as AnthropicStream} from '@anthropic-ai/sdk/streaming'; import {createTokenMeter} from '../token-meter'; import type {MeterConfig} from '../types'; @@ -231,6 +232,71 @@ describe('TokenMeter - Anthropic Provider', () => { }); describe('Messages - Streaming', () => { + it('does not consume the source before the returned stream is read', async () => { + const meter = createTokenMeter(TEST_API_KEY, config); + let pulls = 0; + + async function* chunks() { + pulls += 1; + yield { + type: 'message_start', + message: { + id: 'msg_lazy', + model: 'claude-3-5-sonnet-20241022', + usage: {input_tokens: 10, output_tokens: 0}, + }, + }; + } + + const source = new AnthropicStream(chunks, new AbortController()); + const wrapped = meter.trackUsageStreamAnthropic(source as any, 'cus_123'); + + await new Promise(resolve => setImmediate(resolve)); + + expect(pulls).toBe(0); + expect(wrapped).toBeInstanceOf(AnthropicStream); + }); + + it('closes the source when the returned stream is abandoned', async () => { + const meter = createTokenMeter(TEST_API_KEY, config); + let finalized = false; + let release!: () => void; + const blocked = new Promise(resolve => { + release = resolve; + }); + + async function* chunks() { + try { + yield { + type: 'message_start', + message: { + id: 'msg_cancel', + model: 'claude-3-5-sonnet-20241022', + usage: {input_tokens: 10, output_tokens: 0}, + }, + }; + await blocked; + yield {type: 'message_stop'}; + } finally { + finalized = true; + } + } + + const source = new AnthropicStream(chunks, new AbortController()); + const wrapped = meter.trackUsageStreamAnthropic(source as any, 'cus_123'); + + try { + for await (const _chunk of wrapped) { + break; + } + await new Promise(resolve => setImmediate(resolve)); + + expect(finalized).toBe(true); + } finally { + release(); + } + }); + it('should track usage from basic streaming message', async () => { const meter = createTokenMeter(TEST_API_KEY, config); diff --git a/llm/token-meter/tests/token-meter-openai.test.ts b/llm/token-meter/tests/token-meter-openai.test.ts index 4a2dd51a..b624398d 100644 --- a/llm/token-meter/tests/token-meter-openai.test.ts +++ b/llm/token-meter/tests/token-meter-openai.test.ts @@ -3,6 +3,7 @@ */ import Stripe from 'stripe'; +import {Stream as OpenAIStream} from 'openai/streaming'; import {createTokenMeter} from '../token-meter'; import type {MeterConfig} from '../types'; @@ -227,6 +228,99 @@ describe('TokenMeter - OpenAI Provider', () => { }); describe('Chat Completions - Streaming', () => { + it('does not consume the source before the returned stream is read', async () => { + const meter = createTokenMeter(TEST_API_KEY, config); + let pulls = 0; + + async function* chunks() { + pulls += 1; + yield { + id: 'chatcmpl-lazy', + object: 'chat.completion.chunk', + created: Date.now(), + model: 'gpt-4o-mini', + choices: [], + }; + } + + const source = new OpenAIStream(chunks, new AbortController()); + const wrapped = meter.trackUsageStreamOpenAI(source as any, 'cus_123'); + + await new Promise(resolve => setImmediate(resolve)); + + expect(pulls).toBe(0); + expect(wrapped).toBeInstanceOf(OpenAIStream); + }); + + it('closes the source when the returned stream is abandoned', async () => { + const meter = createTokenMeter(TEST_API_KEY, config); + let finalized = false; + let release!: () => void; + const blocked = new Promise(resolve => { + release = resolve; + }); + + async function* chunks() { + try { + yield { + id: 'chatcmpl-cancel', + object: 'chat.completion.chunk', + created: Date.now(), + model: 'gpt-4o-mini', + choices: [], + }; + await blocked; + yield { + id: 'chatcmpl-unused', + object: 'chat.completion.chunk', + created: Date.now(), + model: 'gpt-4o-mini', + choices: [], + }; + } finally { + finalized = true; + } + } + + const source = new OpenAIStream(chunks, new AbortController()); + const wrapped = meter.trackUsageStreamOpenAI(source as any, 'cus_123'); + + try { + for await (const _chunk of wrapped) { + break; + } + await new Promise(resolve => setImmediate(resolve)); + + expect(finalized).toBe(true); + } finally { + release(); + } + }); + + it('propagates a source error once and finalizes the source', async () => { + const meter = createTokenMeter(TEST_API_KEY, config); + const sourceError = new Error('source failed before the first chunk'); + let finalized = false; + + async function* chunks(): AsyncGenerator { + try { + throw sourceError; + } finally { + finalized = true; + } + } + + const source = new OpenAIStream(chunks, new AbortController()); + const wrapped = meter.trackUsageStreamOpenAI(source as any, 'cus_123'); + const iterator = wrapped[Symbol.asyncIterator](); + + await expect(iterator.next()).rejects.toBe(sourceError); + await expect(iterator.next()).resolves.toEqual({done: true, value: undefined}); + await new Promise(resolve => setImmediate(resolve)); + + expect(finalized).toBe(true); + }); + it('should track usage from basic streaming chat', async () => { const meter = createTokenMeter(TEST_API_KEY, config); diff --git a/llm/token-meter/token-meter.ts b/llm/token-meter/token-meter.ts index 7a3b30d3..3185c6c5 100644 --- a/llm/token-meter/token-meter.ts +++ b/llm/token-meter/token-meter.ts @@ -4,7 +4,7 @@ import Stripe from 'stripe'; import type OpenAI from 'openai'; -import type {Stream} from 'openai/streaming'; +import type {Stream as OpenAIStream} from 'openai/streaming'; import type Anthropic from '@anthropic-ai/sdk'; import type {Stream as AnthropicStream} from '@anthropic-ai/sdk/streaming'; import type { @@ -16,12 +16,30 @@ import {logUsageEvent} from './meter-event-logging'; import { detectResponse, isGeminiStream, - extractUsageFromChatStream, - extractUsageFromResponseStream, - extractUsageFromAnthropicStream, type DetectedResponse, } from './utils/type-detection'; +function wrapSdkStream & {controller: AbortController}>( + source: T, + iterator: () => AsyncIterator +): T { + const StreamConstructor = source.constructor as new ( + iterator: () => AsyncIterator, + controller: AbortController + ) => T; + + // The SDK streams are class instances. Keeping this fallback makes the + // structural test doubles accepted by the existing public type usable too. + if ((StreamConstructor as unknown) === Object) { + return { + controller: source.controller, + [Symbol.asyncIterator]: iterator, + } as T; + } + + return new StreamConstructor(iterator, source.controller); +} + /** * Supported response types from all AI providers */ @@ -36,8 +54,8 @@ export type SupportedResponse = * Supported stream types from all AI providers */ export type SupportedStream = - | Stream - | Stream + | OpenAIStream + | OpenAIStream | AnthropicStream | GenerateContentStreamResult; @@ -58,8 +76,8 @@ export interface TokenMeter { */ trackUsageStreamOpenAI< T extends - | Stream - | Stream + | OpenAIStream + | OpenAIStream >( stream: T, stripeCustomerId: string @@ -172,78 +190,126 @@ export function createTokenMeter( trackUsageStreamOpenAI< T extends - | Stream - | Stream + | OpenAIStream + | OpenAIStream >(stream: T, stripeCustomerId: string): T { - const [peekStream, stream2] = stream.tee(); + const meteredStream = wrapSdkStream( + stream, + async function* () { + let streamType: 'chat_completion' | 'response_api' | null = null; + let model = ''; + let inputTokens = 0; + let outputTokens = 0; - (async () => { - // Peek at the first chunk to determine stream type - const [stream1, meterStream] = peekStream.tee(); - const reader = stream1[Symbol.asyncIterator](); - const firstChunk = await reader.next(); - - let detected: DetectedResponse | null = null; - - if (!firstChunk.done && firstChunk.value) { - const chunk = firstChunk.value as any; - - // Check if it's an OpenAI Chat stream (has choices array) - if ('choices' in chunk && Array.isArray(chunk.choices)) { - detected = await extractUsageFromChatStream(meterStream as any); - } - // Check if it's an OpenAI Response API stream (has type starting with 'response.') - else if (chunk.type && typeof chunk.type === 'string' && chunk.type.startsWith('response.')) { - detected = await extractUsageFromResponseStream(meterStream as any); - } - else { - console.warn('Unable to detect OpenAI stream type from first chunk:', chunk); + for await (const value of stream) { + const chunk = value as any; + + if (!streamType) { + if ('choices' in chunk && Array.isArray(chunk.choices)) { + streamType = 'chat_completion'; + } else if ( + chunk.type && + typeof chunk.type === 'string' && + chunk.type.startsWith('response.') + ) { + streamType = 'response_api'; + } + } + + if (streamType === 'chat_completion') { + model = chunk.model || model; + if (chunk.usage) { + inputTokens = chunk.usage.prompt_tokens ?? 0; + outputTokens = chunk.usage.completion_tokens ?? 0; + } + } else if (streamType === 'response_api' && chunk.response) { + model = chunk.response.model || model; + if (chunk.response.usage) { + inputTokens = chunk.response.usage.input_tokens ?? 0; + outputTokens = chunk.response.usage.output_tokens ?? 0; + } + } + + yield value; } - } - if (detected) { - logUsageEvent(stripeClient, config, { - model: detected.model, - provider: detected.provider, - usage: { - inputTokens: detected.inputTokens, - outputTokens: detected.outputTokens, - }, - stripeCustomerId, - }); - } else { - console.warn('Unable to extract usage from OpenAI stream'); + const detected: DetectedResponse | null = + streamType && model + ? { + provider: 'openai', + type: streamType, + model, + inputTokens, + outputTokens, + } + : null; + + if (detected) { + logUsageEvent(stripeClient, config, { + model: detected.model, + provider: detected.provider, + usage: { + inputTokens: detected.inputTokens, + outputTokens: detected.outputTokens, + }, + stripeCustomerId, + }); + } else { + console.warn('Unable to extract usage from OpenAI stream'); + } } - })(); + ); - return stream2 as T; + return meteredStream as T; }, trackUsageStreamAnthropic( stream: AnthropicStream, stripeCustomerId: string ): AnthropicStream { - const [peekStream, stream2] = stream.tee(); - - (async () => { - const detected = await extractUsageFromAnthropicStream(peekStream); - - if (detected) { - logUsageEvent(stripeClient, config, { - model: detected.model, - provider: detected.provider, - usage: { - inputTokens: detected.inputTokens, - outputTokens: detected.outputTokens, - }, - stripeCustomerId, - }); - } else { - console.warn('Unable to extract usage from Anthropic stream'); - } - })(); + return wrapSdkStream( + stream, + async function* () { + let model = ''; + let inputTokens = 0; + let outputTokens = 0; + + for await (const chunk of stream) { + if (chunk.type === 'message_start') { + model = chunk.message.model; + inputTokens = chunk.message.usage.input_tokens ?? 0; + } else if (chunk.type === 'message_delta') { + outputTokens = chunk.usage.output_tokens ?? 0; + } + + yield chunk; + } - return stream2; + const detected: DetectedResponse | null = model + ? { + provider: 'anthropic', + type: 'chat_completion', + model, + inputTokens, + outputTokens, + } + : null; + + if (detected) { + logUsageEvent(stripeClient, config, { + model: detected.model, + provider: detected.provider, + usage: { + inputTokens: detected.inputTokens, + outputTokens: detected.outputTokens, + }, + stripeCustomerId, + }); + } else { + console.warn('Unable to extract usage from Anthropic stream'); + } + } + ) as AnthropicStream; }, }; }