fix(sse): record partial token usage when a stream is aborted mid-answer

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

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": {

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;

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);
 }
 

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 {

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<object, PartialUsageReader>();
+
+/** Associate a stream with the reader of its partial usage. */
+export function registerStreamPartialUsage<T extends object>(
+  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<T extends object>(
+  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<string, number>)
+      ? (reported as Record<string, unknown>)
+      : 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;
+}