chore: import upstream snapshot with attribution

apps/sim/lib/tokenization/streaming.ts

@@ -0,0 +1,149 @@
+/**
+ * Streaming-specific tokenization helpers
+ */
+
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { calculateStreamingCost } from '@/lib/tokenization/calculators'
+import { TOKENIZATION_CONFIG } from '@/lib/tokenization/constants'
+import {
+  extractTextContent,
+  hasRealCostData,
+  hasRealTokenData,
+  isTokenizableBlockType,
+  logTokenizationDetails,
+} from '@/lib/tokenization/utils'
+import type { BlockLog } from '@/executor/types'
+
+const logger = createLogger('StreamingTokenization')
+
+/**
+ * Processes a block log and adds tokenization data if needed
+ */
+export function processStreamingBlockLog(log: BlockLog, streamedContent: string): boolean {
+  // Check if this block should be tokenized
+  if (!isTokenizableBlockType(log.blockType)) {
+    return false
+  }
+
+  // Check if we already have meaningful token/cost data
+  if (hasRealTokenData(log.output?.tokens) && hasRealCostData(log.output?.cost)) {
+    return false
+  }
+
+  // Skip recalculation if cost was explicitly set by the billing layer (e.g. BYOK zero cost)
+  if (log.output?.cost?.pricing) {
+    return false
+  }
+
+  // Check if we have content to tokenize
+  if (!streamedContent?.trim()) {
+    return false
+  }
+
+  try {
+    // Determine model to use
+    const model = getModelForBlock(log)
+
+    // Prepare input text from log
+    const inputText = extractTextContent(log.input)
+
+    // Calculate streaming cost
+    const systemPrompt =
+      typeof log.input?.systemPrompt === 'string' ? log.input.systemPrompt : undefined
+    const context = typeof log.input?.context === 'string' ? log.input.context : undefined
+    const messages = Array.isArray(log.input?.messages)
+      ? (log.input.messages as Array<{ role: string; content: string }>)
+      : undefined
+    const result = calculateStreamingCost(
+      model,
+      inputText,
+      streamedContent,
+      systemPrompt,
+      context,
+      messages
+    )
+
+    // Update the log output with tokenization data
+    if (!log.output) {
+      log.output = {}
+    }
+
+    log.output.tokens = result.tokens
+    log.output.cost = result.cost
+    log.output.model = result.model
+
+    logTokenizationDetails(`Streaming tokenization completed for ${log.blockType}`, {
+      blockId: log.blockId,
+      blockType: log.blockType,
+      model: result.model,
+      provider: result.provider,
+      inputLength: inputText.length,
+      outputLength: streamedContent.length,
+      tokens: result.tokens,
+      cost: result.cost,
+      method: result.method,
+    })
+
+    return true
+  } catch (error) {
+    logger.error(`Streaming tokenization failed for block ${log.blockId}`, {
+      blockType: log.blockType,
+      error: toError(error).message,
+      contentLength: streamedContent?.length || 0,
+    })
+
+    // Don't throw - graceful degradation
+    return false
+  }
+}
+
+/**
+ * Determines the appropriate model for a block
+ */
+function getModelForBlock(log: BlockLog): string {
+  // Try to get model from output first
+  if (log.output?.model?.trim()) {
+    return log.output.model
+  }
+
+  // Try to get model from input
+  const inputModel = log.input?.model
+  if (typeof inputModel === 'string' && inputModel.trim()) {
+    return inputModel
+  }
+
+  // Use block type specific defaults
+  const blockType = log.blockType
+  if (blockType === 'agent' || blockType === 'router' || blockType === 'evaluator') {
+    return TOKENIZATION_CONFIG.defaults.model
+  }
+
+  // Final fallback
+  return TOKENIZATION_CONFIG.defaults.model
+}
+
+/**
+ * Processes multiple block logs for streaming tokenization
+ */
+export function processStreamingBlockLogs(
+  logs: BlockLog[],
+  streamedContentMap: Map<string, string>
+): number {
+  let processedCount = 0
+
+  for (const log of logs) {
+    const content = streamedContentMap.get(log.blockId)
+    if (content && processStreamingBlockLog(log, content)) {
+      processedCount++
+    }
+  }
+
+  logger.info(`Streaming tokenization summary`, {
+    totalLogs: logs.length,
+    processedBlocks: processedCount,
+    streamedBlocks: streamedContentMap.size,
+  })
+
+  return processedCount
+}