diff --git a/changelog.d/fixes/16521-aborted-stream-partial-usage.md b/changelog.d/fixes/16521-aborted-stream-partial-usage.md new file mode 100644 index 00000000000..9ed511aa9f0 --- /dev/null +++ b/changelog.d/fixes/16521-aborted-stream-partial-usage.md @@ -0,0 +1 @@ +- **fix(sse):** record partial token usage (provider-reported, or estimated from the streamed output and flagged as estimated) for streams the client aborts or that fail mid-answer, instead of zero; status 499/success semantics unchanged ([#16521](https://github.com/diegosouzapw/OmniRoute/pull/16521)) — thanks @brimob-sowax, @ntdat812, @ntdatt812 diff --git a/config/quality/eslint-suppressions.json b/config/quality/eslint-suppressions.json index 800987162fc..798b24484f0 100644 --- a/config/quality/eslint-suppressions.json +++ b/config/quality/eslint-suppressions.json @@ -645,7 +645,7 @@ }, "open-sse/utils/stream.ts": { "@typescript-eslint/no-unused-vars": { - "count": 2 + "count": 1 } }, "open-sse/utils/streamHelpers.ts": { diff --git a/open-sse/handlers/chatCore/streamingTail.ts b/open-sse/handlers/chatCore/streamingTail.ts index dc44627a0ee..b86751d509c 100644 --- a/open-sse/handlers/chatCore/streamingTail.ts +++ b/open-sse/handlers/chatCore/streamingTail.ts @@ -13,6 +13,7 @@ import { import { incrementTokenUsage } from "../../services/geminiRateLimitTracker.ts"; import { forgetClaudeThinkingForClientResponse } from "../../services/claudeThinkingStore.ts"; import { requestTtftMs, streamEmittedOutput } from "../../utils/streamTiming.ts"; +import { readStreamPartialUsage } from "../../utils/streamPartialUsage.ts"; import { clearPostOutputFailureStreak } from "../../services/accountFallback/postOutputFailureStreak.ts"; import { createRoutingEvent, @@ -355,6 +356,7 @@ export async function runStreamingTail(deps: StreamingTailDeps) { firstOutputMs, itlMs: streamItlMs, interrupted: _streamInterrupted, + partialUsage: streamPartialUsage, }) => { const ttft = requestTtftMs(streamTimingOriginOffsetMs, firstOutputMs); const normalizedStreamStatus = streamStatus || 200; @@ -510,7 +512,9 @@ export async function runStreamingTail(deps: StreamingTailDeps) { persistAttemptLogs({ status: normalizedStreamStatus, error: streamError || undefined, - tokens: streamUsage || {}, + // A stream that ended early (499 / mid-stream failure) has no `usage`, but its call log + // still shows the tokens already generated (estimates are flagged via `estimated`). + tokens: streamUsage || streamPartialUsage || {}, responseBody: streamResponseBody ?? undefined, providerRequest: finalBody || translatedBody, providerResponse: providerPayload, @@ -609,6 +613,7 @@ export async function runStreamingTail(deps: StreamingTailDeps) { persistFailureUsage, onStreamFailure, hasEmittedOutput: () => streamEmittedOutput(transformStream), + readPartialUsage: () => readStreamPartialUsage(transformStream), }); const handleStreamFailure = streamFailureFinalizers.handleStreamFailure; onPipelineStreamError = streamFailureFinalizers.onPipelineStreamError; diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 4d2cf060e26..9e7adb469c4 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -92,6 +92,7 @@ import { restoreClaudeToolName } from "../services/claudeCodeToolRemapper.ts"; import { normalizeFinalOpenAIStreamChunk } from "./openAIStreamChunk.ts"; import { collectClaudeDelta } from "./streamClaudeDelta.ts"; import { createStreamTiming, registerStreamTiming, type StreamTiming } from "./streamTiming.ts"; +import { trackStreamPartialUsage } from "./streamPartialUsage.ts"; import { buildUsageOnlyChunk } from "./usageOnlyChunk.ts"; import { normalizeArrayContentChunk } from "./arrayContentDelta.ts"; @@ -3240,12 +3241,11 @@ export function createSSEStream(options: StreamOptions = {}) { console.log(`[STREAM] Error in flush (${model || "unknown"}):`, error.message || error); } }, - cancel(reason) { - clearIdleTimer(); - }, + cancel: () => clearIdleTimer(), }, ...createByteLengthQueueStrategies(streamBufferBytes) ); + trackStreamPartialUsage(sseStream, body, () => [state?.usage ?? usage, totalContentLength]); return registerStreamTiming(sseStream, timing); } diff --git a/open-sse/utils/streamFailureFinalization.ts b/open-sse/utils/streamFailureFinalization.ts index 0ee307b0a3d..0f1fad8bf28 100644 --- a/open-sse/utils/streamFailureFinalization.ts +++ b/open-sse/utils/streamFailureFinalization.ts @@ -6,6 +6,26 @@ import { import { HTTP_STATUS } from "../config/constants.ts"; import { buildErrorBody } from "./error.ts"; import { sanitizeErrorMessage } from "./errorSanitization.ts"; +import type { PartialStreamUsage } from "./streamPartialUsage.ts"; + +/** Token aggregate accepted by the failure usage row (chatCore/failureUsage.ts shape). */ +export type PartialFailureUsageAggregate = { + prompt_tokens: number; + completion_tokens: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + reasoning_tokens?: number; +}; + +function toPartialFailureAggregate(usage: PartialStreamUsage): PartialFailureUsageAggregate { + return { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + reasoning_tokens: usage.reasoning_tokens, + }; +} export type StreamCompletionPayload = { status: number; @@ -16,6 +36,13 @@ export type StreamCompletionPayload = { error?: string | null; errorCode?: string | null; ttft?: number | null; + /** + * Tokens already generated when a stream ended early (client disconnect / mid-stream + * failure): provider-reported, or estimated from the streamed output (`estimated: true`). + * Kept apart from `usage` so a failed stream does not trigger the success-only usage row, + * cost and calibration side effects; the call log reads it for its token columns. + */ + partialUsage?: PartialStreamUsage; }; export type StreamFailurePayload = { @@ -195,19 +222,39 @@ export function createStreamFailureFinalizers({ persistFailureUsage, onStreamFailure, hasEmittedOutput = () => false, + readPartialUsage = () => null, }: { isFailureCompletionRecorded: () => boolean; isStreamCompletionRecorded?: () => boolean; onStreamComplete: (payload: StreamCompletionPayload) => void; - persistFailureUsage: (status: number, errorCode?: string) => void; + persistFailureUsage: ( + status: number, + errorCode?: string, + aggregate?: PartialFailureUsageAggregate + ) => void; onStreamFailure?: ((failure: StreamFailurePayload) => void) | null; hasEmittedOutput?: () => boolean; + /** Usage the stream accumulated before it ended early (see streamPartialUsage.ts). */ + readPartialUsage?: () => PartialStreamUsage | null; }) { + const readPartialUsageSafely = (): PartialStreamUsage | null => { + try { + return readPartialUsage() ?? null; + } catch { + return null; // best-effort accounting must never block finalization + } + }; + const handleStreamFailure = (failure: StreamFailurePayload) => { if (isStreamCompletionRecorded()) { return true; } + // The client hung up / the upstream died mid-answer: the provider already generated (and + // billed) these tokens. Status and success semantics stay as they are — only the token + // counts of the failure records stop being zero. + const partialUsage = readPartialUsageSafely(); + const status = failure.status || HTTP_STATUS.BAD_GATEWAY; const message = failure.message || "Upstream stream error"; const classification = @@ -225,10 +272,15 @@ export function createStreamFailureFinalizers({ error: message, errorCode: projectedCode, ttft: 0, + ...(partialUsage ? { partialUsage } : {}), }); } - persistFailureUsage(status, projectedCode); + if (partialUsage) { + persistFailureUsage(status, projectedCode, toPartialFailureAggregate(partialUsage)); + } else { + persistFailureUsage(status, projectedCode); + } try { onStreamFailure?.({ ...failure, outputEmitted: hasEmittedOutput() }); } catch { diff --git a/open-sse/utils/streamPartialUsage.ts b/open-sse/utils/streamPartialUsage.ts new file mode 100644 index 00000000000..622df13a8d5 --- /dev/null +++ b/open-sse/utils/streamPartialUsage.ts @@ -0,0 +1,144 @@ +/** + * Partial token usage of a stream that ended early. + * + * When the client hangs up (499 `client_disconnected`) or the upstream fails mid-answer, the SSE + * transform's `flush()` never runs, so the usage it would have resolved there is lost and the + * failure record used to carry zero tokens — even though the provider already generated (and + * billed) part of the answer. `createSSEStream` registers a lazy reader here; the failure + * finalizer reads it once, at finalization time. Nothing is computed per chunk: the reader only + * closes over the transform's live `usage` / `totalContentLength` variables. + * + * Same WeakMap-side-table pattern as `registerStreamTiming` (streamTiming.ts): the stream object + * is not monkey-patched and the entry dies with it. + */ +import { + estimateInputTokens, + estimateOutputTokens, + hasValidUsage, + normalizeUsage, +} from "./usageTracking.ts"; + +export type PartialStreamUsage = { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + reasoning_tokens?: number; + /** Set when any side of the count is a local estimate rather than provider-reported. */ + estimated?: true; +}; + +type PartialUsageReader = () => PartialStreamUsage | null; + +const partialUsageByStream = new WeakMap(); + +/** Associate a stream with the reader of its partial usage. */ +export function registerStreamPartialUsage( + stream: T, + reader: PartialUsageReader +): T { + partialUsageByStream.set(stream, reader); + return stream; +} + +/** Partial usage of the stream so far, or null when it has none (or is unknown). */ +export function readStreamPartialUsage( + stream: object | null | undefined +): PartialStreamUsage | null { + if (!stream) return null; + const reader = partialUsageByStream.get(stream); + if (!reader) return null; + try { + return reader(); + } catch { + return null; // accounting is best-effort — never break failure finalization + } +} + +/** + * Wire `createSSEStream` to the side table: `read` returns the transform's live + * `[reportedUsage, streamedContentLength]` pair (translate mode accumulates provider usage on + * `state.usage`, passthrough on `usage`). Nothing is computed until the failure finalizer + * reads it, so a stream that ends before `flush()` (client disconnect / mid-stream failure) + * still reports the tokens it already generated. + */ +export function trackStreamPartialUsage( + stream: T, + body: unknown, + read: () => readonly [reported: unknown, contentLength: number] +): T { + return registerStreamPartialUsage(stream, () => { + const [reported, contentLength] = read(); + return resolvePartialStreamUsage({ reported, contentLength, body }); + }); +} + +function positive(value: unknown): number { + const numeric = typeof value === "number" ? value : Number(value); + return Number.isFinite(numeric) && numeric > 0 ? numeric : 0; +} + +/** + * Resolve the usage to record for a stream that ended before its terminal event. + * + * - Provider-reported usage wins (normalized to `prompt_tokens` / `completion_tokens`). + * - A side the provider has not reported yet (typically the output count, which Claude sends in + * the final `message_delta`) is estimated from what was actually streamed, and the result is + * flagged `estimated`. + * - Nothing reported and nothing streamed → null: no tokens are fabricated. + */ +export function resolvePartialStreamUsage({ + reported, + contentLength, + body, +}: { + reported: unknown; + contentLength: number; + body: unknown; +}): PartialStreamUsage | null { + const streamedChars = positive(contentLength); + const record = + reported && typeof reported === "object" && hasValidUsage(reported as Record) + ? (reported as Record) + : null; + if (!record && streamedChars === 0) return null; + + const normalized = record ? (normalizeUsage(record) ?? {}) : {}; + let prompt = + positive(normalized.prompt_tokens) || + positive(normalized.input_tokens) || + positive(record?.promptTokenCount); + let completion = + positive(normalized.completion_tokens) || + positive(normalized.output_tokens) || + positive(record?.candidatesTokenCount); + + let estimated = record?.estimated === true; + if (streamedChars > 0) { + if (completion === 0) { + completion = estimateOutputTokens(streamedChars); + estimated = true; + } + if (prompt === 0) { + prompt = estimateInputTokens(body); + estimated = true; + } + } + if (prompt === 0 && completion === 0) return null; + + const usage: PartialStreamUsage = { + prompt_tokens: prompt, + completion_tokens: completion, + total_tokens: prompt + completion, + }; + const cacheRead = + positive(normalized.cache_read_input_tokens) || positive(normalized.cached_tokens); + if (cacheRead) usage.cache_read_input_tokens = cacheRead; + const cacheCreation = positive(normalized.cache_creation_input_tokens); + if (cacheCreation) usage.cache_creation_input_tokens = cacheCreation; + const reasoning = positive(normalized.reasoning_tokens); + if (reasoning) usage.reasoning_tokens = reasoning; + if (estimated) usage.estimated = true; + return usage; +} diff --git a/tests/unit/stream-aborted-partial-usage.test.ts b/tests/unit/stream-aborted-partial-usage.test.ts new file mode 100644 index 00000000000..fdc65c31964 --- /dev/null +++ b/tests/unit/stream-aborted-partial-usage.test.ts @@ -0,0 +1,379 @@ +/** + * A stream that ends early — the client hangs up (499 `client_disconnected`) or the upstream + * fails mid-answer — used to be finalized with ZERO tokens even after the provider had already + * generated (and billed) part of the answer. The failure finalizer now carries the partial usage + * the SSE transform accumulated: provider-reported usage when it arrived, otherwise an estimate + * from the streamed output length (flagged `estimated`). Status / success semantics are unchanged. + */ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { FORMATS } from "../../open-sse/translator/formats.ts"; +import { + createPassthroughStreamWithLogger, + createSSETransformStreamWithLogger, +} from "../../open-sse/utils/stream.ts"; +import { createStreamFailureFinalizers } from "../../open-sse/utils/streamFailureFinalization.ts"; +import { + readStreamPartialUsage, + resolvePartialStreamUsage, + trackStreamPartialUsage, +} from "../../open-sse/utils/streamPartialUsage.ts"; + +const encoder = new TextEncoder(); + +const REQUEST_BODY = { + model: "gpt-test", + stream: true, + messages: [{ role: "user", content: "Write a long story about a lighthouse keeper." }], +}; + +function openAiChunk(delta: Record, extra: Record = {}) { + return `data: ${JSON.stringify({ + id: "chatcmpl-partial", + object: "chat.completion.chunk", + model: "gpt-test", + choices: [{ index: 0, delta, finish_reason: null }], + ...extra, + })}\n\n`; +} + +type TransformLike = { readable: ReadableStream; writable: WritableStream }; + +const teardowns: Array<() => Promise> = []; +test.afterEach(async () => { + // The client "hangs up": cancelling the readable runs the transform's cancel() (clears its + // idle timer) — never its flush(), exactly like a real disconnect. + for (const teardown of teardowns.splice(0)) await teardown(); +}); + +/** Push SSE text through the transform and read what it forwards, WITHOUT closing it (no flush). */ +async function pumpWithoutFlush(stream: TransformLike, frames: string[]): Promise { + const writer = stream.writable.getWriter(); + const reader = stream.readable.getReader(); + teardowns.push(async () => { + await reader.cancel("client_disconnected").catch(() => {}); + await writer.abort("client_disconnected").catch(() => {}); + }); + // Keep a reader draining (like a live client socket) so backpressure never stalls the writes. + void (async () => { + try { + while (!(await reader.read()).done) { + /* drain */ + } + } catch { + /* stream torn down */ + } + })(); + for (const frame of frames) { + await writer.write(encoder.encode(frame)); + } +} + +type PersistCall = [number, string | undefined, Record | undefined]; + +function finalizersWith(readPartialUsage?: () => unknown) { + const persisted: unknown[][] = []; + const completions: Record[] = []; + const finalizers = createStreamFailureFinalizers({ + isFailureCompletionRecorded: () => false, + onStreamComplete: (payload) => { + completions.push(payload as unknown as Record); + }, + persistFailureUsage: (...args: unknown[]) => { + persisted.push(args); + }, + readPartialUsage: readPartialUsage as never, + }); + return { ...finalizers, persisted, completions }; +} + +// ── finalizer: carries the partial usage into the 499 record ──────────────────────────────── + +test("client disconnect records the provider-reported partial usage instead of zero", () => { + const { handleStreamFailure, persisted, completions } = finalizersWith(() => ({ + prompt_tokens: 100, + completion_tokens: 42, + total_tokens: 142, + })); + + handleStreamFailure({ + status: 499, + message: "Client disconnected: request_signal_aborted", + code: "client_disconnected", + type: "client_disconnected", + }); + + assert.equal(persisted.length, 1); + const [status, code, aggregate] = persisted[0] as PersistCall; + assert.equal(status, 499); + assert.equal(code, "client_disconnected"); + assert.ok(aggregate, "the failure usage row must receive the partial token aggregate"); + assert.equal(aggregate.prompt_tokens, 100); + assert.equal(aggregate.completion_tokens, 42); + + assert.equal(completions.length, 1); + assert.equal(completions[0].status, 499); + // `usage` stays null so the success-only side effects (usage row, cost, calibration) do not + // fire a second time; the call log reads the partial usage from its own field. + assert.equal(completions[0].usage, null); + assert.deepEqual(completions[0].partialUsage, { + prompt_tokens: 100, + completion_tokens: 42, + total_tokens: 142, + }); +}); + +test("a mid-stream upstream failure also carries the partial usage", () => { + const { onPipelineStreamError, persisted } = finalizersWith(() => ({ + prompt_tokens: 12, + completion_tokens: 7, + total_tokens: 19, + })); + + onPipelineStreamError({ message: "terminated", statusCode: 502 }); + + assert.equal(persisted.length, 1); + const [status, code, aggregate] = persisted[0] as PersistCall; + assert.equal(status, 502); + assert.equal(code, "stream_terminated"); + assert.equal(aggregate?.prompt_tokens, 12); + assert.equal(aggregate?.completion_tokens, 7); +}); + +test("no partial usage keeps the zero-token failure record unchanged", () => { + const { handleStreamFailure, persisted, completions } = finalizersWith(() => null); + + handleStreamFailure({ status: 499, message: "gone", code: "client_disconnected" }); + + assert.deepEqual(persisted, [[499, "client_disconnected"]]); + assert.equal(completions[0].usage, null); + assert.equal("partialUsage" in completions[0], false); +}); + +test("a throwing partial-usage reader never blocks the failure finalization", () => { + const { handleStreamFailure, persisted, completions } = finalizersWith(() => { + throw new Error("boom"); + }); + + handleStreamFailure({ status: 499, message: "gone", code: "client_disconnected" }); + + assert.deepEqual(persisted, [[499, "client_disconnected"]]); + assert.equal(completions.length, 1); +}); + +test("a recorded completion wins: the disconnect finalizer neither reads nor writes usage", () => { + let reads = 0; + const persisted: unknown[][] = []; + const { handleStreamFailure } = createStreamFailureFinalizers({ + isFailureCompletionRecorded: () => false, + isStreamCompletionRecorded: () => true, + onStreamComplete: () => assert.fail("completion already recorded"), + persistFailureUsage: (...args: unknown[]) => { + persisted.push(args); + }, + readPartialUsage: () => { + reads++; + return { prompt_tokens: 1, completion_tokens: 1 }; + }, + }); + + handleStreamFailure({ status: 499, message: "late", code: "client_disconnected" }); + + assert.equal(reads, 0); + assert.deepEqual(persisted, []); +}); + +// ── pure resolver ─────────────────────────────────────────────────────────────────────────── + +test("resolvePartialStreamUsage: reported usage is returned verbatim, not estimated", () => { + const usage = resolvePartialStreamUsage({ + reported: { prompt_tokens: 100, completion_tokens: 42, total_tokens: 142 }, + contentLength: 5000, + body: REQUEST_BODY, + }); + assert.deepEqual(usage, { prompt_tokens: 100, completion_tokens: 42, total_tokens: 142 }); +}); + +test("resolvePartialStreamUsage: no reported usage + emitted output → flagged estimate", () => { + const usage = resolvePartialStreamUsage({ + reported: null, + contentLength: 350, + body: REQUEST_BODY, + }); + assert.ok(usage); + assert.ok(usage.prompt_tokens > 0, "input estimated from the request body"); + assert.equal(usage.completion_tokens, 100, "350 chars ≈ 100 output tokens"); + assert.equal(usage.total_tokens, usage.prompt_tokens + usage.completion_tokens); + assert.equal(usage.estimated, true); +}); + +test("resolvePartialStreamUsage: nothing reported and nothing emitted → null (no fabrication)", () => { + assert.equal( + resolvePartialStreamUsage({ reported: null, contentLength: 0, body: REQUEST_BODY }), + null + ); + assert.equal( + resolvePartialStreamUsage({ reported: {}, contentLength: 0, body: REQUEST_BODY }), + null + ); +}); + +test("resolvePartialStreamUsage: Claude-style input-only usage gets the output estimated", () => { + // message_start reported the prompt; the client left before message_delta's output count. + const usage = resolvePartialStreamUsage({ + reported: { prompt_tokens: 900, cache_read_input_tokens: 300 }, + contentLength: 70, + body: REQUEST_BODY, + }); + assert.ok(usage); + assert.equal(usage.prompt_tokens, 900, "reported input is kept"); + assert.equal(usage.cache_read_input_tokens, 300); + assert.equal(usage.completion_tokens, 20); + assert.equal(usage.total_tokens, 920); + assert.equal(usage.estimated, true); +}); + +test("resolvePartialStreamUsage: input_tokens/output_tokens shape is normalized", () => { + const usage = resolvePartialStreamUsage({ + reported: { input_tokens: 30, output_tokens: 9 }, + contentLength: 0, + body: REQUEST_BODY, + }); + assert.deepEqual(usage, { prompt_tokens: 30, completion_tokens: 9, total_tokens: 39 }); +}); + +// ── real SSE transform: the snapshot reflects what was streamed so far ───────────────────── + +test("passthrough stream: content without a usage frame yields an estimated partial usage", async () => { + const stream = createPassthroughStreamWithLogger( + "openai", + null, + null, + "gpt-test", + null, + REQUEST_BODY, + () => assert.fail("flush must not run — the client is still mid-stream"), + null, + null, + FORMATS.OPENAI + ); + assert.equal(readStreamPartialUsage(stream), null, "nothing streamed yet"); + + await pumpWithoutFlush(stream, [ + openAiChunk({ role: "assistant", content: "Once upon a time, " }), + openAiChunk({ content: "a keeper climbed the stairs of the old lighthouse." }), + ]); + + const usage = readStreamPartialUsage(stream); + assert.ok(usage, "partial usage must be available before flush"); + assert.ok(usage.completion_tokens > 0); + assert.ok(usage.prompt_tokens > 0); + assert.equal(usage.estimated, true); +}); + +test("passthrough stream: a provider usage frame is preferred over the estimate", async () => { + const stream = createPassthroughStreamWithLogger( + "openai", + null, + null, + "gpt-test", + null, + REQUEST_BODY, + null, + null, + null, + FORMATS.OPENAI + ); + + await pumpWithoutFlush(stream, [ + openAiChunk( + { role: "assistant", content: "Partial answer" }, + { usage: { prompt_tokens: 100, completion_tokens: 42, total_tokens: 142 } } + ), + ]); + + const usage = readStreamPartialUsage(stream); + assert.equal(usage?.prompt_tokens, 100); + assert.equal(usage?.completion_tokens, 42); + assert.equal(usage?.estimated, undefined); +}); + +test("translate stream (OpenAI upstream → Claude client): partial output is accounted", async () => { + const stream = createSSETransformStreamWithLogger( + FORMATS.OPENAI, + FORMATS.CLAUDE, + "openai", + null, + null, + "gpt-test", + null, + REQUEST_BODY, + () => assert.fail("flush must not run — the client is still mid-stream") + ); + + await pumpWithoutFlush(stream, [ + openAiChunk({ role: "assistant", content: "The lamp turned slowly over the black water." }), + ]); + + const usage = readStreamPartialUsage(stream); + assert.ok(usage); + assert.ok(usage.completion_tokens > 0); +}); + +test("end to end: disconnect after streamed output persists non-zero tokens with status 499", async () => { + const stream = createPassthroughStreamWithLogger( + "openai", + null, + null, + "gpt-test", + null, + REQUEST_BODY, + null, + null, + null, + FORMATS.OPENAI + ); + const { handleStreamFailure, persisted } = finalizersWith(() => readStreamPartialUsage(stream)); + + await pumpWithoutFlush(stream, [openAiChunk({ role: "assistant", content: "x".repeat(700) })]); + handleStreamFailure({ + status: 499, + message: "Client disconnected: request_signal_aborted", + code: "client_disconnected", + type: "client_disconnected", + }); + + const [status, code, aggregate] = persisted[0] as PersistCall; + assert.equal(status, 499); + assert.equal(code, "client_disconnected"); + assert.equal(aggregate?.completion_tokens, 200); + assert.ok((aggregate?.prompt_tokens as number) > 0); +}); + +test("readStreamPartialUsage: unknown / missing streams read as null", () => { + assert.equal(readStreamPartialUsage(null), null); + assert.equal(readStreamPartialUsage(undefined), null); + assert.equal(readStreamPartialUsage({}), null); +}); + +test("trackStreamPartialUsage: reads the live usage/length lazily, only when finalized", () => { + const stream = {}; + let reads = 0; + let reported: unknown = null; + let contentLength = 0; + const returned = trackStreamPartialUsage(stream, REQUEST_BODY, () => { + reads += 1; + return [reported, contentLength]; + }); + assert.equal(returned, stream, "returns the same stream for chaining"); + assert.equal(reads, 0, "nothing is computed at registration time"); + + reported = { prompt_tokens: 12, completion_tokens: 34 }; + contentLength = 200; + const usage = readStreamPartialUsage(stream); + assert.equal(reads, 1); + assert.equal(usage?.prompt_tokens, 12); + assert.equal(usage?.completion_tokens, 34); + assert.equal(usage?.estimated, undefined, "provider-reported usage is not flagged"); +});