improvement(trace): billing trace span typing

Author: icecrasher321

apps/sim/lib/logs/execution/logging-factory.test.ts

@@ -466,4 +466,37 @@ describe('calculateCostSummary', () => {
     expect(result.totalCost).toBe(0.03 + BASE_EXECUTION_CHARGE)
     expect(result.models['gpt-4o'].total).toBe(0.03)
   })
+
+  test('preserves parent toolCost while skipping model breakdown children', () => {
+    const traceSpans = [
+      {
+        id: 'agent-span',
+        type: 'agent',
+        model: 'gpt-4o',
+        cost: { input: 0.01, output: 0.02, toolCost: 0.015, total: 0.045 },
+        tokens: { input: 1000, output: 2000, total: 3000 },
+        children: [
+          {
+            id: 'agent-span-model-segment',
+            type: 'model',
+            model: 'gpt-4o',
+            cost: { input: 0.01, output: 0.02, total: 0.03 },
+            tokens: { input: 1000, output: 2000, total: 3000 },
+          },
+          {
+            id: 'agent-span-tool-segment',
+            type: 'tool',
+            name: 'firecrawl_scrape',
+          },
+        ],
+      },
+    ]
+
+    const result = calculateCostSummary(traceSpans)
+
+    expect(result.modelCost).toBe(0.045)
+    expect(result.totalCost).toBe(0.045 + BASE_EXECUTION_CHARGE)
+    expect(result.models['gpt-4o'].total).toBe(0.045)
+    expect(result.models['gpt-4o'].toolCost).toBe(0.015)
+  })
 })

apps/sim/lib/logs/execution/logging-factory.ts

@@ -1,7 +1,12 @@
 import { db, workflow } from '@sim/db'
 import { eq } from 'drizzle-orm'
 import { BASE_EXECUTION_CHARGE } from '@/lib/billing/constants'
-import type { ExecutionEnvironment, ExecutionTrigger, WorkflowState } from '@/lib/logs/types'
+import type {
+  ExecutionEnvironment,
+  ExecutionTrigger,
+  TraceSpan,
+  WorkflowState,
+} from '@/lib/logs/types'
 import {
   loadDeployedWorkflowState,
   loadWorkflowFromNormalizedTables,
@@ -80,7 +85,22 @@ export async function loadDeployedWorkflowStateForLogging(
   }
 }
 
