preserve streaming usage estimates

Author: Bill Leoutsakos

apps/sim/executor/execution/block-executor.test.ts

@@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { clearLargeValueCacheForTests } from '@/lib/execution/payloads/cache'
 import { isLargeArrayManifest } from '@/lib/execution/payloads/large-array-manifest-metadata'
 import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
+import { calculateStreamingCost } from '@/lib/tokenization'
 import { BlockType } from '@/executor/constants'
 import type { DAGNode } from '@/executor/dag/builder'
 import { BlockExecutor } from '@/executor/execution/block-executor'
@@ -954,6 +955,51 @@ describe('BlockExecutor streaming pump', () => {
     expect(state.getBlockOutput(block.id)?.content).toBe('offline answer')
   })
 
+  it('estimates missing streaming usage from resolved input before sanitizing logs', async () => {
+    const secret = `sk-${'resolved-secret-'.repeat(20)}`
+    const params = { prompt: '{{OPENAI_API_KEY}}', model: 'gpt-4o' }
+    const handler: BlockHandler = {
+      canHandle: () => true,
+      execute: async (_ctx, _block, resolvedInputs) => ({
+        stream: new ReadableStream({
+          start(controller) {
+            controller.enqueue(new TextEncoder().encode('streamed answer'))
+            controller.close()
+          },
+        }),
+        execution: {
+          success: true,
+          output: { content: '', model: 'gpt-4o' },
+          logs: [],
+          metadata: {
+            startTime: new Date().toISOString(),
+            endTime: new Date().toISOString(),
+            duration: 1,
+          },
+        },
+      }),
+    }
+    const { executor, block, state } = createExecutor(handler, params)
+    const ctx = createContext(state)
+    ctx.environmentVariables = { OPENAI_API_KEY: secret }
+
+    await executor.execute(ctx, createNode(block), block)
+
+    const expected = calculateStreamingCost(
+      'gpt-4o',
+      JSON.stringify({ prompt: secret, model: 'gpt-4o' }),
+      'streamed answer'
+    )
+    expect(state.getBlockOutput(block.id)?.tokens).toEqual(expected.tokens)
+    expect(state.getBlockOutput(block.id)?.cost).toEqual(expected.cost)
+    expect(ctx.blockLogs[0].input).toEqual({
+      prompt: '{{OPENAI_API_KEY}}',
+      model: 'gpt-4o',
+    })
+    expect(ctx.blockLogs[0].output?.tokens).toEqual(expected.tokens)
+    expect(JSON.stringify(ctx.blockLogs[0])).not.toContain(secret)
+  })
+
   it('throws on mid-stream provider error (no truncated success)', async () => {
     const handler = createAgentEventsStreamingHandler({
       failAfterText: 'partial',

apps/sim/executor/execution/block-executor.ts

@@ -7,6 +7,7 @@ import { getBaseUrl } from '@/lib/core/utils/urls'
 import { compactExecutionPayload } from '@/lib/execution/payloads/serializer'
 import { redactLargeValueRefsInValue } from '@/lib/logs/execution/pii-large-values'
 import { redactObjectStrings } from '@/lib/logs/execution/pii-redaction'
+import { processStreamingBlockLog } from '@/lib/tokenization'
 import {
   containsUserFileWithMetadata,
   hydrateUserFilesWithBase64,
@@ -1037,6 +1038,19 @@ export class BlockExecutor {
       if (!parsedForFormat) {
         executionOutput.content = fullContent
       }
+
+      // Fallback usage estimation must happen while the resolved input is
+      // still available. The log copy is sanitized later, and estimating from
+      // `{{ENV_VAR}}` placeholders would skew token counts and cost.
+      processStreamingBlockLog(
+        {
+          blockId,
+          blockType: block.metadata?.id,
+          input: resolvedInputs,
+          output: streamingExec.execution.output,
+        },
+        fullContent
+      )
     }
 
     if (streamingExec.onFullContent) {

apps/sim/lib/tokenization/streaming.test.ts

@@ -0,0 +1,25 @@
+import { describe, expect, it } from 'vitest'
+import { processStreamingBlockLog } from '@/lib/tokenization/streaming'
+
+describe('processStreamingBlockLog', () => {
+  it('does not estimate usage from sanitized environment references', () => {
+    const log = {
+      blockId: 'agent-1',
+      blockType: 'agent',
+      input: {
+        prompt: 'Use {{OPENAI_API_KEY}} to answer',
+        model: 'gpt-4o',
+      },
+      output: {
+        content: 'streamed answer',
+        model: 'gpt-4o',
+      },
+    }
+
+    expect(processStreamingBlockLog(log, 'streamed answer')).toBe(false)
+    expect(log.output).toEqual({
+      content: 'streamed answer',
+      model: 'gpt-4o',
+    })
+  })
+})

apps/sim/lib/tokenization/streaming.ts

@@ -16,11 +16,17 @@ import {
 import type { BlockLog } from '@/executor/types'
 
 const logger = createLogger('StreamingTokenization')
+const ENVIRONMENT_REFERENCE_PATTERN = /\{\{[^}]+\}\}/
+
+type StreamingTokenizationLog = Pick<BlockLog, 'blockId' | 'blockType' | 'input' | 'output'>
 
 /**
  * Processes a block log and adds tokenization data if needed
  */
-export function processStreamingBlockLog(log: BlockLog, streamedContent: string): boolean {
+export function processStreamingBlockLog(
+  log: StreamingTokenizationLog,
+  streamedContent: string
+): boolean {
   // Check if this block should be tokenized
   if (!isTokenizableBlockType(log.blockType)) {
     return false
@@ -47,6 +53,12 @@ export function processStreamingBlockLog(log: BlockLog, streamedContent: string)
 
     // Prepare input text from log
     const inputText = extractTextContent(log.input)
+    // Environment values are restored to references before logs leave the
+    // executor. Never use those shorter placeholders as a billing estimate:
+    // the executor performs this fallback once with the raw resolved input.
+    if (ENVIRONMENT_REFERENCE_PATTERN.test(inputText)) {
+      return false
+    }
 
     // Calculate streaming cost
     const systemPrompt =
@@ -101,7 +113,7 @@ export function processStreamingBlockLog(log: BlockLog, streamedContent: string)
 /**
  * Determines the appropriate model for a block
  */
-function getModelForBlock(log: BlockLog): string {
+function getModelForBlock(log: StreamingTokenizationLog): string {
   // Try to get model from output first
   if (log.output?.model?.trim()) {
     return log.output.model