Fix Codex token count saturation

rust/src/codex_costs.rs

@@ -60,7 +60,7 @@ pub(crate) fn add_codex_records_to_summary(
 
 /// Merge billable records into a day→model→`[input,cached,output]` map.
 pub(crate) fn merge_codex_records_into_days(
-    days: &mut std::collections::HashMap<String, std::collections::HashMap<String, Vec<i32>>>,
+    days: &mut std::collections::HashMap<String, std::collections::HashMap<String, Vec<i64>>>,
     records: &[CodexUsageRecord],
 ) {
     for record in records {
@@ -77,7 +77,7 @@ pub(crate) fn merge_codex_records_into_days(
 pub(crate) fn add_codex_packed_tokens_to_summary(
     summary: &mut CostSummary,
     model: &str,
-    packed: &[i32],
+    packed: &[i64],
     pricing_day: Option<NaiveDate>,
 ) -> Option<f64> {
     let input = packed.first().copied().unwrap_or(0);
@@ -99,7 +99,7 @@ pub(crate) fn add_codex_packed_tokens_to_summary(
 /// Returns `(session_cost, has_tokens)` — caller adds cost to `total_cost_usd`.
 pub(crate) fn add_codex_days_map_to_summary(
     summary: &mut CostSummary,
-    days: &std::collections::HashMap<String, std::collections::HashMap<String, Vec<i32>>>,
+    days: &std::collections::HashMap<String, std::collections::HashMap<String, Vec<i64>>>,
     range: &CostUsageDayRange,
 ) -> (f64, bool) {
     let mut total_cost = 0.0;
@@ -153,12 +153,12 @@ struct CodexTokenCounts {
 }
 
 impl CodexTokenCounts {
-    fn from_values(input: i32, cached: i32, output: i32) -> Self {
-        let input = input.max(0) as u64;
+    fn from_values(input: i64, cached: i64, output: i64) -> Self {
+        let input = u64::try_from(input.max(0)).unwrap_or(0);
         Self {
             input,
-            cached: (cached.max(0) as u64).min(input),
-            output: output.max(0) as u64,
+            cached: u64::try_from(cached.max(0)).unwrap_or(0).min(input),
+            output: u64::try_from(output.max(0)).unwrap_or(0),
             reasoning: None,
         }
     }
@@ -404,35 +404,11 @@ fn codex_cost_usd_for_day(
 
     let normalized = CostUsagePricing::normalize_codex_model(model);
     if normalized.contains("fast") || normalized.contains("priority") {
-        // Fast pricing takes i32 token counts; usage-record counts fit far
-        // below i32::MAX, and the callee re-checks the long-context threshold
-        // against the original u64 magnitude.
-        #[allow(
-            clippy::cast_possible_truncation,
-            reason = "token counts from usage records fit i32"
-        )]
-        #[allow(
-            clippy::cast_possible_wrap,
-            reason = "token counts are non-negative; wrapping is impossible"
-        )]
         let fast = pricing_day
             .and_then(|day| {
-                CostUsagePricing::codex_fast_cost_usd_at_date(
-                    model,
-                    input as i32,
-                    cached as i32,
-                    output as i32,
-                    day,
-                )
+                CostUsagePricing::codex_fast_cost_usd_at_date(model, input, cached, output, day)
             })
-            .or_else(|| {
-                CostUsagePricing::codex_fast_cost_usd(
-                    model,
-                    input as i32,
-                    cached as i32,
-                    output as i32,
-                )
-            });
+            .or_else(|| CostUsagePricing::codex_fast_cost_usd(model, input, cached, output));
         if let Some(cost) = fast {
             return cost;
         }

rust/src/core/cost_cache_budget.rs

@@ -71,6 +71,12 @@ fn touches_window(entry: &CostUsageFileUsage, since_key: &str, until_key: &str)
         .any(|day| CostUsageDayRange::is_in_range(day, since_key, until_key))
 }
 
+/// Conservative JSON byte width of one packed token value. The totals table was
+/// widened from `i32` to `i64`, so a single value can serialize to up to 20
+/// bytes (19 digits plus a sign); budgeting 10 bytes would under-estimate the
+/// artifact and let pruning admit a cache that is already over budget.
+const PACKED_VALUE_BYTES: usize = 20;
+
 /// Cheap per-entry byte estimate (conservative overhead) so the save path can
 /// decide whether to prune *before* materializing the encoded document. Mirrors
 /// upstream `estimatedCodexCacheBytes`'s per-entry shape; it deliberately
@@ -84,7 +90,7 @@ fn estimated_entry_bytes(entry: &CostUsageFileUsage) -> usize {
     for (day, models) in &entry.days {
         bytes += day.len() + 32;
         for (model, packed) in models {
-            bytes += model.len() + 40 + packed.len() * 10;
+            bytes += model.len() + 40 + packed.len() * PACKED_VALUE_BYTES;
         }
     }
     bytes
@@ -93,7 +99,7 @@ fn estimated_entry_bytes(entry: &CostUsageFileUsage) -> usize {
 /// Conservative estimate of the encoded artifact size.
 pub fn estimated_cache_bytes(
     files: &HashMap<String, CostUsageFileUsage>,
-    days: &HashMap<String, HashMap<String, Vec<i32>>>,
+    days: &HashMap<String, HashMap<String, Vec<i64>>>,
 ) -> usize {
     let mut bytes = 4096;
     bytes += files.len() * 160;
@@ -103,7 +109,7 @@ pub fn estimated_cache_bytes(
     for (day, models) in days {
         bytes += day.len() + 32;
         for (model, packed) in models {
-            bytes += model.len() + 40 + packed.len() * 10;
+            bytes += model.len() + 40 + packed.len() * PACKED_VALUE_BYTES;
         }
     }
     bytes
@@ -121,7 +127,7 @@ pub fn estimated_cache_bytes(
 /// shape (no fork lineages, no discovery/lookback state).
 pub fn prune_out_of_window_for_budget(
     files: &mut HashMap<String, CostUsageFileUsage>,
-    days: &mut HashMap<String, HashMap<String, Vec<i32>>>,
+    days: &mut HashMap<String, HashMap<String, Vec<i64>>>,
     scan_since_key: Option<&str>,
     scan_until_key: Option<&str>,
     force: bool,
@@ -171,7 +177,7 @@ pub fn prune_out_of_window_for_budget(
 /// cache shape.
 pub fn trim_in_window_for_budget(
     files: &mut HashMap<String, CostUsageFileUsage>,
-    days: &mut HashMap<String, HashMap<String, Vec<i32>>>,
+    days: &mut HashMap<String, HashMap<String, Vec<i64>>>,
     scan_since_key: Option<&str>,
     scan_until_key: Option<&str>,
     max_bytes: usize,
@@ -248,8 +254,8 @@ pub fn trim_in_window_for_budget(
 /// of the scanner's `rebuild_cache_days` accumulation), so pruned entries do
 /// not inflate totals.
 fn subtract_entry_days(
-    days: &mut HashMap<String, HashMap<String, Vec<i32>>>,
-    entry_days: &HashMap<String, HashMap<String, Vec<i32>>>,
+    days: &mut HashMap<String, HashMap<String, Vec<i64>>>,
+    entry_days: &HashMap<String, HashMap<String, Vec<i64>>>,
 ) {
     let mut empty_days = Vec::new();
     for (day, models) in entry_days {
@@ -261,7 +267,7 @@ fn subtract_entry_days(
             let Some(dest) = day_entry.get_mut(model) else {
                 continue;
             };
-            for (i, value) in packed.iter().take(3).enumerate() {
+            for (i, value) in packed.iter().take(4).enumerate() {
                 if i < dest.len() {
                     dest[i] = dest[i].saturating_sub(*value);
                 }
@@ -315,7 +321,7 @@ mod tests {
     use super::*;
 
     fn entry(days: &[&str], parsed: Option<i64>, size: i64) -> CostUsageFileUsage {
-        let mut day_map: HashMap<String, HashMap<String, Vec<i32>>> = HashMap::new();
+        let mut day_map: HashMap<String, HashMap<String, Vec<i64>>> = HashMap::new();
         for day in days {
             day_map.insert(
                 (*day).to_string(),
@@ -342,12 +348,12 @@ mod tests {
 
     type TestCache = (
         HashMap<String, CostUsageFileUsage>,
-        HashMap<String, HashMap<String, Vec<i32>>>,
+        HashMap<String, HashMap<String, Vec<i64>>>,
     );
 
     fn cache(files: &[(&str, CostUsageFileUsage)]) -> TestCache {
         let mut file_map = HashMap::new();
-        let mut days: HashMap<String, HashMap<String, Vec<i32>>> = HashMap::new();
+        let mut days: HashMap<String, HashMap<String, Vec<i64>>> = HashMap::new();
         for (key, entry) in files {
             for (day, models) in &entry.days {
                 let day_entry = days.entry(day.clone()).or_default();
@@ -367,6 +373,24 @@ mod tests {
         (file_map, days)
     }
 
+    #[test]
+    fn subtract_entry_days_removes_reasoning_slot() {
+        let day = "2026-01-10".to_string();
+        let model = "gpt-5.6-sol".to_string();
+        let mut days = HashMap::from([(
+            day.clone(),
+            HashMap::from([(model.clone(), vec![30, 12, 9, 6])]),
+        )]);
+        let entry_days = HashMap::from([(
+            day.clone(),
+            HashMap::from([(model.clone(), vec![10, 4, 3, 2])]),
+        )]);
+
+        subtract_entry_days(&mut days, &entry_days);
+
+        assert_eq!(days[&day][&model], vec![20, 8, 6, 4]);
+    }
+
     #[test]
     fn out_of_window_entries_are_pruned_and_days_subtracted() {
         let (mut files, mut days) = cache(&[
@@ -526,6 +550,45 @@ mod tests {
         assert!(files.contains_key("c"), "newest kept");
     }
 
+    #[test]
+    fn estimated_entry_bytes_reserves_twenty_bytes_per_packed_i64_value() {
+        let day = "2026-01-09";
+        let model = "gpt-5.6-sol";
+        let three = entry(&[day], None, 100);
+        let mut four = entry(&[day], None, 100);
+        four.days
+            .get_mut(day)
+            .unwrap()
+            .insert(model.to_string(), vec![1, 2, 3, 4]);
+
+        assert_eq!(
+            estimated_entry_bytes(&four) - estimated_entry_bytes(&three),
+            20,
+            "each packed i64 slot must reserve its full 20-byte JSON width"
+        );
+    }
+
+    #[test]
+    fn estimated_cache_bytes_reserves_twenty_bytes_per_packed_i64_value_in_days() {
+        let day = "2026-01-09";
+        let model = "gpt-5.6-sol";
+        let files: HashMap<String, CostUsageFileUsage> = HashMap::new();
+        let three = HashMap::from([(
+            day.to_string(),
+            HashMap::from([(model.to_string(), vec![1_i64, 2, 3])]),
+        )]);
+        let four = HashMap::from([(
+            day.to_string(),
+            HashMap::from([(model.to_string(), vec![1_i64, 2, 3, 4])]),
+        )]);
+
+        assert_eq!(
+            estimated_cache_bytes(&files, &four) - estimated_cache_bytes(&files, &three),
+            20,
+            "each packed i64 slot in the aggregate day map must reserve 20 bytes"
+        );
+    }
+
     #[test]
     fn is_unpriced_codex_routing_model_flags_auto_review_and_unattributed() {
         assert!(is_unpriced_codex_routing_model("codex-auto-review"));

rust/src/core/cost_pricing.rs

@@ -703,22 +703,17 @@ impl CostUsagePricing {
     /// suffixes), then applies the Fast multiplier. Returns `None` when the
     /// model has no Fast lane or when a model without Astra's published
     /// long-context Fast rates exceeds the 272 000 threshold.
-    pub fn codex_fast_cost_usd(model: &str, input: i32, cached: i32, output: i32) -> Option<f64> {
+    pub fn codex_fast_cost_usd(model: &str, input: u64, cached: u64, output: u64) -> Option<f64> {
         let multiplier = Self::codex_api_fast_multiplier(model)?;
         // Older models do not offer Fast for long-context requests. Astra
         // publishes a Fast rate for the same whole-request long-context tier.
-        if (input.max(0) as u64) > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD
+        if input > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD
             && !codex_pricing::codex_fast_allows_long_context(model)
         {
             return None;
         }
         let base = Self::codex_fast_base_model(model);
-        let base_cost = Self::codex_cost_usd(
-            &base,
-            input.max(0) as u64,
-            cached.max(0) as u64,
-            output.max(0) as u64,
-        )?;
+        let base_cost = Self::codex_cost_usd(&base, input, cached, output)?;
         Some(base_cost * multiplier)
     }
 
@@ -782,25 +777,19 @@ impl CostUsagePricing {
 
     pub fn codex_fast_cost_usd_at_date(
         model: &str,
-        input: i32,
-        cached: i32,
-        output: i32,
+        input: u64,
+        cached: u64,
+        output: u64,
         pricing_date: NaiveDate,
     ) -> Option<f64> {
         let multiplier = Self::codex_api_fast_multiplier(model)?;
-        if (input.max(0) as u64) > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD
+        if input > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD
             && !codex_pricing::codex_fast_allows_long_context(model)
         {
             return None;
         }
         let base = Self::codex_fast_base_model(model);
-        let base_cost = Self::codex_cost_usd_at_date(
-            &base,
-            input.max(0) as u64,
-            cached.max(0) as u64,
-            output.max(0) as u64,
-            pricing_date,
-        )?;
+        let base_cost = Self::codex_cost_usd_at_date(&base, input, cached, output, pricing_date)?;
         Some(base_cost * multiplier)
     }