From 3e83fb398cead855dcb950fad2b042edf9acb7d5 Mon Sep 17 00:00:00 2001
From: Waleed <walif6@gmail.com>
Date: Tue, 2 Dec 2025 20:36:16 -0800
Subject: [PATCH] fix(trace-spans): fix input/output token count in trace spans
 (#2168)

---
 apps/sim/lib/logs/execution/logging-factory.ts | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/apps/sim/lib/logs/execution/logging-factory.ts b/apps/sim/lib/logs/execution/logging-factory.ts
index 0cc5a243d..d70b750ff 100644
--- a/apps/sim/lib/logs/execution/logging-factory.ts
+++ b/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
     }
   }