fix(providers): record hosted-key failure when a hosted stream errors mid-drain

Author: TheodoreSpeaks

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

@@ -3,14 +3,20 @@
  */
 import { describe, expect, it, vi } from 'vitest'
 
-const { mockRecordCostCharged } = vi.hoisted(() => ({ mockRecordCostCharged: vi.fn() }))
+const { mockRecordCostCharged, mockRecordFailed } = vi.hoisted(() => ({
+  mockRecordCostCharged: vi.fn(),
+  mockRecordFailed: vi.fn(),
+}))
 
 vi.mock('@/providers/utils', () => ({
   calculateCost: vi.fn(() => ({ input: 1, output: 2, total: 3, pricing: {} })),
 }))
 vi.mock('@/lib/core/config/env-flags', () => ({ getCostMultiplier: () => 1 }))
 vi.mock('@/lib/monitoring/metrics', () => ({
-  hostedKeyMetrics: { recordCostCharged: mockRecordCostCharged },
+  hostedKeyMetrics: { recordCostCharged: mockRecordCostCharged, recordFailed: mockRecordFailed },
+}))
+vi.mock('@/lib/api-key/hosted-cost', () => ({
+  classifyHostedKeyFailure: () => 'other',
 }))
 
 import type { NormalizedBlockOutput } from '@/executor/types'
@@ -81,6 +87,41 @@ describe('createStreamingExecution', () => {
       provider: 'openai',
       tool: 'test-model',
     })
+    expect(mockRecordFailed).not.toHaveBeenCalled()
+  })
+
+  it('records a hosted-key failure (not cost) when the stream errors mid-drain', async () => {
+    mockRecordCostCharged.mockClear()
+    mockRecordFailed.mockClear()
+    const boom = new Error('upstream 500')
+    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 reader = result.stream.getReader()
+    await expect(reader.read()).rejects.toThrow('upstream 500')
+
+    // Failure recorded once; no cost charged for a failed stream.
+    expect(mockRecordFailed).toHaveBeenCalledTimes(1)
+    expect(mockRecordFailed).toHaveBeenCalledWith({
+      provider: 'openai',
+      tool: 'test-model',
+      key: 'OPENAI_API_KEY_1',
+      reason: 'other',
+    })
+    expect(mockRecordCostCharged).not.toHaveBeenCalled()
   })
 
   it('assembles the simple (no-tools) shape and finalizes timing on drain', () => {

apps/sim/providers/streaming-execution.ts

@@ -1,35 +1,42 @@
+import { classifyHostedKeyFailure } from '@/lib/api-key/hosted-cost'
 import { getCostMultiplier } from '@/lib/core/config/env-flags'
 import { hostedKeyMetrics } from '@/lib/monitoring/metrics'
 import type { NormalizedBlockOutput, StreamingExecution } from '@/executor/types'
 import type { TimeSegment } from '@/providers/types'
 import { calculateCost } from '@/providers/utils'
 
 /**
- * Passthrough of `source` that invokes `onDrain` exactly once, after the source
- * completes naturally (not on error or cancel). Lets the hosted-key cost settle
- * when final tokens are known, regardless of whether the provider's drain
- * callback finalized timing.
+ * 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.
  */
 function settleHostedCostOnStreamDrain(
   source: ReadableStream,
-  onDrain: () => void
+  onDrain: () => void,
+  onError: (error: unknown) => void
 ): ReadableStream {
   const reader = source.getReader()
-  let settled = false
+  let finished = false
   return new ReadableStream({
     async pull(controller) {
       try {
         const { done, value } = await reader.read()
         if (done) {
-          if (!settled) {
-            settled = true
+          if (!finished) {
+            finished = true
             onDrain()
           }
           controller.close()
           return
         }
         controller.enqueue(value)
       } catch (error) {
+        if (!finished) {
+          finished = true
+          onError(error)
+        }
         controller.error(error)
       }
     },
@@ -247,8 +254,16 @@ export function createStreamingExecution(
   // the provider's drain callback before the stream closes). Recomputes the
   // authoritative cost with the multiplier and emits the cost metric exactly once.
   const stream = hostedKey
-    ? settleHostedCostOnStreamDrain(baseStream, () =>
-        settleStreamingLlmCost(output, model, hostedKey, cached ?? false)
+    ? settleHostedCostOnStreamDrain(
+        baseStream,
+        () => settleStreamingLlmCost(output, model, hostedKey, cached ?? false),
+        (error) =>
+          hostedKeyMetrics.recordFailed({
+            provider: hostedKey.provider,
+            tool: model,
+            key: hostedKey.envVar,
+            reason: classifyHostedKeyFailure(error),
+          })
       )
     : baseStream