fix(providers): record hosted-stream failures provider-agnostically at the chokepoint (covers gemini)

Author: TheodoreSpeaks

apps/sim/providers/index.ts

@@ -10,6 +10,7 @@ import {
   uploadLargeFilesToProvider,
 } from '@/providers/file-attachments.server'
 import { getProviderExecutor } from '@/providers/registry'
+import { recordHostedStreamFailure } from '@/providers/streaming-execution'
 import type { ProviderId, ProviderRequest, ProviderResponse } from '@/providers/types'
 import {
   calculateCost,
@@ -229,14 +230,17 @@ export async function executeProviderRequest(
     if (isBYOK) {
       zeroCostForBYOK(response)
     } else if (hostedKeyEnvVar) {
-      // Hosted key used: record usage now; cost is settled on stream drain via
-      // settleStreamingLlmCost (from createStreamingExecution, or the provider's
-      // bespoke stream finalizer for gemini).
-      hostedKeyMetrics.recordUsed({
-        provider: providerId,
-        tool: response.execution.output?.model ?? request.model,
-        key: hostedKeyEnvVar,
-      })
+      // Hosted key used: record usage at dispatch. Cost is settled per-provider on
+      // successful stream drain (createStreamingExecution, or gemini's bespoke
+      // finalizer); a mid-stream error records a failure here — provider-agnostic,
+      // so it covers gemini and any non-wrapper stream too.
+      const model = response.execution.output?.model ?? request.model
+      hostedKeyMetrics.recordUsed({ provider: providerId, tool: model, key: hostedKeyEnvVar })
+      response.stream = recordHostedStreamFailure(
+        response.stream,
+        { provider: providerId, envVar: hostedKeyEnvVar },
+        model
+      )
     }
     return response
   }

apps/sim/providers/streaming-execution.test.ts

@@ -20,7 +20,10 @@ vi.mock('@/lib/api-key/hosted-cost', () => ({
 }))
 
 import type { NormalizedBlockOutput } from '@/executor/types'
-import { createStreamingExecution } from '@/providers/streaming-execution'
+import {
+  createStreamingExecution,
+  recordHostedStreamFailure,
+} from '@/providers/streaming-execution'
 
 /**
  * Builds a fake stream factory mirroring the providers' `createReadableStreamFrom*`
@@ -90,27 +93,19 @@ describe('createStreamingExecution', () => {
     expect(mockRecordFailed).not.toHaveBeenCalled()
   })
 
-  it('records a hosted-key failure (not cost) when the stream errors mid-drain', async () => {
+  it('recordHostedStreamFailure records a failure (not cost) when the stream errors', async () => {
     mockRecordCostCharged.mockClear()
     mockRecordFailed.mockClear()
     const boom = new Error('upstream 500')
-    const sourceStream = new ReadableStream({
-      pull: (c) => c.error(boom),
-    })
+    const sourceStream = new ReadableStream({ pull: (c) => c.error(boom) })
 
-    const result = createStreamingExecution({
-      model: 'test-model',
-      providerStartTime,
-      providerStartTimeISO,
-      timing: { kind: 'simple', segmentName: 'test-model' },
-      initialTokens: { input: 0, output: 0, total: 0 },
-      initialCost: { input: 0, output: 0, total: 0 },
-      hostedKey: { provider: 'openai', envVar: 'OPENAI_API_KEY_1' },
-      cached: false,
-      createStream: () => sourceStream,
-    })
+    const wrapped = recordHostedStreamFailure(
+      sourceStream,
+      { provider: 'openai', envVar: 'OPENAI_API_KEY_1' },
+      'test-model'
+    )
 
-    const reader = result.stream.getReader()
+    const reader = wrapped.getReader()
     await expect(reader.read()).rejects.toThrow('upstream 500')
 
     // Failure recorded once; no cost charged for a failed stream.
@@ -124,6 +119,20 @@ describe('createStreamingExecution', () => {
     expect(mockRecordCostCharged).not.toHaveBeenCalled()
   })
 
+  it('recordHostedStreamFailure does not record a failure when the stream completes', async () => {
+    mockRecordFailed.mockClear()
+    const wrapped = recordHostedStreamFailure(
+      new ReadableStream({ start: (c) => c.close() }),
+      { provider: 'openai', envVar: 'OPENAI_API_KEY_1' },
+      'test-model'
+    )
+    const reader = wrapped.getReader()
+    while (!(await reader.read()).done) {
+      // drain
+    }
+    expect(mockRecordFailed).not.toHaveBeenCalled()
+  })
+
   it('assembles the simple (no-tools) shape and finalizes timing on drain', () => {
     const drainTime = 5_000
     vi.spyOn(Date, 'now').mockReturnValue(drainTime)

apps/sim/providers/streaming-execution.ts

@@ -6,16 +6,13 @@ import type { TimeSegment } from '@/providers/types'
 import { calculateCost } from '@/providers/utils'
 
 /**
- * Passthrough of `source` that runs exactly one terminal callback: `onDrain`
- * when it completes normally (cost is settled then), or `onError` when a read
- * errors mid-stream (records the hosted-key failure). A client `cancel` is not a
- * key failure, so it runs neither. Lets hosted-key cost/failure metrics stay
- * symmetric regardless of whether the provider's drain callback finalized timing.
+ * Passthrough of `source` that runs at most one terminal callback: `onDrain` when
+ * it completes normally, or `onError` when a read errors mid-stream. A client
+ * `cancel` runs neither (an abort is not a key failure).
  */
-function settleHostedCostOnStreamDrain(
+function tapStreamTermination(
   source: ReadableStream,
-  onDrain: () => void,
-  onError: (error: unknown) => void
+  callbacks: { onDrain?: () => void; onError?: (error: unknown) => void }
 ): ReadableStream {
   const reader = source.getReader()
   let finished = false
@@ -26,7 +23,7 @@ function settleHostedCostOnStreamDrain(
         if (done) {
           if (!finished) {
             finished = true
-            onDrain()
+            callbacks.onDrain?.()
           }
           controller.close()
           return
@@ -35,7 +32,7 @@ function settleHostedCostOnStreamDrain(
       } catch (error) {
         if (!finished) {
           finished = true
-          onError(error)
+          callbacks.onError?.(error)
         }
         controller.error(error)
       }
@@ -46,6 +43,29 @@ function settleHostedCostOnStreamDrain(
   })
 }
 
+/**
+ * Wrap a hosted-key streaming response so a mid-stream read error records a
+ * hosted-key failure metric. Applied provider-agnostically at the chokepoint
+ * (`executeProviderRequest`) so it covers every provider — including ones that
+ * build streams bespoke (gemini) and don't go through {@link createStreamingExecution}.
+ * Cost on success is settled per-provider; this only handles the failure leg.
+ */
+export function recordHostedStreamFailure(
+  source: ReadableStream,
+  hostedKey: { provider: string; envVar: string },
+  model: string
+): ReadableStream {
+  return tapStreamTermination(source, {
+    onError: (error) =>
+      hostedKeyMetrics.recordFailed({
+        provider: hostedKey.provider,
+        tool: model,
+        key: hostedKey.envVar,
+        reason: classifyHostedKeyFailure(error),
+      }),
+  })
+}
+
 /**
  * Settle the authoritative streaming LLM cost onto `output.cost` from its final
  * tokens (the single cost seam shared with the non-streaming path), and — on the
@@ -253,18 +273,11 @@ export function createStreamingExecution(
   // returned stream and settle once the source completes (final tokens are set by
   // the provider's drain callback before the stream closes). Recomputes the
   // authoritative cost with the multiplier and emits the cost metric exactly once.
+  // Failure on error is handled provider-agnostically in executeProviderRequest.
   const stream = hostedKey
-    ? settleHostedCostOnStreamDrain(
-        baseStream,
-        () => settleStreamingLlmCost(output, model, hostedKey, cached ?? false),
-        (error) =>
-          hostedKeyMetrics.recordFailed({
-            provider: hostedKey.provider,
-            tool: model,
-            key: hostedKey.envVar,
-            reason: classifyHostedKeyFailure(error),
-          })
-      )
+    ? tapStreamTermination(baseStream, {
+        onDrain: () => settleStreamingLlmCost(output, model, hostedKey, cached ?? false),
+      })
     : baseStream
 
   return {