fix(providers): settle hosted-key streaming cost on stream drain, not finalizeTiming

Author: TheodoreSpeaks

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

@@ -2,6 +2,17 @@
  * @vitest-environment node
  */
 import { describe, expect, it, vi } from 'vitest'
+
+const { mockRecordCostCharged } = vi.hoisted(() => ({ mockRecordCostCharged: 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 },
+}))
+
 import type { NormalizedBlockOutput } from '@/executor/types'
 import { createStreamingExecution } from '@/providers/streaming-execution'
 
@@ -27,6 +38,51 @@ describe('createStreamingExecution', () => {
   const providerStartTime = 1_000
   const providerStartTimeISO = new Date(providerStartTime).toISOString()
 
+  it('settles hosted-key cost on stream drain even when finalizeTiming is never called (post-tool path)', async () => {
+    mockRecordCostCharged.mockClear()
+    // A source stream that closes immediately, mirroring a drained provider stream.
+    const sourceStream = new ReadableStream({ start: (c) => c.close() })
+
+    const result = createStreamingExecution({
+      model: 'test-model',
+      providerStartTime,
+      providerStartTimeISO,
+      timing: {
+        kind: 'accumulated',
+        modelTime: 1,
+        toolsTime: 0,
+        firstResponseTime: 1,
+        iterations: 1,
+        timeSegments: [],
+      },
+      initialTokens: { input: 0, output: 0, total: 0 },
+      initialCost: { input: 0, output: 0, total: 0 },
+      hostedKey: { provider: 'openai', envVar: 'OPENAI_API_KEY_1' },
+      cached: false,
+      // Post-tool streaming path: sets final tokens but never calls finalizeTiming.
+      createStream: ({ output }) => {
+        output.tokens = { input: 100, output: 50, total: 150 }
+        return sourceStream
+      },
+    })
+
+    // Cost not settled until the stream is actually drained.
+    expect(mockRecordCostCharged).not.toHaveBeenCalled()
+
+    const reader = result.stream.getReader()
+    while (!(await reader.read()).done) {
+      // drain
+    }
+
+    // Settlement ran on drain: cost recomputed from final tokens, metric emitted once.
+    expect(result.execution.output.cost).toEqual({ input: 1, output: 2, total: 3, pricing: {} })
+    expect(mockRecordCostCharged).toHaveBeenCalledTimes(1)
+    expect(mockRecordCostCharged).toHaveBeenCalledWith(3, {
+      provider: 'openai',
+      tool: 'test-model',
+    })
+  })
+
   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

@@ -4,6 +4,41 @@ 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.
+ */
+function settleHostedCostOnStreamDrain(
+  source: ReadableStream,
+  onDrain: () => void
+): ReadableStream {
+  const reader = source.getReader()
+  let settled = false
+  return new ReadableStream({
+    async pull(controller) {
+      try {
+        const { done, value } = await reader.read()
+        if (done) {
+          if (!settled) {
+            settled = true
+            onDrain()
+          }
+          controller.close()
+          return
+        }
+        controller.enqueue(value)
+      } catch (error) {
+        controller.error(error)
+      }
+    },
+    cancel(reason) {
+      return reader.cancel(reason)
+    },
+  })
+}
+
 /**
  * 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
@@ -200,19 +235,23 @@ export function createStreamingExecution(
   }
 
   const timingKind = timing.kind
-  const stream = createStream({
+  const baseStream = createStream({
     output,
-    finalizeTiming: () => {
-      // Hosted-key path: the wrapper owns the authoritative cost — recompute from
-      // the now-final tokens (with the cost multiplier the per-provider streaming
-      // paths omit) and emit the hosted-key cost metric exactly once on drain.
-      if (hostedKey) {
-        settleStreamingLlmCost(output, model, hostedKey, cached ?? false)
-      }
-      finalizeTiming(output, providerStartTime, timingKind)
-    },
+    finalizeTiming: () => finalizeTiming(output, providerStartTime, timingKind),
   })
 
+  // Settle hosted-key cost on actual stream drain. This must NOT hang off the
+  // provider's `finalizeTiming` call — the post-tool streaming paths
+  // (`createStream: ({ output }) => …`) never invoke it — so instead we wrap the
+  // 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.
+  const stream = hostedKey
+    ? settleHostedCostOnStreamDrain(baseStream, () =>
+        settleStreamingLlmCost(output, model, hostedKey, cached ?? false)
+      )
+    : baseStream
+
   return {
     stream,
     execution: {