Fix capped large token totals in TUI stats (#130)

Author: mike1858

src/contribution_cache/single_message.rs

@@ -118,11 +118,11 @@ impl PackedStatsDate {
     #[inline]
     pub fn to_tui_stats(self) -> TuiStats {
         TuiStats {
-            input_tokens: self.input_tokens(),
-            output_tokens: self.output_tokens(),
-            reasoning_tokens: self.reasoning_tokens(),
-            cached_tokens: self.cached_tokens(),
-            cost_cents: self.cost_cents() as u32,
+            input_tokens: self.input_tokens() as u64,
+            output_tokens: self.output_tokens() as u64,
+            reasoning_tokens: self.reasoning_tokens() as u64,
+            cached_tokens: self.cached_tokens() as u64,
+            cost_cents: self.cost_cents() as u64,
             tool_calls: self.tool_calls() as u32,
         }
     }

src/mcp/server.rs

@@ -349,7 +349,7 @@ impl SplitrailMcpServer {
                     .sum();
                 let total_tokens: u64 = filtered_stats
                     .iter()
-                    .map(|(_, ds)| (ds.stats.input_tokens as u64) + (ds.stats.output_tokens as u64))
+                    .map(|(_, ds)| ds.stats.input_tokens + ds.stats.output_tokens)
                     .sum();
                 let total_tool_calls: u32 = filtered_stats
                     .iter()

src/mcp/types.rs

@@ -158,9 +158,9 @@ impl DailySummary {
             ai_messages: ds.ai_messages,
             conversations: ds.conversations,
             total_cost: ds.stats.cost(),
-            input_tokens: ds.stats.input_tokens as u64,
-            output_tokens: ds.stats.output_tokens as u64,
-            cache_read_tokens: ds.stats.cached_tokens as u64,
+            input_tokens: ds.stats.input_tokens,
+            output_tokens: ds.stats.output_tokens,
+            cache_read_tokens: ds.stats.cached_tokens,
             tool_calls: ds.stats.tool_calls,
             files_read: file_ops.files_read,
             files_edited: file_ops.files_edited,

src/tui.rs

@@ -1021,15 +1021,15 @@ fn draw_daily_stats_table(
     // Find best values for highlighting
     // TODO: Let's refactor this.
 
-    let mut best_cost_cents: u32 = 0;
+    let mut best_cost_cents: u64 = 0;
     let mut best_cost_i = 0;
-    let mut best_cached_tokens: u32 = 0;
+    let mut best_cached_tokens: u64 = 0;
     let mut best_cached_tokens_i = 0;
-    let mut best_input_tokens: u32 = 0;
+    let mut best_input_tokens: u64 = 0;
     let mut best_input_tokens_i = 0;
-    let mut best_output_tokens: u32 = 0;
+    let mut best_output_tokens: u64 = 0;
     let mut best_output_tokens_i = 0;
-    let mut best_reasoning_tokens: u32 = 0;
+    let mut best_reasoning_tokens: u64 = 0;
     let mut best_reasoning_tokens_i = 0;
     let mut best_conversations = 0;
     let mut best_conversations_i = 0;
@@ -1089,11 +1089,11 @@ fn draw_daily_stats_table(
             continue;
         }
 
-        total_cost_cents += day_stats.stats.cost_cents as u64;
-        total_cached += day_stats.stats.cached_tokens as u64;
-        total_input += day_stats.stats.input_tokens as u64;
-        total_output += day_stats.stats.output_tokens as u64;
-        total_reasoning += day_stats.stats.reasoning_tokens as u64;
+        total_cost_cents += day_stats.stats.cost_cents;
+        total_cached += day_stats.stats.cached_tokens;
+        total_input += day_stats.stats.input_tokens;
+        total_output += day_stats.stats.output_tokens;
+        total_reasoning += day_stats.stats.reasoning_tokens;
         total_tool_calls += day_stats.stats.tool_calls as u64;
         total_conversations += day_stats.conversations as u64;
 
@@ -1607,11 +1607,11 @@ fn draw_session_stats_table(
             best_tool_calls_i = Some(idx);
         }
 
-        total_cost_cents += session.stats.cost_cents as u64;
-        total_input_tokens += session.stats.input_tokens as u64;
-        total_output_tokens += session.stats.output_tokens as u64;
-        total_cached_tokens += session.stats.cached_tokens as u64;
-        total_reasoning_tokens += session.stats.reasoning_tokens as u64;
+        total_cost_cents += session.stats.cost_cents;
+        total_input_tokens += session.stats.input_tokens;
+        total_output_tokens += session.stats.output_tokens;
+        total_cached_tokens += session.stats.cached_tokens;
+        total_reasoning_tokens += session.stats.reasoning_tokens;
         total_tool_calls += session.stats.tool_calls as u64;
 
         for &(model, _) in session.models.iter() {
@@ -1921,11 +1921,11 @@ fn draw_summary_stats(
                 continue;
             }
 
-            total_cost_cents += day_stats.stats.cost_cents as u64;
-            total_cached += day_stats.stats.cached_tokens as u64;
-            total_input += day_stats.stats.input_tokens as u64;
-            total_output += day_stats.stats.output_tokens as u64;
-            total_reasoning += day_stats.stats.reasoning_tokens as u64;
+            total_cost_cents += day_stats.stats.cost_cents;
+            total_cached += day_stats.stats.cached_tokens;
+            total_input += day_stats.stats.input_tokens;
+            total_output += day_stats.stats.output_tokens;
+            total_reasoning += day_stats.stats.reasoning_tokens;
             total_tool_calls += day_stats.stats.tool_calls as u64;
 
             // Collect unique days across all tools that have actual data

src/tui/logic.rs

@@ -8,12 +8,10 @@ pub use crate::types::SessionAggregate;
 /// Accumulate TUI-relevant stats from a full Stats into a TuiStats.
 /// Only copies the 6 fields displayed in the TUI.
 pub fn accumulate_tui_stats(dst: &mut TuiStats, src: &Stats) {
-    dst.input_tokens = dst.input_tokens.saturating_add(src.input_tokens as u32);
-    dst.output_tokens = dst.output_tokens.saturating_add(src.output_tokens as u32);
-    dst.reasoning_tokens = dst
-        .reasoning_tokens
-        .saturating_add(src.reasoning_tokens as u32);
-    dst.cached_tokens = dst.cached_tokens.saturating_add(src.cached_tokens as u32);
+    dst.input_tokens = dst.input_tokens.saturating_add(src.input_tokens);
+    dst.output_tokens = dst.output_tokens.saturating_add(src.output_tokens);
+    dst.reasoning_tokens = dst.reasoning_tokens.saturating_add(src.reasoning_tokens);
+    dst.cached_tokens = dst.cached_tokens.saturating_add(src.cached_tokens);
     dst.add_cost(src.cost);
     dst.tool_calls = dst.tool_calls.saturating_add(src.tool_calls);
 }

src/tui/tests.rs

@@ -325,6 +325,35 @@ fn test_large_tui_stats_accumulation() {
     assert!((dst.cost() - 10.0).abs() < 0.01);
 }
 
+#[test]
+fn test_tui_stats_accumulation_exceeds_u32_max() {
+    let mut dst = TuiStats::default();
+    let src = Stats {
+        input_tokens: u32::MAX as u64,
+        output_tokens: 1,
+        reasoning_tokens: 2,
+        cached_tokens: 3,
+        ..Stats::default()
+    };
+
+    accumulate_tui_stats(&mut dst, &src);
+    accumulate_tui_stats(
+        &mut dst,
+        &Stats {
+            input_tokens: 1,
+            output_tokens: u32::MAX as u64,
+            reasoning_tokens: u32::MAX as u64,
+            cached_tokens: u32::MAX as u64,
+            ..Stats::default()
+        },
+    );
+
+    assert_eq!(dst.input_tokens, u32::MAX as u64 + 1);
+    assert_eq!(dst.output_tokens, u32::MAX as u64 + 1);
+    assert_eq!(dst.reasoning_tokens, u32::MAX as u64 + 2);
+    assert_eq!(dst.cached_tokens, u32::MAX as u64 + 3);
+}
+
 // ============================================================================
 // COMPREHENSIVE DATA INTEGRITY TESTS
 // ============================================================================

src/types.rs

@@ -459,11 +459,11 @@ impl std::ops::SubAssign for Stats {
 #[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TuiStats {
-    pub input_tokens: u32,
-    pub output_tokens: u32,
-    pub reasoning_tokens: u32,
-    pub cached_tokens: u32,
-    pub cost_cents: u32, // Store as cents to avoid f32 precision issues
+    pub input_tokens: u64,
+    pub output_tokens: u64,
+    pub reasoning_tokens: u64,
+    pub cached_tokens: u64,
+    pub cost_cents: u64, // Store as cents to avoid f32 precision issues
     pub tool_calls: u32,
 }
 
@@ -477,26 +477,26 @@ impl TuiStats {
     /// Set cost from f64 dollars
     #[inline]
     pub fn set_cost(&mut self, dollars: f64) {
-        self.cost_cents = (dollars * 100.0).round() as u32;
+        self.cost_cents = (dollars * 100.0).round() as u64;
     }
 
     /// Add cost from f64 dollars
     #[inline]
     pub fn add_cost(&mut self, dollars: f64) {
         self.cost_cents = self
             .cost_cents
-            .saturating_add((dollars * 100.0).round() as u32);
+            .saturating_add((dollars * 100.0).round() as u64);
     }
 }
 
 impl From<&Stats> for TuiStats {
     fn from(s: &Stats) -> Self {
         TuiStats {
-            input_tokens: s.input_tokens as u32,
-            output_tokens: s.output_tokens as u32,
-            reasoning_tokens: s.reasoning_tokens as u32,
-            cached_tokens: s.cached_tokens as u32,
-            cost_cents: (s.cost * 100.0).round() as u32,
+            input_tokens: s.input_tokens,
+            output_tokens: s.output_tokens,
+            reasoning_tokens: s.reasoning_tokens,
+            cached_tokens: s.cached_tokens,
+            cost_cents: (s.cost * 100.0).round() as u64,
             tool_calls: s.tool_calls,
         }
     }

src/utils.rs

@@ -223,19 +223,19 @@ pub fn aggregate_by_date(entries: &[ConversationMessage]) -> BTreeMap<String, Da
                 daily_stats_entry.stats.input_tokens = daily_stats_entry
                     .stats
                     .input_tokens
-                    .saturating_add(entry.stats.input_tokens as u32);
+                    .saturating_add(entry.stats.input_tokens);
                 daily_stats_entry.stats.output_tokens = daily_stats_entry
                     .stats
                     .output_tokens
-                    .saturating_add(entry.stats.output_tokens as u32);
+                    .saturating_add(entry.stats.output_tokens);
                 daily_stats_entry.stats.reasoning_tokens = daily_stats_entry
                     .stats
                     .reasoning_tokens
-                    .saturating_add(entry.stats.reasoning_tokens as u32);
+                    .saturating_add(entry.stats.reasoning_tokens);
                 daily_stats_entry.stats.cached_tokens = daily_stats_entry
                     .stats
                     .cached_tokens
-                    .saturating_add(entry.stats.cached_tokens as u32);
+                    .saturating_add(entry.stats.cached_tokens);
                 daily_stats_entry.stats.tool_calls = daily_stats_entry
                     .stats
                     .tool_calls