fix(trace-spans): fix input/output token count in trace spans (#2168)

Author: waleedlatif1

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

@@ -80,7 +80,6 @@ export function calculateCostSummary(traceSpans: any[]): {
     }
   }
 
-  // Recursively collect all spans with cost information from the trace span tree
   const collectCostSpans = (spans: any[]): any[] => {
     const costSpans: any[] = []
 
@@ -119,12 +118,12 @@ export function calculateCostSummary(traceSpans: any[]): {
     totalCost += span.cost.total || 0
     totalInputCost += span.cost.input || 0
     totalOutputCost += span.cost.output || 0
-    // Tokens are at span.tokens, not span.cost.tokens
+    const promptTokens = span.tokens?.prompt ?? span.tokens?.input ?? 0
+    const completionTokens = span.tokens?.completion ?? span.tokens?.output ?? 0
     totalTokens += span.tokens?.total || 0
-    totalPromptTokens += span.tokens?.prompt || 0
-    totalCompletionTokens += span.tokens?.completion || 0
+    totalPromptTokens += promptTokens
+    totalCompletionTokens += completionTokens
 
-    // Aggregate model-specific costs - model is at span.model, not span.cost.model
     if (span.model) {
       const model = span.model
       if (!models[model]) {
@@ -138,8 +137,8 @@ export function calculateCostSummary(traceSpans: any[]): {
       models[model].input += span.cost.input || 0
       models[model].output += span.cost.output || 0
       models[model].total += span.cost.total || 0
-      models[model].tokens.prompt += span.tokens?.prompt || 0
-      models[model].tokens.completion += span.tokens?.completion || 0
+      models[model].tokens.prompt += promptTokens
+      models[model].tokens.completion += completionTokens
       models[model].tokens.total += span.tokens?.total || 0
     }
   }