From e8e9f5c79c028c3297ab3dcf6bafcd00e547e0c1 Mon Sep 17 00:00:00 2001
From: Naomi Carrigan <commits@nhcarrigan.com>
Date: Fri, 6 Feb 2026 10:13:19 -0800
Subject: [PATCH] fix: more accurate cost tracking hopefully

---
 src-tauri/src/wsl_bridge.rs | 102 ++++++++++++++++++++----------------
 1 file changed, 57 insertions(+), 45 deletions(-)

diff --git a/src-tauri/src/wsl_bridge.rs b/src-tauri/src/wsl_bridge.rs
index 1fa7102..6d0c54e 100644
--- a/src-tauri/src/wsl_bridge.rs
+++ b/src-tauri/src/wsl_bridge.rs
@@ -717,60 +717,72 @@ fn process_json_line(
 
             // Only update stats if we have usage information
             if let Some(usage) = &message.usage {
-                if let Some(model) = &message.model {
-                    // Calculate cost for historical tracking (including cache tokens)
-                    let cost_usd = calculate_cost(
+                // Get model from message, or fall back to last known model from stats
+                let model = message.model.clone().or_else(|| {
+                    let stats_guard = stats.read();
+                    stats_guard.model.clone()
+                }).unwrap_or_else(|| {
+                    println!("[WARNING] No model info available for cost calculation, using default");
+                    "claude-sonnet-4-5-20250929".to_string()
+                });
+
+                // Calculate cost for historical tracking (including cache tokens)
+                let cost_usd = calculate_cost(
+                    usage.input_tokens,
+                    usage.output_tokens,
+                    &model,
+                    usage.cache_creation_input_tokens,
+                    usage.cache_read_input_tokens,
+                );
+
+                println!("Assistant message tokens - input: {}, output: {}, cache_creation: {:?}, cache_read: {:?}, cost: ${:.4}",
+                    usage.input_tokens,
+                    usage.output_tokens,
+                    usage.cache_creation_input_tokens,
+                    usage.cache_read_input_tokens,
+                    cost_usd
+                );
+
+                // Store cost for later use in output events
+                message_cost = Some(MessageCost {
+                    input_tokens: usage.input_tokens,
+                    output_tokens: usage.output_tokens,
+                    cost_usd,
+                });
+
+                // Batch all stats updates in a single write lock
+                {
+                    let mut stats_guard = stats.write();
+                    stats_guard.increment_messages();
+                    stats_guard.add_usage(
                         usage.input_tokens,
                         usage.output_tokens,
-                        model,
+                        &model,
                         usage.cache_creation_input_tokens,
                         usage.cache_read_input_tokens,
                     );
+                    stats_guard.get_session_duration();
 
-                    // Store cost for later use in output events
-                    message_cost = Some(MessageCost {
-                        input_tokens: usage.input_tokens,
-                        output_tokens: usage.output_tokens,
-                        cost_usd,
-                    });
-
-                    // Batch all stats updates in a single write lock
-                    {
-                        let mut stats_guard = stats.write();
-                        stats_guard.increment_messages();
-                        stats_guard.add_usage(
-                            usage.input_tokens,
-                            usage.output_tokens,
-                            model,
-                            usage.cache_creation_input_tokens,
-                            usage.cache_read_input_tokens,
-                        );
-                        stats_guard.get_session_duration();
-
-                        // Attribute tokens to tools if any tools were used in this message
-                        if !tools_in_message.is_empty() {
-                            let per_tool_input = usage.input_tokens / tools_in_message.len() as u64;
-                            let per_tool_output = usage.output_tokens / tools_in_message.len() as u64;
-                            for tool in &tools_in_message {
-                                stats_guard.add_tool_tokens(tool, per_tool_input, per_tool_output);
-                            }
+                    // Attribute tokens to tools if any tools were used in this message
+                    if !tools_in_message.is_empty() {
+                        let per_tool_input = usage.input_tokens / tools_in_message.len() as u64;
+                        let per_tool_output = usage.output_tokens / tools_in_message.len() as u64;
+                        for tool in &tools_in_message {
+                            stats_guard.add_tool_tokens(tool, per_tool_input, per_tool_output);
                         }
                     }
-
-                    // Record to historical cost tracking
-                    let app_clone = app.clone();
-                    let input = usage.input_tokens;
-                    let output = usage.output_tokens;
-                    tauri::async_runtime::spawn(async move {
-                        record_cost(&app_clone, input, output, cost_usd).await;
-                    });
-
-                    // Don't emit here - we'll emit on Result message instead
-                    // This reduces the frequency of updates
-                } else {
-                    // Just increment message count if no usage info
-                    stats.write().increment_messages();
                 }
+
+                // Record to historical cost tracking
+                let app_clone = app.clone();
+                let input = usage.input_tokens;
+                let output = usage.output_tokens;
+                tauri::async_runtime::spawn(async move {
+                    record_cost(&app_clone, input, output, cost_usd).await;
+                });
+
+                // Don't emit here - we'll emit on Result message instead
+                // This reduces the frequency of updates
             } else {
                 // Just increment message count if no usage info
                 stats.write().increment_messages();