fix(stream): persist mid-stream failures

open-sse/handlers/chatCore.ts

@@ -20,6 +20,7 @@ import { ensureStreamReadiness } from "../utils/streamReadiness.ts";
 import { synthesizeOpenAiSseFromJson } from "../utils/jsonToSse.ts";
 import { resolveStreamReadinessTimeout } from "../utils/streamReadinessPolicy.ts";
 import { createStreamController, pipeWithDisconnect } from "../utils/streamHandler.ts";
+import * as streamFailure from "../utils/streamFailureFinalization.ts";
 import { createSseHeartbeatTransform, shapeForClientFormat } from "../utils/sseHeartbeat.ts";
 import { addBufferToUsage, filterUsageForFormat, estimateUsage } from "../utils/usageTracking.ts";
 import {
@@ -295,10 +296,6 @@ function cloneBoundedChatLogPayload(value: unknown, depth = 0): unknown {
 }
 
 import { estimateSizeFast, isSmallEnoughForSemanticCache } from "../utils/estimateSize.ts";
-import {
-  finalizeMostRecentPendingRequest,
-  finalizePendingRequestById,
-} from "@/lib/usage/usageHistory.ts";
 
 const MAX_LOG_BODY_CHARS = 8 * 1024; // 8KB cap for logged request/response bodies
 /**
@@ -3688,9 +3685,12 @@ export async function handleChatCore({
     };
   };
 
+  let onPipelineStreamError: streamFailure.PipelineStreamErrorHandler | null = null;
+
   // Create stream controller for disconnect detection
   const streamController = createStreamController({
     onDisconnect,
+    onError: (event) => onPipelineStreamError?.(event),
     provider,
     model,
     connectionId,
@@ -5495,18 +5495,27 @@ export async function handleChatCore({
     (finalBody as Record<string, unknown> | null | undefined) ?? null
   );
 
+  let streamFailureCompletionRecorded = false;
+
   // Callback to save call log when stream completes (include responseBody when provided by stream)
   const onStreamComplete = ({
     status: streamStatus,
     usage: streamUsage,
     responseBody: streamResponseBody,
     providerPayload,
     clientPayload,
+    error: streamError,
+    errorCode: streamErrorCode,
     ttft,
   }) => {
+    const normalizedStreamStatus = streamStatus || 200;
+    if (normalizedStreamStatus !== 200) {
+      if (streamFailureCompletionRecorded) return;
+      streamFailureCompletionRecorded = true;
+    }
     const cacheUsageLogMeta = buildCacheUsageLogMeta(streamUsage);
 
-    if (streamStatus === 200) {
+    if (normalizedStreamStatus === 200) {
       void maybeSyncClaudeExtraUsageState({
         provider,
         connectionId,
@@ -5517,7 +5526,7 @@ export async function handleChatCore({
 
     // Reasoning Replay Cache (#1628): Capture reasoning_content from streaming responses
     // with tool_calls so it can be replayed on subsequent turns (DeepSeek V4, Kimi K2, etc.)
-    if (streamStatus === 200 && streamResponseBody) {
+    if (normalizedStreamStatus === 200 && streamResponseBody) {
       try {
         const body = streamResponseBody as Record<string, unknown>;
         const choices = body.choices as { message?: Record<string, unknown> }[] | undefined;
@@ -5532,23 +5541,17 @@ export async function handleChatCore({
     }
     effectiveServiceTier = resolveReportedServiceTier(streamResponseBody) ?? effectiveServiceTier;
 
-    try {
-      const finalizedConnId = connectionId || credentials?.connectionId || null;
-      const completedById = finalizePendingRequestById(pendingRequestId, {
-        providerResponse: providerPayload ?? streamResponseBody ?? undefined,
-        clientResponse: clientPayload ?? streamResponseBody ?? undefined,
-      });
-      if (!completedById) {
-        finalizeMostRecentPendingRequest(model, provider, finalizedConnId, {
-          providerResponse: providerPayload ?? streamResponseBody ?? undefined,
-          clientResponse: clientPayload ?? streamResponseBody ?? undefined,
-        });
-      }
-    } catch (e) {
-      try {
-        console.warn("finalizeMostRecentPendingRequest failed:", e && (e.message || e));
-      } catch {}
-    }
+    streamFailure.finalizeStreamRequestLog({
+      pendingRequestId,
+      model,
+      provider,
+      connectionId: connectionId || credentials?.connectionId || null,
+      providerResponse: providerPayload ?? streamResponseBody ?? undefined,
+      clientResponse: clientPayload ?? streamResponseBody ?? undefined,
+      status: normalizedStreamStatus,
+      error: streamError,
+      errorCode: streamErrorCode,
+    });
 
     // Track cache token metrics for streaming responses
     if (streamUsage && typeof streamUsage === "object") {
@@ -5558,11 +5561,12 @@ export async function handleChatCore({
         provider: provider || "unknown",
         model: model || "unknown",
         tokens: streamUsage,
-        status: String(streamStatus || 200),
-        success: streamStatus === 200,
+        status: String(normalizedStreamStatus),
+        success: normalizedStreamStatus === 200,
         latencyMs: Date.now() - startTime,
         timeToFirstTokenMs: ttft,
-        errorCode: null,
+        errorCode:
+          normalizedStreamStatus === 200 ? null : streamErrorCode || String(normalizedStreamStatus),
         timestamp: new Date().toISOString(),
         connectionId: connectionId || undefined,
         apiKeyId: apiKeyInfo?.id || undefined,
@@ -5573,7 +5577,7 @@ export async function handleChatCore({
         console.error("Failed to save usage stats:", err.message);
       });
 
-      if (apiKeyInfo?.id && streamStatus === 200) {
+      if (apiKeyInfo?.id && normalizedStreamStatus === 200) {
         try {
           const billable = computeBillableTokens(streamUsage);
           if (billable > 0)
@@ -5585,7 +5589,8 @@ export async function handleChatCore({
     }
 
     persistAttemptLogs({
-      status: streamStatus || 200,
+      status: normalizedStreamStatus,
+      error: streamError || undefined,
       tokens: streamUsage || {},
       responseBody: streamResponseBody ?? undefined,
       providerRequest: finalBody || translatedBody,
@@ -5608,7 +5613,7 @@ export async function handleChatCore({
     // Resolve the real per-request cost (calculateCost) so USD-unit pools accrue
     // on streaming traffic too; this previously recorded usd:0 hardcoded, which
     // meant DeepSeek-style `usd/monthly` shared pools never blocked on streams.
-    if (apiKeyInfo?.id && credentials?.connectionId && streamStatus === 200) {
+    if (apiKeyInfo?.id && credentials?.connectionId && normalizedStreamStatus === 200) {
       const quotaApiKeyId = apiKeyInfo.id;
       const quotaConnectionId = credentials.connectionId;
       // onStreamComplete is sync — use .then() (fire-and-forget, fail-open) instead of await
@@ -5621,7 +5626,7 @@ export async function handleChatCore({
               provider,
               model,
               streamUsage,
-              streamStatus,
+              streamStatus: normalizedStreamStatus,
               serviceTier: effectiveServiceTier,
             },
             { calculateCost, log }
@@ -5681,19 +5686,14 @@ export async function handleChatCore({
     }
   };
 
-  const handleStreamFailure = (failure: {
-    status: number;
-    message: string;
-    code?: string;
-    type?: string;
-  }) => {
-    persistFailureUsage(failure.status || HTTP_STATUS.BAD_GATEWAY, failure.code || failure.type);
-    try {
-      onStreamFailure?.(failure);
-    } catch {
-      // Best-effort fallback state update only.
-    }
-  };
+  const streamFailureFinalizers = streamFailure.createStreamFailureFinalizers({
+    isFailureCompletionRecorded: () => streamFailureCompletionRecorded,
+    onStreamComplete,
+    persistFailureUsage,
+    onStreamFailure,
+  });
+  const handleStreamFailure = streamFailureFinalizers.handleStreamFailure;
+  onPipelineStreamError = streamFailureFinalizers.onPipelineStreamError;
 
   // For providers using Responses API format, translate stream back to openai (Chat Completions) format
   // UNLESS client is Droid CLI which expects openai-responses format back

open-sse/translator/response/claude-to-openai.ts

@@ -115,6 +115,21 @@ export function claudeToOpenAIResponse(chunk, state) {
       // Extract usage from message_delta event (Claude native format)
       // Normalize to OpenAI format (prompt_tokens/completion_tokens) for consistent logging
       if (chunk.usage && typeof chunk.usage === "object") {
+        const previousUsage = state.usage && typeof state.usage === "object" ? state.usage : {};
+        const previousInputTokens =
+          typeof previousUsage.input_tokens === "number"
+            ? previousUsage.input_tokens
+            : typeof previousUsage.prompt_tokens === "number"
+              ? previousUsage.prompt_tokens
+              : 0;
+        const previousCacheReadTokens =
+          typeof previousUsage.cache_read_input_tokens === "number"
+            ? previousUsage.cache_read_input_tokens
+            : 0;
+        const previousCacheCreationTokens =
+          typeof previousUsage.cache_creation_input_tokens === "number"
+            ? previousUsage.cache_creation_input_tokens
+            : 0;
         const inputTokens =
           typeof chunk.usage.input_tokens === "number" ? chunk.usage.input_tokens : 0;
         const outputTokens =
@@ -136,7 +151,10 @@ export function claudeToOpenAIResponse(chunk, state) {
         // minimum, so a 2-token "hi" can be reported as ~2008 prompt_tokens and
         // inflate downstream billing ~250x. cache_creation is still exposed
         // separately via prompt_tokens_details.cache_creation_tokens below.
-        const billableInputTokens = inputTokens + cacheReadTokens;
+        const billableInputTokens =
+          inputTokens > 0 || cacheReadTokens > 0 || cacheCreationTokens > 0
+            ? inputTokens + cacheReadTokens
+            : previousInputTokens;
         state.usage = {
           prompt_tokens: billableInputTokens,
           completion_tokens: outputTokens,
@@ -145,11 +163,13 @@ export function claudeToOpenAIResponse(chunk, state) {
         };
 
         // Store cache tokens if present (needed for prompt_tokens_details in final chunk)
-        if (cacheReadTokens > 0) {
-          state.usage.cache_read_input_tokens = cacheReadTokens;
+        const effectiveCacheReadTokens = cacheReadTokens || previousCacheReadTokens;
+        const effectiveCacheCreationTokens = cacheCreationTokens || previousCacheCreationTokens;
+        if (effectiveCacheReadTokens > 0) {
+          state.usage.cache_read_input_tokens = effectiveCacheReadTokens;
         }
-        if (cacheCreationTokens > 0) {
-          state.usage.cache_creation_input_tokens = cacheCreationTokens;
+        if (effectiveCacheCreationTokens > 0) {
+          state.usage.cache_creation_input_tokens = effectiveCacheCreationTokens;
         }
       }
 

open-sse/utils/stream.ts

@@ -93,6 +93,9 @@ type StreamCompletePayload = {
   responseBody?: unknown;
   providerPayload?: unknown;
   clientPayload?: unknown;
+  error?: string | null;
+  errorCode?: string | null;
+  ttft?: number | null;
 };
 
 type StreamFailurePayload = {
@@ -116,7 +119,7 @@ type StreamOptions = {
   apiKeyInfo?: unknown;
   body?: unknown;
   onComplete?: ((payload: StreamCompletePayload) => void) | null;
-  onFailure?: ((payload: StreamFailurePayload) => void | Promise<void>) | null;
+  onFailure?: ((payload: StreamFailurePayload) => boolean | void | Promise<void>) | null;
 };
 
 type TranslateState = ReturnType<typeof initState> & {
@@ -846,12 +849,13 @@ export function createSSEStream(options: StreamOptions = {}) {
     reqLogger?.appendConvertedChunk?.(errOutput);
     clientPayloadCollector.push(errorEvent);
     controller.enqueue(encoder.encode(errOutput));
+    let failureHandled = false;
     if (onFailure) {
       try {
-        void onFailure({ status: 502, message: msg, code: "empty_response" });
+        failureHandled = onFailure({ status: 502, message: msg, code: "empty_response" }) === true;
       } catch {}
     }
-    if (decrementPendingRequest) {
+    if (decrementPendingRequest && !failureHandled) {
       clearPendingRequestFromStream();
     }
     controller.error(markPendingRequestCleared(new Error(msg)));
@@ -1063,7 +1067,21 @@ export function createSSEStream(options: StreamOptions = {}) {
               clearIdleTimer();
               const timeoutMsg = `[STREAM] Idle timeout: no data from ${provider || "provider"} for ${STREAM_IDLE_TIMEOUT_MS}ms (model: ${model || "unknown"})`;
               console.warn(timeoutMsg);
-              clearPendingRequestFromStream();
+              let failureHandled = false;
+              if (onFailure) {
+                try {
+                  failureHandled =
+                    onFailure({
+                      status: HTTP_STATUS.GATEWAY_TIMEOUT,
+                      message: timeoutMsg,
+                      code: "stream_idle_timeout",
+                      type: "timeout_error",
+                    }) === true;
+                } catch {}
+              }
+              if (!failureHandled) {
+                clearPendingRequestFromStream();
+              }
               appendRequestLog({
                 model,
                 provider,
@@ -1766,13 +1784,16 @@ export function createSSEStream(options: StreamOptions = {}) {
             reqLogger?.appendConvertedChunk?.(output);
             controller.enqueue(encoder.encode(output));
             if (failurePayload) {
+              let failureHandled = false;
               if (onFailure) {
                 try {
-                  void onFailure(failurePayload);
+                  failureHandled = onFailure(failurePayload) === true;
                 } catch {}
               }
               clearIdleTimer();
-              clearPendingRequestFromStream();
+              if (!failureHandled) {
+                clearPendingRequestFromStream();
+              }
               controller.error(
                 markPendingRequestCleared(new Error(failurePayload.message || "Upstream failure"))
               );
@@ -1929,7 +1950,25 @@ export function createSSEStream(options: StreamOptions = {}) {
 
           // Extract usage
           const extracted = extractUsage(parsed);
-          if (extracted) state.usage = extracted; // Keep original usage for logging
+          if (extracted) {
+            if (!state.usage) {
+              state.usage = extracted;
+            } else {
+              const su = state.usage as Record<string, number>;
+              const eu = extracted as Record<string, number>;
+              if (eu.prompt_tokens > 0) su.prompt_tokens = eu.prompt_tokens;
+              if (eu.completion_tokens > 0) su.completion_tokens = eu.completion_tokens;
+              if (eu.total_tokens > 0) su.total_tokens = eu.total_tokens;
+              if (eu.input_tokens > 0) su.input_tokens = eu.input_tokens;
+              if (eu.output_tokens > 0) su.output_tokens = eu.output_tokens;
+              if (eu.cache_read_input_tokens > 0)
+                su.cache_read_input_tokens = eu.cache_read_input_tokens;
+              if (eu.cache_creation_input_tokens > 0)
+                su.cache_creation_input_tokens = eu.cache_creation_input_tokens;
+              if (eu.cached_tokens > 0) su.cached_tokens = eu.cached_tokens;
+              if (eu.reasoning_tokens > 0) su.reasoning_tokens = eu.reasoning_tokens;
+            }
+          }
 
           // Translate: targetFormat -> openai -> sourceFormat
           const translated = translateResponse(targetFormat, sourceFormat, parsed, state);
@@ -2309,6 +2348,8 @@ export function createSSEStream(options: StreamOptions = {}) {
                   if (eu.prompt_tokens > 0) su.prompt_tokens = eu.prompt_tokens;
                   if (eu.completion_tokens > 0) su.completion_tokens = eu.completion_tokens;
                   if (eu.total_tokens > 0) su.total_tokens = eu.total_tokens;
+                  if (eu.input_tokens > 0) su.input_tokens = eu.input_tokens;
+                  if (eu.output_tokens > 0) su.output_tokens = eu.output_tokens;
                   if (eu.cache_read_input_tokens > 0)
                     su.cache_read_input_tokens = eu.cache_read_input_tokens;
                   if (eu.cache_creation_input_tokens > 0)
@@ -2336,15 +2377,16 @@ export function createSSEStream(options: StreamOptions = {}) {
 
           if (state?.upstreamError) {
             const err = state.upstreamError;
-            clearPendingRequestFromStream();
+            let failureHandled = false;
             if (onFailure) {
               try {
-                void onFailure({
-                  status: err.status,
-                  message: err.message,
-                  code: err.code,
-                  type: err.type,
-                });
+                failureHandled =
+                  onFailure({
+                    status: err.status,
+                    message: err.message,
+                    code: err.code,
+                    type: err.type,
+                  }) === true;
               } catch {}
             }
 
@@ -2355,6 +2397,8 @@ export function createSSEStream(options: StreamOptions = {}) {
                   status: err.status,
                   usage: state?.usage,
                   responseBody: errorBody,
+                  error: err.message,
+                  errorCode: err.code,
                   providerPayload: providerPayloadCollector.build(
                     buildStreamSummaryFromEvents(
                       providerPayloadCollector.getEvents(),
@@ -2367,10 +2411,14 @@ export function createSSEStream(options: StreamOptions = {}) {
                     includeEvents: false,
                   }),
                 });
+                failureHandled = true;
               } catch {}
             }
 
             clearIdleTimer();
+            if (!failureHandled) {
+              clearPendingRequestFromStream();
+            }
             controller.error(
               markPendingRequestCleared(new Error(err.message || "Upstream failure"))
             );