-export function calculateCostSummary(traceSpans: any[]): {
+type CostTraceSpan = Pick<TraceSpan, 'cost' | 'model' | 'tokens'> & {
+  type?: TraceSpan['type']
+  children?: CostTraceSpan[]
+}
+
+type BillableTraceSpan = CostTraceSpan & { cost: NonNullable<TraceSpan['cost']> }
+
+function hasBillableCost(span: CostTraceSpan): span is BillableTraceSpan {
+  return span.cost !== undefined
+}
+
+function isModelBreakdownSpan(span: CostTraceSpan): boolean {
+  return span.type === 'model'
+}
+
+export function calculateCostSummary(traceSpans: CostTraceSpan[] | undefined): {
   totalCost: number
   totalInputCost: number
   totalOutputCost: number
@@ -122,17 +142,17 @@ export function calculateCostSummary(traceSpans: any[]): {
    * avoid double-counting. The parent cost is set by the provider response
    * (and is correctly zeroed by `executeProviderRequest` for BYOK calls);
    * model children only carry per-segment cost from the trace enrichers,
-   * which is unaware of BYOK status. Tool children do not carry cost and
-   * are unaffected.
+   * which is unaware of BYOK status. Non-model children are still visited
+   * so standalone nested costs remain billable.
    *
    * Spans without their own `cost` (e.g. parent workflow spans for
    * subworkflow blocks) still recurse so nested billable spans are counted.
    */
-  const collectCostSpans = (spans: any[]): any[] => {
-    const costSpans: any[] = []
+  const collectCostSpans = (spans: CostTraceSpan[]): BillableTraceSpan[] => {
+    const costSpans: BillableTraceSpan[] = []
 
     for (const span of spans) {
-      const hasOwnCost = !!span.cost
+      const hasOwnCost = hasBillableCost(span)
       if (hasOwnCost) {
         costSpans.push(span)
       }
@@ -142,7 +162,7 @@ export function calculateCostSummary(traceSpans: any[]): {
           // Parent already accounts for its model segments; only recurse into
           // non-model children (e.g. nested workflow spans) to find further
           // billable units.
-          const nonModelChildren = span.children.filter((c: any) => c?.type !== 'model')
+          const nonModelChildren = span.children.filter((child) => !isModelBreakdownSpan(child))
           costSpans.push(...collectCostSpans(nonModelChildren))
         } else {
           costSpans.push(...collectCostSpans(span.children))

apps/sim/lib/logs/execution/trace-spans/span-factory.ts

@@ -117,8 +117,13 @@ function enrichWithProviderMetadata(span: TraceSpan, log: ValidBlockLog): void {
   }
 
   if (output.cost) {
-    const { input, output: out, total } = output.cost
-    span.cost = { input, output: out, total }
+    const { input, output: out, total, toolCost } = output.cost
+    span.cost = {
+      input,
+      output: out,
+      total,
+      ...(typeof toolCost === 'number' && toolCost > 0 ? { toolCost } : {}),
+    }
   }
 
   if (output.tokens) {

apps/sim/lib/logs/execution/trace-spans/trace-spans.test.ts

@@ -2206,4 +2206,38 @@ describe('nested subflow grouping via parentIterations', () => {
     expect(secondModel.errorMessage).toBe('too many requests')
     expect(secondModel.status).toBe('error')
   })
+
+  it.concurrent('preserves parent toolCost on trace span cost', () => {
+    const result: ExecutionResult = {
+      success: true,
+      output: { content: 'done' },
+      logs: [
+        {
+          blockId: 'agent-tool-cost',
+          blockName: 'Agent With Tool Cost',
+          blockType: 'agent',
+          startedAt: '2024-01-01T10:00:00.000Z',
+          endedAt: '2024-01-01T10:00:02.000Z',
+          durationMs: 2000,
+          success: true,
+          input: {},
+          output: {
+            content: 'done',
+            model: 'gpt-4o',
+            tokens: { input: 100, output: 50, total: 150 },
+            cost: { input: 0.001, output: 0.002, toolCost: 0.015, total: 0.018 },
+          },
+        },
+      ],
+    }
+
+    const { traceSpans } = buildTraceSpans(result)
+
+    expect(traceSpans[0].cost).toEqual({
+      input: 0.001,
+      output: 0.002,
+      toolCost: 0.015,
+      total: 0.018,
+    })
+  })
 })

apps/sim/lib/logs/types.ts

@@ -223,6 +223,7 @@ export interface TraceSpan {
     input?: number
     output?: number
     total?: number
+    toolCost?: number
   }
   providerTiming?: ProviderTiming
   loopId?: string

apps/sim/providers/index.test.ts

@@ -149,9 +149,7 @@ describe('executeProviderRequest — BYOK regression', () => {
             startTime: 1777584457940,
             endTime: 1777584458000,
             duration: 60,
-            // Tool segments do not currently carry `cost` (tool cost lives on
-            // the parent's response.cost.toolCost), but if a future provider
-            // ever wrote a tool-segment cost we must NOT zero it.
+            cost: { total: 0.01 },
           },
         ],
       },
@@ -166,8 +164,7 @@ describe('executeProviderRequest — BYOK regression', () => {
     const [model, tool] = result.timing!.timeSegments!
     expect(model.cost?.total).toBe(0)
     expect(tool.type).toBe('tool')
-    // Helper only zeroes type==='model'; the tool segment is untouched.
-    expect((tool as { cost?: unknown }).cost).toBeUndefined()
+    expect(tool.cost?.total).toBe(0.01)
   })
 
   it('zeroes per-segment cost on streaming responses for BYOK callers', async () => {