From c6efa0760e743b4fbc6eb63f8c7f895c3066cf5f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 00:34:40 +0700 Subject: [PATCH 01/62] Harden provider numeric boundaries --- rust/src/providers/amp/mod.rs | 18 ++++ rust/src/providers/amp/subscription.rs | 52 ++++++++-- rust/src/providers/kilo/mod.rs | 66 +++++++++--- rust/src/providers/longcat/mod.rs | 138 ++++++++++++++++++++----- 4 files changed, 231 insertions(+), 43 deletions(-) diff --git a/rust/src/providers/amp/mod.rs b/rust/src/providers/amp/mod.rs index d248d22d38..591112e31c 100755 --- a/rust/src/providers/amp/mod.rs +++ b/rust/src/providers/amp/mod.rs @@ -501,4 +501,22 @@ period 2026-09-13 to 2026-10-13, resets upon renewal in 27 days"; assert!(snapshot.secondary.is_none()); assert_eq!(snapshot.primary.used_percent, 85.0); } + + #[test] + fn tier_keeps_explicit_period_when_renewal_count_overflows() { + let now = Utc.with_ymd_and_hms(2026, 9, 16, 12, 0, 0).unwrap(); + let text = "Amp Example Tier: agent usage $10 of $20 remaining - \ +period 2026-09-13 to 2026-10-13, resets upon renewal in 999999999999999999999999999999 days"; + let sub = parse_amp_subscription_usage(text, now).expect("tier"); + assert_eq!( + sub.resets_at(), + Some(Utc.with_ymd_and_hms(2026, 10, 13, 0, 0, 0).unwrap()) + ); + assert_eq!( + sub.reset_description, + "renews in 999999999999999999999999999999 days" + ); + let snapshot = usage_snapshot_from_amp_display_text(text, now).expect("snapshot"); + assert_eq!(snapshot.primary.used_percent, 50.0); + } } diff --git a/rust/src/providers/amp/subscription.rs b/rust/src/providers/amp/subscription.rs index 449b98199a..d471cea1f8 100644 --- a/rust/src/providers/amp/subscription.rs +++ b/rust/src/providers/amp/subscription.rs @@ -197,10 +197,21 @@ pub(super) fn parse_amp_subscription_usage( let agent_remaining = parse_amp_number(caps.get(2)?.as_str())?; let agent_limit = parse_amp_number(caps.get(3)?.as_str())?; let details = caps.get(4)?.as_str(); - let renewal_value: i64 = caps.get(5)?.as_str().replace(',', "").parse().ok()?; + let renewal_text = caps.get(5)?.as_str().replace(',', ""); + let renewal_value = renewal_text.parse::().ok(); let renewal_unit = caps.get(6)?.as_str().to_ascii_lowercase(); - let reset_description = amp_renewal_description(renewal_value, &renewal_unit); - let (period_start, resets_at) = parse_amp_tier_period(details).unzip(); + let period = parse_amp_tier_period(details); + let has_period_text = details.to_ascii_lowercase().contains("period "); + let resets_at = period.map(|(_, end)| end).or_else(|| { + (!has_period_text) + .then(|| { + renewal_value + .and_then(|value| subscription_reset_date(value, &renewal_unit, now)) + }) + .flatten() + }); + let reset_description = amp_renewal_description_text(&renewal_text, &renewal_unit); + let period_start = period.map(|(start, _)| start); let orb = orb_re.captures(details).and_then(|orb_caps| { let remaining = parse_amp_number(orb_caps.get(1)?.as_str())?; let limit = parse_amp_number(orb_caps.get(2)?.as_str())?; @@ -242,11 +253,7 @@ pub(super) fn parse_amp_subscription_usage( continue; } let unit = caps.get(5)?.as_str().to_ascii_lowercase(); - let resets_at = if unit.starts_with("month") { - add_calendar_months(now, renewal_value)? - } else { - now + chrono::Duration::days(renewal_value) - }; + let resets_at = subscription_reset_date(renewal_value, &unit, now)?; let reset_description = amp_renewal_description(renewal_value, &unit); return Some(AmpSubscriptionUsage { plan: plan.to_string(), @@ -274,6 +281,35 @@ fn amp_renewal_description(value: i64, unit: &str) -> String { } } +fn amp_renewal_description_text(value: &str, unit: &str) -> String { + let singular = if unit.starts_with("month") { + "month" + } else { + "day" + }; + if value == "1" { + format!("renews in 1 {singular}") + } else { + format!("renews in {value} {singular}s") + } +} + +fn subscription_reset_date( + value: i64, + unit: &str, + now: chrono::DateTime, +) -> Option> { + if value < 0 { + return None; + } + if unit.starts_with("month") { + add_calendar_months(now, value) + } else { + let seconds = value.checked_mul(24 * 60 * 60)?; + now.checked_add_signed(chrono::Duration::seconds(seconds)) + } +} + fn parse_amp_tier_period( text: &str, ) -> Option<(chrono::DateTime, chrono::DateTime)> { diff --git a/rust/src/providers/kilo/mod.rs b/rust/src/providers/kilo/mod.rs index 5aa8bd330a..4b10d32d40 100644 --- a/rust/src/providers/kilo/mod.rs +++ b/rust/src/providers/kilo/mod.rs @@ -111,8 +111,8 @@ impl KiloProvider { has_blocks = !arr.is_empty(); for block in arr { if let Ok(b) = serde_json::from_value::(block.clone()) { - total += b.amount_m_usd.unwrap_or(0.0); - remaining += b.balance_m_usd.unwrap_or(0.0); + add_finite(&mut total, b.amount_m_usd); + add_finite(&mut remaining, b.balance_m_usd); } } } @@ -121,23 +121,27 @@ impl KiloProvider { && let Some(balance_m_usd) = payload.get("totalBalance_mUsd").and_then(|v| v.as_f64()) { - total = balance_m_usd; - remaining = balance_m_usd; + total = finite_or_zero(Some(balance_m_usd)); + remaining = total; } } let total_usd = total / 1_000_000.0; let remaining_usd = remaining / 1_000_000.0; - let used_usd = (total_usd - remaining_usd).max(0.0); - let percent = if total_usd > 0.0 { - ((used_usd / total_usd) * 100.0).clamp(0.0, 100.0) + let primary = if total_usd.is_finite() && remaining_usd.is_finite() { + let used_usd = (total_usd - remaining_usd).max(0.0); + let percent = if total_usd > 0.0 { + ((used_usd / total_usd) * 100.0).clamp(0.0, 100.0) + } else { + 0.0 + }; + let mut window = RateWindow::new(percent); + window.reset_description = Some(format!("${used_usd:.2}/${total_usd:.2}")); + window } else { - 0.0 + RateWindow::informational("Credit usage unavailable") }; - let mut primary = RateWindow::new(percent); - primary.reset_description = Some(format!("${:.2}/${:.2}", used_usd, total_usd)); - let mut snap = UsageSnapshot::new(primary); // --- Kilo Pass (secondary window) --- @@ -145,17 +149,20 @@ impl KiloProvider { let usage = pass .get("currentPeriodUsageUsd") .and_then(|v| v.as_f64()) + .filter(|value| value.is_finite()) .unwrap_or(0.0); let base = pass .get("currentPeriodBaseCreditsUsd") .and_then(|v| v.as_f64()) + .filter(|value| value.is_finite()) .unwrap_or(0.0); let bonus = pass .get("currentPeriodBonusCreditsUsd") .and_then(|v| v.as_f64()) + .filter(|value| value.is_finite()) .unwrap_or(0.0); let pass_total = base + bonus; - if pass_total > 0.0 { + if pass_total.is_finite() && pass_total > 0.0 { let pass_pct = ((usage / pass_total) * 100.0).clamp(0.0, 100.0); let mut secondary = RateWindow::new(pass_pct); secondary.reset_description = Some(format!("${:.2}/${:.2}", usage, pass_total)); @@ -221,6 +228,17 @@ impl KiloProvider { } } +fn finite_or_zero(value: Option) -> f64 { + value.filter(|value| value.is_finite()).unwrap_or(0.0) +} + +fn add_finite(total: &mut f64, value: Option) { + let Some(value) = value.filter(|value| value.is_finite()) else { + return; + }; + *total += value; +} + fn direct_kilo_api_key(api_key: Option<&str>) -> Option { api_key.filter(|key| !key.is_empty()).map(str::to_string) } @@ -373,4 +391,28 @@ mod tests { ); assert!(snap.secondary.is_none()); } + + #[test] + fn makes_overflowed_credit_totals_unavailable() { + let credit_blocks = serde_json::json!({ + "creditBlocks": [ + { "amount_mUsd": 1e308, "balance_mUsd": 1e308 }, + { "amount_mUsd": 1e308, "balance_mUsd": 1e308 } + ] + }); + let snap = KiloProvider::build_snapshot(Some(&credit_blocks), None).unwrap(); + assert!(snap.primary.is_informational); + assert!(snap.primary.used_percent.is_finite()); + } + + #[test] + fn omits_pass_window_when_usage_arithmetic_overflows() { + let pass = serde_json::json!({ + "currentPeriodUsageUsd": 1e308, + "currentPeriodBaseCreditsUsd": 1e308, + "currentPeriodBonusCreditsUsd": 1e308 + }); + let snap = KiloProvider::build_snapshot(None, Some(&pass)).unwrap(); + assert!(snap.secondary.is_none()); + } } diff --git a/rust/src/providers/longcat/mod.rs b/rust/src/providers/longcat/mod.rs index f3ef3f6ce4..1bbd777608 100644 --- a/rust/src/providers/longcat/mod.rs +++ b/rust/src/providers/longcat/mod.rs @@ -125,7 +125,7 @@ impl Provider for LongCatProvider { }; let account = self.get_json(USER_CURRENT, &cookie).await?; // Meituan-style envelope may return HTTP 200 with business 401. - if let Some(code) = envelope_code(&account) + if let Some(code) = envelope_code(&account)? && (code == 401 || code == 403) { return Err(ProviderError::AuthRequired); @@ -181,11 +181,16 @@ fn normalize_cookie_header(raw: &str) -> Option { (!header.is_empty()).then_some(header) } -fn envelope_code(value: &Value) -> Option { +fn envelope_code(value: &Value) -> Result, ProviderError> { value .get("code") - .and_then(|c| c.as_i64()) - .or_else(|| value.get("status").and_then(|c| c.as_i64())) + .or_else(|| value.get("status")) + .map(|raw| { + json_integer(raw).ok_or_else(|| { + ProviderError::Parse("LongCat response code was not a valid integer".into()) + }) + }) + .transpose() } fn envelope_data(value: &Value) -> &Value { @@ -197,10 +202,54 @@ fn json_f64(value: &Value, key: &str) -> Option { } fn json_number(value: &Value) -> Option { - value + let number = value .as_f64() .or_else(|| value.as_i64().map(|number| number as f64)) + .or_else(|| value.as_str()?.trim().parse().ok())?; + number.is_finite().then_some(number) +} + +fn json_integer(value: &Value) -> Option { + value + .as_i64() + .or_else(|| value.as_f64().and_then(truncate_to_i64)) .or_else(|| value.as_str()?.trim().parse().ok()) + .or_else(|| { + value + .as_str()? + .trim() + .parse::() + .ok() + .and_then(truncate_to_i64) + }) +} + +fn truncate_to_i64(value: f64) -> Option { + if !value.is_finite() { + return None; + } + let truncated = value.trunc(); + // `i64::MAX as f64` rounds up to 2^63, so keep that boundary exclusive. + if truncated < i64::MIN as f64 || truncated >= i64::MAX as f64 { + return None; + } + #[expect( + clippy::cast_possible_truncation, + reason = "the finite value was range-checked before conversion" + )] + Some(truncated as i64) +} + +fn whole_number(value: f64) -> Option { + if !value.is_finite() { + return None; + } + let normalized = if value.trunc() == 0.0 { + 0.0 + } else { + value.trunc() + }; + Some(format!("{normalized:.0}")) } fn json_str(value: &Value, key: &str) -> Option { @@ -259,19 +308,16 @@ fn build_snapshot( )); }; - let primary = if total > 0.0 { + let primary = if total.is_finite() && total > 0.0 && used.is_finite() { let mut w = RateWindow::new(((used / total) * 100.0).clamp(0.0, 100.0)); - // Display-only rendering of token counts; values beyond i64 are - // unrealistic quota sizes and would only affect this label. - #[allow( - clippy::cast_possible_truncation, - reason = "display-only quota label; token counts beyond i64 are unrealistic" - )] - let desc = format!("{}/{}", used as i64, total as i64); - w.reset_description = Some(desc); + if let (Some(used_text), Some(total_text)) = (whole_number(used), whole_number(total)) { + w.reset_description = Some(format!("{used_text}/{total_text}")); + } w - } else { + } else if total.is_finite() && total <= 0.0 { RateWindow::informational("No token quota") + } else { + RateWindow::informational("Token quota unavailable") }; let account_name = json_str(account_data, "name") @@ -286,19 +332,20 @@ fn build_snapshot( if let Some(fuel_raw) = fuel_raw { let fuel_data = envelope_data(fuel_raw); if let Some((total_fuel, remaining_fuel, expiry)) = parse_fuel(fuel_data) + && total_fuel.is_finite() + && remaining_fuel.is_finite() && total_fuel > 0.0 { let used_fuel = (total_fuel - remaining_fuel).max(0.0); let mut secondary = RateWindow::new(((used_fuel / total_fuel) * 100.0).clamp(0.0, 100.0)); secondary.resets_at = expiry; - // Same display-only label for fuel-pack counts. - #[allow( - clippy::cast_possible_truncation, - reason = "display-only fuel label; fuel counts beyond i64 are unrealistic" - )] - let fuel_desc = format!("Fuel pack: {}/{}", remaining_fuel as i64, total_fuel as i64); - secondary.reset_description = Some(fuel_desc); + if let (Some(remaining_text), Some(total_text)) = + (whole_number(remaining_fuel), whole_number(total_fuel)) + { + secondary.reset_description = + Some(format!("Fuel pack: {remaining_text}/{total_text}")); + } snap = snap.with_secondary(secondary); } } @@ -374,7 +421,7 @@ fn parse_fuel_timestamp(value: &Value) -> Option> { } else { number * 1000.0 }; - if millis <= 1_000_000_000_000.0 || millis > i64::MAX as f64 { + if millis <= 1_000_000_000_000.0 || millis >= i64::MAX as f64 { return None; } #[expect( @@ -494,4 +541,49 @@ mod tests { Some("a=1; b=2") ); } + + #[test] + fn rejects_unrepresentable_response_codes() { + let error = envelope_code(&json!({ "code": "Infinity" })).unwrap_err(); + assert!(matches!(error, ProviderError::Parse(_))); + assert_eq!(envelope_code(&json!({ "code": 200.9 })).unwrap(), Some(200)); + } + + #[test] + fn formats_large_counts_without_i64_saturation() { + let account = json!({ "code": 0, "data": { "name": "cat" } }); + let usage = json!({ + "code": 0, + "data": { + "usage": { + "totalToken": 2e20, + "availableToken": 1e20 + } + } + }); + let snapshot = build_snapshot(&account, None, Some(&usage), None).unwrap(); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("100000000000000000000/200000000000000000000") + ); + assert!(snapshot.primary.used_percent.is_finite()); + } + + #[test] + fn omits_fuel_window_when_counts_overflow() { + let account = json!({ "code": 0 }); + let usage = json!({ + "data": { "usage": { "totalToken": 100, "availableToken": 50 } } + }); + let fuel = json!({ + "totalQuota": 1e308, + "list": [ + { "availableToken": 1e308 }, + { "availableToken": 1e308 } + ] + }); + let snapshot = build_snapshot(&account, None, Some(&usage), Some(&fuel)).unwrap(); + assert!(snapshot.secondary.is_none()); + assert!(snapshot.primary.used_percent.is_finite()); + } } From 703e0d54db3293d4102e8730e4a2e6f2b53f56d1 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 00:52:18 +0700 Subject: [PATCH 02/62] Harden oversized provider usage values --- rust/src/providers/chutes/mod.rs | 219 ++++++++++++++++-- rust/src/providers/kimi/mod.rs | 37 ++- .../src/providers/minimax/coding_plan_html.rs | 38 ++- rust/src/providers/perplexity/mod.rs | 60 +++-- 4 files changed, 309 insertions(+), 45 deletions(-) diff --git a/rust/src/providers/chutes/mod.rs b/rust/src/providers/chutes/mod.rs index 5c69951749..c6c75c4727 100644 --- a/rust/src/providers/chutes/mod.rs +++ b/rust/src/providers/chutes/mod.rs @@ -141,9 +141,15 @@ fn collect_windows(value: &Value, out: &mut Vec) { fn window_from_object(map: &serde_json::Map) -> Option { let percent = percent_from_object(map)?; let detail = quota_count_description(map); + let window_minutes = window_minutes_from_object(map); // Only parse dedicated reset timestamp keys — never raw quota counts. let resets_at = first_reset_timestamp(map); - Some(RateWindow::with_details(percent, None, resets_at, detail)) + Some(RateWindow::with_details( + percent, + window_minutes, + resets_at, + detail, + )) } fn percent_from_object(map: &serde_json::Map) -> Option { @@ -162,7 +168,9 @@ fn percent_from_object(map: &serde_json::Map) -> Option { // 0..=1 as fractions turned a real 1% into a false 100% exhausted // state (#408; same class as #247 / upstream #3216, fixed for // opencodego in #407). - return Some(v.clamp(0.0, 100.0)); + if v.is_finite() { + return Some(v.clamp(0.0, 100.0)); + } } let used = first_f64(map, &["used", "usage", "current_usage", "currentUsage"]); let limit = first_f64( @@ -171,13 +179,21 @@ fn percent_from_object(map: &serde_json::Map) -> Option { ); let remaining = first_f64(map, &["remaining", "remaining_quota", "remainingQuota"]); match (used, limit, remaining) { - (Some(used), Some(limit), _) if limit > 0.0 => Some((used / limit) * 100.0), + (Some(used), Some(limit), _) if limit > 0.0 => { + let percent = (used / limit) * 100.0; + percent.is_finite().then_some(percent.clamp(0.0, 100.0)) + } (None, Some(limit), Some(remaining)) if limit > 0.0 => { - Some(((limit - remaining).max(0.0) / limit) * 100.0) + let percent = ((limit - remaining).max(0.0) / limit) * 100.0; + percent.is_finite().then_some(percent.clamp(0.0, 100.0)) } (Some(used), None, Some(remaining)) => { let limit = used + remaining; - (limit > 0.0).then_some((used / limit) * 100.0) + if limit <= 0.0 { + return None; + } + let percent = (used / limit) * 100.0; + percent.is_finite().then_some(percent.clamp(0.0, 100.0)) } _ => None, } @@ -196,7 +212,7 @@ fn quota_count_description(map: &serde_json::Map) -> Option 0.0)?; + let limit = limit.filter(|l| l.is_finite() && *l > 0.0)?; let used = match used { Some(u) => u, None => remaining.map(|r| (limit - r).max(0.0))?, @@ -210,6 +226,123 @@ fn quota_count_description(map: &serde_json::Map) -> Option) -> Option { + for (keys, multiplier) in [ + ( + [ + "window_minutes", + "windowMinutes", + "period_minutes", + "periodMinutes", + "duration_minutes", + "durationMinutes", + ] + .as_slice(), + 1.0, + ), + ( + [ + "window_hours", + "windowHours", + "period_hours", + "periodHours", + "duration_hours", + "durationHours", + ] + .as_slice(), + 60.0, + ), + ( + [ + "window_days", + "windowDays", + "period_days", + "periodDays", + "duration_days", + "durationDays", + ] + .as_slice(), + 24.0 * 60.0, + ), + ( + [ + "window_seconds", + "windowSeconds", + "period_seconds", + "periodSeconds", + "duration_seconds", + "durationSeconds", + ] + .as_slice(), + 1.0 / 60.0, + ), + ] { + if let Some(minutes) = keys.iter().find_map(|key| { + map.get(*key) + .and_then(numeric_value) + .and_then(|value| rounded_window_minutes(value * multiplier)) + }) { + return Some(minutes); + } + } + + ["window", "period", "interval", "duration"] + .iter() + .find_map(|key| map.get(*key).and_then(Value::as_str)) + .and_then(parse_window_duration_text) +} + +fn numeric_value(value: &Value) -> Option { + match value { + Value::Number(number) => number.as_f64().filter(|value| value.is_finite()), + Value::String(text) => text + .trim() + .parse::() + .ok() + .filter(|value| value.is_finite()), + _ => None, + } +} + +fn rounded_window_minutes(value: f64) -> Option { + if !value.is_finite() || value <= 0.0 { + return None; + } + let rounded = value.round(); + if rounded <= 0.0 || rounded > u32::MAX as f64 { + return None; + } + #[expect( + clippy::cast_possible_truncation, + reason = "rounded value is bounded by u32::MAX" + )] + Some(rounded as u32) +} + +fn parse_window_duration_text(raw: &str) -> Option { + let compact: String = raw + .chars() + .filter(|character| !character.is_whitespace()) + .collect(); + let split_at = compact.find(|character: char| { + !character.is_ascii_digit() && !matches!(character, '.' | '+' | '-' | 'e' | 'E') + })?; + let (number, suffix) = compact.split_at(split_at); + let value = number.parse::().ok()?; + let multiplier = if suffix.starts_with("min") || suffix == "m" { + 1.0 + } else if suffix.starts_with("hour") || suffix.starts_with("hr") || suffix == "h" { + 60.0 + } else if suffix.starts_with("day") || suffix == "d" { + 24.0 * 60.0 + } else if suffix.starts_with("month") || suffix == "mo" { + 30.0 * 24.0 * 60.0 + } else { + return None; + }; + rounded_window_minutes(value * multiplier) +} + fn first_reset_timestamp(map: &serde_json::Map) -> Option> { for key in [ "resets_at", @@ -274,8 +407,11 @@ fn epoch_to_datetime(value: f64) -> Option> { } fn first_f64(map: &serde_json::Map, keys: &[&str]) -> Option { - keys.iter() - .find_map(|k| map.get(*k).and_then(Value::as_f64)) + keys.iter().find_map(|k| { + map.get(*k) + .and_then(Value::as_f64) + .filter(|v| v.is_finite()) + }) } fn first_str<'a>(map: &'a serde_json::Map, keys: &[&str]) -> Option<&'a str> { @@ -286,14 +422,18 @@ fn first_str<'a>(map: &'a serde_json::Map, keys: &[&str]) -> Opti } fn format_quota_amount(value: f64) -> String { - if (value - value.round()).abs() < 0.0001 { + if !value.is_finite() { + return "unknown".to_string(); + } + let rounded = value.round(); + if (value - rounded).abs() < 0.0001 && rounded >= i64::MIN as f64 && rounded < i64::MAX as f64 { // Guarded above: value is within 0.0001 of a whole number, so the - // fractional part is zero. - #[expect( + // fractional part is zero and the rounded value fits in i64. + #[allow( clippy::cast_possible_truncation, - reason = "whole-number guard above; fractional part is zero" + reason = "finite rounded value is bounded to the i64 range above" )] - let whole = value.round() as i64; + let whole = rounded as i64; format!("{}", whole) } else { let mut text = format!("{value:.2}"); @@ -393,4 +533,57 @@ mod tests { })); assert_eq!(snapshot.primary.used_percent, 100.0); } + + #[test] + fn large_quota_amounts_keep_their_description() { + let snapshot = snapshot_from_usage(&serde_json::json!({ + "rolling_window": {"used": 1e20, "limit": 2e20, "unit": "credits"} + })); + assert_eq!(snapshot.primary.used_percent, 50.0); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("100000000000000000000/200000000000000000000 credits") + ); + } + + #[test] + fn duration_fields_populate_window_minutes() { + let snapshot = snapshot_from_usage(&serde_json::json!({ + "quotas": [ + {"used": 25, "limit": 100, "duration": "4 hours"}, + {"used": 1, "limit": 2, "window_seconds": 1800} + ] + })); + assert_eq!(snapshot.primary.window_minutes, Some(240)); + assert_eq!( + snapshot.secondary.as_ref().unwrap().window_minutes, + Some(30) + ); + } + + #[test] + fn unrepresentable_duration_keeps_usage_with_unknown_window() { + let snapshot = snapshot_from_usage(&serde_json::json!({ + "rolling_window": { + "used": 25, + "limit": 100, + "window_hours": "1e308" + } + })); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!(snapshot.primary.window_minutes, None); + } + + #[test] + fn non_finite_amount_formatting_is_safe() { + assert_eq!(format_quota_amount(f64::INFINITY), "unknown"); + assert_eq!(format_quota_amount(f64::NAN), "unknown"); + } + + #[test] + fn oversized_integral_amount_is_not_saturated_to_i64_max() { + let value = 2_f64.powi(63); + + assert_eq!(format_quota_amount(value), "9223372036854775808"); + } } diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index abc47c76e7..d96d3a36fc 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -411,8 +411,8 @@ fn kimi_window_minutes(window: &KimiWindow) -> Option { match unit.as_str() { "second" | "seconds" => Some((window.duration / 60).max(1)), "minute" | "minutes" => Some(window.duration), - "hour" | "hours" => Some(window.duration.saturating_mul(60)), - "day" | "days" => Some(window.duration.saturating_mul(24 * 60)), + "hour" | "hours" => window.duration.checked_mul(60), + "day" | "days" => window.duration.checked_mul(24 * 60), _ => None, } } @@ -583,14 +583,22 @@ fn ascii_header_value(raw: &str) -> String { } fn format_usage_amount(value: f64) -> String { - if (value.fract()).abs() < f64::EPSILON { - // Value verified integral to f64 precision; the i64 cast loses nothing. + if value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value < i64::MAX as f64 + { + // The strict upper bound excludes 2^63, which is representable as f64 + // but has no exact i64 representation. #[allow( clippy::cast_possible_truncation, - reason = "guarded by the fract() == 0 check above" + reason = "finite integral value is bounded to the i64 range above" )] let integral = value as i64; format!("{integral}") + } else if value.is_finite() && value.fract() == 0.0 { + // Preserve a large integral value without saturating it to i64::MAX. + format!("{value:.0}") } else { format!("{value:.2}") } @@ -849,4 +857,23 @@ mod tests { assert_eq!(cleaned_owned("'token'").as_deref(), Some("token")); assert!(cleaned_owned(" ").is_none()); } + + #[test] + fn oversized_integral_usage_amount_is_not_saturated_to_i64_max() { + let value = 2_f64.powi(63); + let formatted = format_usage_amount(value); + + assert!(formatted.starts_with("9223372036854775808")); + assert_ne!(formatted, i64::MAX.to_string()); + } + + #[test] + fn overflowing_window_units_are_omitted() { + let window = KimiWindow { + duration: u32::MAX, + time_unit: "hours".to_string(), + }; + + assert_eq!(kimi_window_minutes(&window), None); + } } diff --git a/rust/src/providers/minimax/coding_plan_html.rs b/rust/src/providers/minimax/coding_plan_html.rs index 35bc47f205..0c4ee6c7cc 100644 --- a/rust/src/providers/minimax/coding_plan_html.rs +++ b/rust/src/providers/minimax/coding_plan_html.rs @@ -159,7 +159,7 @@ fn parse_available_usage(text: &str) -> Option<(i64, u32)> { return None; } let duration: f64 = duration_raw.parse().ok()?; - let window_minutes = minutes_from_duration(duration, unit_raw); + let window_minutes = minutes_from_duration(duration, unit_raw)?; if window_minutes == 0 { return None; } @@ -167,14 +167,23 @@ fn parse_available_usage(text: &str) -> Option<(i64, u32)> { } /// Convert a duration + unit to minutes (upstream `minutes(from:unit:)`). -fn minutes_from_duration(value: f64, unit: &str) -> u32 { - // Window lengths come from the provider's own dashboard text and are - // minutes-scale; u32 overflow would need a >8000-year window. - #[allow( - clippy::cast_possible_truncation, - reason = "dashboard window durations are minutes-scale; u32 is far beyond any real window" - )] - let to_minutes = |scaled: f64| -> u32 { scaled.round() as u32 }; +fn minutes_from_duration(value: f64, unit: &str) -> Option { + let to_minutes = |scaled: f64| -> Option { + if !scaled.is_finite() { + return None; + } + let rounded = scaled.round(); + if !(0.0..=u32::MAX as f64).contains(&rounded) { + return None; + } + #[allow( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "finite rounded value is bounded to the non-negative u32 range above" + )] + let rounded = rounded as u32; + Some(rounded) + }; let lower = unit.to_lowercase(); if lower.starts_with('d') { return to_minutes(value * 24.0 * 60.0); @@ -186,9 +195,9 @@ fn minutes_from_duration(value: f64, unit: &str) -> u32 { return to_minutes(value); } if lower.starts_with('s') { - return to_minutes(value / 60.0).max(1); + return to_minutes(value / 60.0).map(|minutes| minutes.max(1)); } - 0 + None } /// Parse "37% used" or "used 37%" (upstream `parseUsedPercent`). @@ -608,4 +617,11 @@ mod tests { assert!((usage.primary.used_percent - 25.0).abs() < 0.01); assert_eq!(usage.login_method.as_deref(), Some("Text Generation Pro")); } + + #[test] + fn oversized_html_duration_is_omitted() { + assert_eq!(minutes_from_duration(f64::MAX, "hours"), None); + assert_eq!(minutes_from_duration(f64::INFINITY, "days"), None); + assert_eq!(minutes_from_duration(5.0, "hours"), Some(300)); + } } diff --git a/rust/src/providers/perplexity/mod.rs b/rust/src/providers/perplexity/mod.rs index 06fbb1c9da..25057b50f6 100644 --- a/rust/src/providers/perplexity/mod.rs +++ b/rust/src/providers/perplexity/mod.rs @@ -72,6 +72,9 @@ impl PerplexityProvider { } fn ts_to_datetime(ts: f64) -> Option> { + if !ts.is_finite() { + return None; + } // Grant expiry epochs are whole-second unix timestamps, far below i64::MAX. #[expect( clippy::cast_possible_truncation, @@ -117,7 +120,7 @@ impl PerplexityProvider { let purchased_used = remaining_usage.min(purchased_total); let pct = |used: f64, total: f64| -> f64 { - if total <= 0.0 { + if !used.is_finite() || !total.is_finite() || total <= 0.0 { 0.0 } else { ((used / total) * 100.0).clamp(0.0, 100.0) @@ -128,33 +131,26 @@ impl PerplexityProvider { let mut primary = RateWindow::new(pct(recurring_used, recurring_total)); primary.resets_at = renewal; - primary.reset_description = Some(format!( - "${:.2}/${:.2}", - recurring_used / 100.0, - recurring_total / 100.0 - )); + primary.reset_description = Self::credit_description(recurring_used, recurring_total); let mut snapshot = UsageSnapshot::new(primary); if bonus_total > 0.0 { let mut secondary = RateWindow::new(pct(bonus_used, bonus_total)); secondary.resets_at = bonus_expiry; - let mut bonus_description = - format!("${:.2}/${:.2}", bonus_used / 100.0, bonus_total / 100.0); - if let Some(expiry) = bonus_expiry { - bonus_description.push_str(&format!(" · exp. {}", expiry.format("%Y-%m-%d"))); + let mut bonus_description = Self::credit_description(bonus_used, bonus_total); + if let Some(expiry) = bonus_expiry + && let Some(description) = bonus_description.as_mut() + { + description.push_str(&format!(" · exp. {}", expiry.format("%Y-%m-%d"))); } - secondary.reset_description = Some(bonus_description); + secondary.reset_description = bonus_description; snapshot = snapshot.with_secondary(secondary); } if purchased_total > 0.0 { let mut tertiary = RateWindow::new(pct(purchased_used, purchased_total)); - tertiary.reset_description = Some(format!( - "${:.2}/${:.2}", - purchased_used / 100.0, - purchased_total / 100.0 - )); + tertiary.reset_description = Self::credit_description(purchased_used, purchased_total); snapshot = snapshot.with_tertiary(tertiary); } @@ -175,6 +171,13 @@ impl PerplexityProvider { Ok(snapshot) } + fn credit_description(used: f64, total: f64) -> Option { + if !used.is_finite() || !total.is_finite() { + return None; + } + Some(format!("${:.2}/${:.2}", used / 100.0, total / 100.0)) + } + async fn fetch_with_cookies( &self, cookie_header: &str, @@ -317,4 +320,29 @@ mod tests { let snap = PerplexityProvider::parse_response(resp).unwrap(); assert_eq!(snap.login_method.as_deref(), Some("Max")); } + + #[test] + fn oversized_credit_totals_do_not_render_non_finite_descriptions() { + let resp = CreditsResponse { + balance_cents: 0.0, + renewal_date_ts: None, + current_period_purchased_cents: 0.0, + credit_grants: vec![ + CreditGrant { + grant_type: "recurring".to_string(), + amount_cents: f64::MAX, + expires_at_ts: None, + }, + CreditGrant { + grant_type: "recurring".to_string(), + amount_cents: f64::MAX, + expires_at_ts: None, + }, + ], + total_usage_cents: 0.0, + }; + + let snap = PerplexityProvider::parse_response(resp).unwrap(); + assert!(snap.primary.reset_description.is_none()); + } } From 33fde92d643f90855bbfa64c5cf635c1b37d0b37 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 01:05:45 +0700 Subject: [PATCH 03/62] Reconcile Kimi zero ratio placeholders --- rust/src/providers/kimi/code_api.rs | 230 +++++++++++++++++++++++++++- 1 file changed, 225 insertions(+), 5 deletions(-) diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 907241d5d9..7e488594e8 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -9,8 +9,9 @@ use std::path::{Path, PathBuf}; use super::web; use super::{ - FetchContext, KimiCodeApiUsageResponse, KimiProvider, ProviderError, UsageSnapshot, - ascii_header_value, cleaned_env, cleaned_owned, kimi_window_minutes, + FetchContext, KimiCodeApiUsageResponse, KimiProvider, KimiRatioPool, KimiUsageDetail, + ProviderError, UsageSnapshot, ascii_header_value, cleaned_env, cleaned_owned, + kimi_window_minutes, }; const KIMI_CODE_API_BASE: &str = "https://api.kimi.com"; @@ -116,16 +117,40 @@ pub(super) fn snapshot_from_code_api_response( response: KimiCodeApiUsageResponse, ) -> Result { let pools_present = response.usages.is_some(); + let legacy_limit = response.limits.as_ref().and_then(|limits| limits.first()); + let legacy_session_minutes = legacy_limit.map(|limit| { + limit + .window + .as_ref() + .and_then(kimi_window_minutes) + .unwrap_or(300) + }); let session_pool = response .usages .as_ref() .and_then(|pools| pools.session.as_ref()) - .and_then(|pool| pool.rate_window(300)); + .and_then(|pool| { + resolved_ratio_window( + &response, + pool, + legacy_limit.map(|limit| &limit.detail), + 300, + legacy_session_minutes, + ) + }); let weekly_pool = response .usages .as_ref() .and_then(|pools| pools.weekly.as_ref()) - .and_then(|pool| pool.rate_window(10_080)); + .and_then(|pool| { + resolved_ratio_window( + &response, + pool, + response.usage.as_ref(), + 10_080, + Some(10_080), + ) + }); let monthly_pool = response .usages .as_ref() @@ -139,7 +164,9 @@ pub(super) fn snapshot_from_code_api_response( response .usage .as_ref() - .and_then(|detail| KimiProvider::rate_window_from_usage_detail(detail, None).ok()) + .and_then(|detail| { + KimiProvider::rate_window_from_usage_detail(detail, Some(10_080)).ok() + }) .ok_or_else(|| { ProviderError::Parse("Kimi Code API has no usable quota window".into()) })? @@ -165,6 +192,54 @@ pub(super) fn snapshot_from_code_api_response( } Ok(usage) } + +/// Resolve a ratio pool while recognizing the mixed legacy response used by +/// Kimi accounts during the pool migration. A zero ratio is authoritative for +/// monthly-pool accounts and for any response without matching reliable count +/// evidence. Only a same-duration, same-reset count window can replace it. +fn resolved_ratio_window( + response: &KimiCodeApiUsageResponse, + pool: &KimiRatioPool, + detail: Option<&KimiUsageDetail>, + window_minutes: u32, + count_window_minutes: Option, +) -> Option { + let ratio_window = pool.rate_window(window_minutes)?; + if ratio_window.used_percent != 0.0 + || response + .usages + .as_ref() + .and_then(|pools| pools.monthly.as_ref()) + .is_some() + || count_window_minutes != Some(window_minutes) + { + return Some(ratio_window); + } + + let Some(detail) = detail else { + return Some(ratio_window); + }; + let Some(used) = + super::value_as_f64(detail.used.as_ref()).filter(|value| value.is_finite() && *value > 0.0) + else { + return Some(ratio_window); + }; + let Some(count_window) = + KimiProvider::rate_window_from_usage_detail(detail, Some(window_minutes)).ok() + else { + return Some(ratio_window); + }; + let (Some(count_reset), Some(ratio_reset)) = (count_window.resets_at, ratio_window.resets_at) + else { + return Some(ratio_window); + }; + + if (count_reset - ratio_reset).num_milliseconds().abs() <= 2_000 && used > 0.0 { + Some(count_window) + } else { + Some(ratio_window) + } +} pub(crate) fn code_api_key(explicit: Option<&str>) -> Result { if let Some(key) = explicit.map(str::trim).filter(|key| !key.is_empty()) { return Ok(key.to_string()); @@ -474,4 +549,149 @@ mod tests { if message.contains("unusable session quota pool") )); } + + #[test] + fn zero_ratio_placeholders_fall_back_to_matching_legacy_counts() { + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { + "limit": "100", + "used": "19", + "remaining": "81", + "resetTime": "2026-09-19T16:45:59.449979Z" + }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { + "limit": "100", + "used": "1", + "remaining": "99", + "resetTime": "2026-09-19T14:45:59.449979Z" + } + }], + "usages": { + "limit_5h": { + "used_ratio": 0, + "reset_time": "2026-09-19T14:45:58Z" + }, + "limit_7d": { + "used_ratio": 0, + "reset_time": "2026-09-19T16:45:58Z" + } + } + })) + .unwrap(); + + let snapshot = snapshot_from_code_api_response(response).unwrap(); + assert_eq!(snapshot.primary.used_percent, 1.0); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + let weekly = snapshot.secondary.expect("weekly count fallback"); + assert_eq!(weekly.used_percent, 19.0); + assert_eq!(weekly.window_minutes, Some(10_080)); + } + + #[test] + fn zero_ratio_with_different_reset_stays_authoritative() { + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { + "limit": "100", + "used": "19", + "resetTime": "2026-09-19T16:45:59Z" + }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { + "limit": "100", + "used": "1", + "resetTime": "2026-09-19T14:45:59Z" + } + }], + "usages": { + "limit_5h": { + "used_ratio": 0, + "reset_time": "2026-09-19T14:46:03Z" + }, + "limit_7d": { + "used_ratio": 0, + "reset_time": "2026-09-19T16:46:03Z" + } + } + })) + .unwrap(); + + let snapshot = snapshot_from_code_api_response(response).unwrap(); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.secondary.unwrap().used_percent, 0.0); + } + + #[test] + fn monthly_pool_keeps_zero_ratios_even_with_matching_counts() { + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { + "limit": "100", + "used": "19", + "resetTime": "2026-09-19T16:45:59Z" + }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { + "limit": "100", + "used": "1", + "resetTime": "2026-09-19T14:45:59Z" + } + }], + "usages": { + "limit_5h": { + "used_ratio": 0, + "reset_time": "2026-09-19T14:45:58Z" + }, + "limit_7d": { + "used_ratio": 0, + "reset_time": "2026-09-19T16:45:58Z" + }, + "limit_month_total": { "used_ratio": 0.0313 } + } + })) + .unwrap(); + + let snapshot = snapshot_from_code_api_response(response).unwrap(); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.secondary.unwrap().used_percent, 0.0); + assert!((snapshot.tertiary.unwrap().used_percent - 3.13).abs() < 0.000_001); + } + + #[test] + fn invalid_legacy_counts_do_not_override_zero_ratio() { + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { + "limit": "100", + "used": "invalid", + "remaining": "99", + "resetTime": "2026-09-19T16:45:59Z" + }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { + "limit": "100", + "used": "-1", + "remaining": "99", + "resetTime": "2026-09-19T14:45:59Z" + } + }], + "usages": { + "limit_5h": { + "used_ratio": 0, + "reset_time": "2026-09-19T14:45:58Z" + }, + "limit_7d": { + "used_ratio": 0, + "reset_time": "2026-09-19T16:45:58Z" + } + } + })) + .unwrap(); + + let snapshot = snapshot_from_code_api_response(response).unwrap(); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.secondary.unwrap().used_percent, 0.0); + } } From 594159caf9d9dd5985a88a574cb58e119a5ccc8c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:31:12 +0700 Subject: [PATCH 04/62] Port standalone Pi history --- .../src-tauri/src/auto_refresh.rs | 2 +- .../src-tauri/src/commands/chart.rs | 1 + .../src-tauri/src/commands/spend_contract.rs | 3 +- .../src-tauri/src/commands/usage_spend.rs | 104 ++- .../src/components/providers/providerIcons.ts | 1 + .../src/lib/providerCharts.test.ts | 1 + apps/desktop-tauri/src/lib/providerCharts.ts | 2 +- .../providers/ProvidersSidebar.test.tsx | 2 +- .../sections/charts/ChartsSection.tsx | 4 +- .../desktop-tauri/src/test/providerCatalog.ts | 1 + rust/src/cli/cost.rs | 31 +- rust/src/cli/serve/dashboard/source.rs | 25 +- rust/src/cli/serve/data.rs | 1 + rust/src/cli/usage.rs | 2 +- rust/src/core/provider.rs | 10 +- rust/src/core/provider_factory.rs | 9 +- rust/src/core/token_accounts.rs | 1 + rust/src/cost_scanner.rs | 109 +++- rust/src/pi_session_cost.rs | 595 ++++++++++++++++-- rust/src/providers/mod.rs | 2 + rust/src/providers/pi.rs | 87 +++ rust/src/spend_contract.rs | 1 + 22 files changed, 883 insertions(+), 111 deletions(-) create mode 100644 rust/src/providers/pi.rs diff --git a/apps/desktop-tauri/src-tauri/src/auto_refresh.rs b/apps/desktop-tauri/src-tauri/src/auto_refresh.rs index 6825f4170d..75128cb00d 100644 --- a/apps/desktop-tauri/src-tauri/src/auto_refresh.rs +++ b/apps/desktop-tauri/src-tauri/src/auto_refresh.rs @@ -218,7 +218,7 @@ fn local_usage_provider_ids(settings: &Settings) -> Vec { .get_enabled_provider_ids() .into_iter() .map(|provider| provider.cli_name().to_string()) - .filter(|provider_id| matches!(provider_id.as_str(), "codex" | "claude" | "muse")) + .filter(|provider_id| matches!(provider_id.as_str(), "codex" | "claude" | "pi" | "muse")) .collect() } diff --git a/apps/desktop-tauri/src-tauri/src/commands/chart.rs b/apps/desktop-tauri/src-tauri/src/commands/chart.rs index 5e66b2d003..4561dd9129 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/chart.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/chart.rs @@ -644,6 +644,7 @@ fn scan_local_cost( match provider_id { "codex" => Some(scanner.scan_codex_with_cancel(cancel)), "claude" => Some(scanner.scan_claude_with_cancel(cancel)), + "pi" => Some(scanner.scan_pi_with_cancel(cancel)), "opencodego" => Some(scanner.scan_opencodego_with_cancel(cancel)), _ => None, } diff --git a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs index a3e3637f12..4403ba7ad0 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs @@ -11,7 +11,7 @@ pub async fn get_spend_contract( include_open_codex: Option, ) -> Result { let provider = provider_id.trim().to_ascii_lowercase(); - if !matches!(provider.as_str(), "codex" | "claude" | "opencodego") { + if !matches!(provider.as_str(), "codex" | "claude" | "pi" | "opencodego") { return Err(format!( "Spend contract is unavailable for provider: {provider}" )); @@ -24,6 +24,7 @@ pub async fn get_spend_contract( let summary = match provider.as_str() { "codex" => scanner.scan_codex(), "claude" => scanner.scan_claude(), + "pi" => scanner.scan_pi(), "opencodego" => scanner.scan_opencodego_with_cancel(None), _ => unreachable!(), }; diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index f55ae71cd2..01f562e111 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -363,6 +363,9 @@ fn build_usage_spend_summary( ) -> UsageSpendSummary { let include_opencodex = settings.open_codex_usage_logs_enabled; let hide_native = settings.hide_native_codex_cost_when_open_codex_present; + let pi_selected = settings.enabled_providers.iter().any(|id| id == "pi") + || cached.iter().any(|snapshot| snapshot.provider_id == "pi"); + let include_pi_in_native = !pi_selected; // Upstream 0.55.0 #3105: independent provider baselines load in parallel. // Keep each provider's 7d/30d scans serial so they can safely share that @@ -372,29 +375,43 @@ fn build_usage_spend_summary( } else { codexbar::core::CostScanOptions::default() }; - let ((codex_7_summary, codex_30_summary), (claude_7_summary, claude_30_summary)) = - std::thread::scope(|scope| { - let codex = scope.spawn(move || { - ( - CostScanner::new(7) - .with_options(codex_scan_options) - .scan_codex(), - CostScanner::new(30) - .with_options(codex_scan_options) - .scan_codex(), - ) - }); - let claude = scope.spawn(|| { - ( - CostScanner::new(7).scan_claude(), - CostScanner::new(30).scan_claude(), - ) - }); + let mut codex_scan_options = codex_scan_options; + codex_scan_options.include_pi_sessions = include_pi_in_native; + let ( + (codex_7_summary, codex_30_summary), + (claude_7_summary, claude_30_summary), + (pi_7_summary, pi_30_summary), + ) = std::thread::scope(|scope| { + let codex = scope.spawn(move || { + ( + CostScanner::new(7) + .with_options(codex_scan_options) + .scan_codex(), + CostScanner::new(30) + .with_options(codex_scan_options) + .scan_codex(), + ) + }); + let claude = scope.spawn(|| { + ( + CostScanner::new(7) + .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native), + CostScanner::new(30) + .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native), + ) + }); + let pi = scope.spawn(|| { ( - codex.join().expect("Codex spend scan worker panicked"), - claude.join().expect("Claude spend scan worker panicked"), + CostScanner::new(7).scan_pi(), + CostScanner::new(30).scan_pi(), ) }); + ( + codex.join().expect("Codex spend scan worker panicked"), + claude.join().expect("Claude spend scan worker panicked"), + pi.join().expect("Pi spend scan worker panicked"), + ) + }); let codex_stale = !codex_30_summary.history_coverage_established; let codex_stale_updated_at = codex_stale @@ -421,6 +438,22 @@ fn build_usage_spend_summary( settings.hide_personal_info, codex_30_summary.clone(), ); + let pi_7_contract = build_local_spend_contract_from_summary( + "pi", + 7, + false, + false, + settings.hide_personal_info, + pi_7_summary.clone(), + ); + let pi_30_contract = build_local_spend_contract_from_summary( + "pi", + 30, + false, + false, + settings.hide_personal_info, + pi_30_summary.clone(), + ); let mut provider_ids: BTreeSet = settings.enabled_providers.iter().cloned().collect(); provider_ids.extend(cached.iter().map(|snapshot| snapshot.provider_id.clone())); @@ -494,6 +527,15 @@ fn build_usage_spend_summary( refreshing: false, stale_updated_at: None, }, + "pi" => SpendValues { + seven_day: pi_7_contract.known_cost_usd, + thirty_day: pi_30_contract.known_cost_usd, + seven_day_tokens: total_token_mix(&pi_7_contract.token_mix), + thirty_day_tokens: total_token_mix(&pi_30_contract.token_mix), + source: "local Pi/OMP history".to_string(), + refreshing: !pi_30_summary.history_coverage_established, + stale_updated_at: None, + }, "opencodego" | "kimi" | "deepseek" if include_opencodex => { let seven = build_local_spend_contract(&provider_id, 7, true); let thirty = build_local_spend_contract(&provider_id, 30, true); @@ -585,8 +627,11 @@ fn build_usage_spend_summary( thirty_day_tokens: spend.thirty_day_tokens, currency, source: spend.source, - included_in_overview: settings.enabled_providers.contains(&provider_id) - || cached_snapshot.is_some(), + included_in_overview: include_in_shared_overview( + &provider_id, + settings.enabled_providers.contains(&provider_id), + cached_snapshot.is_some(), + ), daily, refreshing: spend.refreshing, stale_updated_at: spend.stale_updated_at, @@ -623,6 +668,13 @@ fn build_usage_spend_summary( } } +/// Pi is an alternate local-history view over rows that may already be +/// projected into Codex or Claude. Keep it out of the shared denominator so +/// enabling Pi cannot double-count the same physical usage. +fn include_in_shared_overview(provider_id: &str, enabled: bool, cached: bool) -> bool { + provider_id != "pi" && (enabled || cached) +} + fn last_included_reporting_day(contract: &SpendContract) -> String { contract .daily @@ -774,4 +826,12 @@ mod cache_key_tests { let private = usage_spend_cache_key_with_privacy(&[], 30, false, false, true); assert_ne!(public, private); } + + #[test] + fn pi_history_is_an_alternate_view_not_a_shared_overview_source() { + assert!(!include_in_shared_overview("pi", true, true)); + assert!(include_in_shared_overview("codex", true, false)); + assert!(include_in_shared_overview("claude", false, true)); + assert!(!include_in_shared_overview("codex", false, false)); + } } diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index a2cf5dda8a..ccfa304078 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -166,6 +166,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { antigravity: { id: "antigravity", brandColor: "#60ba7e", fallbackLetter: "◉", svgPath: RAW.antigravity }, augment: { id: "augment", brandColor: "#6366f1", fallbackLetter: "A", svgPath: RAW.augment }, claude: { id: "claude", brandColor: "#cc7c5e", fallbackLetter: "◈", svgPath: RAW.claude }, + pi: { id: "pi", brandColor: "#7c3aed", fallbackLetter: "P" }, codebuff: { id: "codebuff", brandColor: "#44ff00", fallbackLetter: "B", svgPath: RAW.codebuff }, coderabbit: { id: "coderabbit", brandColor: "#ff5c35", fallbackLetter: "C", svgPath: RAW.coderabbit }, codex: { id: "codex", brandColor: "#49a3b0", fallbackLetter: "◆", svgPath: RAW.codex }, diff --git a/apps/desktop-tauri/src/lib/providerCharts.test.ts b/apps/desktop-tauri/src/lib/providerCharts.test.ts index bf8d1179a4..0f3c27a6be 100644 --- a/apps/desktop-tauri/src/lib/providerCharts.test.ts +++ b/apps/desktop-tauri/src/lib/providerCharts.test.ts @@ -7,6 +7,7 @@ describe("providerSupportsChartData", () => { expect(providerSupportsChartData("claude")).toBe(true); expect(providerSupportsChartData("openai")).toBe(true); expect(providerSupportsChartData("muse")).toBe(true); + expect(providerSupportsChartData("pi")).toBe(true); expect(providerSupportsChartData("OpenAI")).toBe(true); expect(providerSupportsChartData("copilot")).toBe(false); diff --git a/apps/desktop-tauri/src/lib/providerCharts.ts b/apps/desktop-tauri/src/lib/providerCharts.ts index f230f80090..6a4bec3586 100644 --- a/apps/desktop-tauri/src/lib/providerCharts.ts +++ b/apps/desktop-tauri/src/lib/providerCharts.ts @@ -1,4 +1,4 @@ -const PROVIDER_CHART_DATA_IDS = new Set(["claude", "codex", "muse", "openai"]); +const PROVIDER_CHART_DATA_IDS = new Set(["claude", "codex", "muse", "openai", "pi"]); export function providerSupportsChartData(providerId: string): boolean { return PROVIDER_CHART_DATA_IDS.has(providerId.toLowerCase()); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx index 8b947a7dc0..e85aa379ce 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx @@ -133,7 +133,7 @@ describe("ProvidersSidebar", () => { container.querySelectorAll(".providers-sidebar__name"), (node) => node.textContent, ); - expect(names.slice(0, 3)).toEqual(["Claude", "Codex", "Cursor"]); + expect(names.slice(0, 3)).toEqual(["Claude", "Codex", "Pi"]); }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.tsx index 86bcb183eb..8257e3069c 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.tsx @@ -89,7 +89,9 @@ export function ChartsSection({ providerId, accountEmail, accentColor, t }: Prop // Upstream 0.50.0 #2930: Codex defaults to exact local token totals. const defaultTab: TabKey = - (providerId === "codex" || providerId === "muse") && hasTokens ? "tokens" : available[0]; + (providerId === "codex" || providerId === "muse" || providerId === "pi") && hasTokens + ? "tokens" + : available[0]; const current: TabKey = active && available.includes(active) ? active : defaultTab; const emptyMsg = t("DetailChartEmpty"); diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index 6eac860900..c8ab2daa5d 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -1,6 +1,7 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["codex", "Codex"], ["claude", "Claude"], + ["pi", "Pi"], ["cursor", "Cursor"], ["factory", "Factory"], ["gemini", "Gemini"], diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index fe12351104..8e9be16dc2 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -16,7 +16,7 @@ use crate::spend_contract::build_local_spend_contract_from_summary; /// Arguments for the cost command #[derive(Args, Debug, Default)] pub struct CostArgs { - /// Provider to query (codex, claude, muse, antigravity, cursor, gemini, copilot, all, both) + /// Provider to query (codex, claude, pi, muse, antigravity, cursor, gemini, copilot, all, both) #[arg(short, long)] pub provider: Option, @@ -109,7 +109,12 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { } let mut scan_options = CostScanOptions::app_driven(); - scan_options.include_pi_sessions = !args.provider_native_only; + let requested_providers = providers.as_list(); + let pi_selected = requested_providers.contains(&ProviderId::Pi); + // When Pi is selected alongside native providers, the standalone Pi row + // owns its mirrored Codex/Claude events. A single native-provider request + // keeps the historical inclusive behavior unless explicitly narrowed. + scan_options.include_pi_sessions = !args.provider_native_only && !pi_selected; let scanner = CostScanner::new(args.days).with_options(scan_options); tracing::debug!( @@ -135,7 +140,21 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { }); } ProviderId::Claude => { - let summary = scanner.scan_claude(); + let summary = if pi_selected || args.provider_native_only { + scanner.scan_claude_with_cancel_and_pi_sessions(None, false) + } else { + scanner.scan_claude() + }; + results.push(CostResult { + provider: provider.cli_name().to_string(), + display_name: provider.display_name().to_string(), + summary, + supported: true, + token_history: None, + }); + } + ProviderId::Pi => { + let summary = scanner.scan_pi(); results.push(CostResult { provider: provider.cli_name().to_string(), display_name: provider.display_name().to_string(), @@ -450,7 +469,7 @@ fn build_json_payloads(results: &[CostResult], days: u32) -> Vec Vec serde_json::Value::String("complete".to_string()), crate::cost_scanner::ModelPricingCompleteness::Partial { unpriced_models } => serde_json::json!({"partial": {"unpriced_models": unpriced_models}}), diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 76dc1dd162..8d28758df4 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -16,6 +16,7 @@ use chrono::{Local, Utc}; use crate::core::{CostScanOptions, FetchContext, ProviderId, SourceMode, instantiate_provider}; use crate::cost_scanner::{self, CostScanner}; use crate::settings::Settings; +use crate::spend_contract::build_local_spend_contract_from_summary; use crate::cli::serve::collection::{ AccountFetchEnvelope, ClaudeAccountsInput, ProviderFetchEnvelope, RawCostPayload, @@ -111,7 +112,7 @@ impl SnapshotProducer { let providers: Vec = indexed.into_iter().map(|(_, envelope)| envelope).collect(); - let costs = collect_costs().await; + let costs = collect_costs(provider_ids.contains(&ProviderId::Pi)).await; let claude_accounts = collect_claude_accounts(provider_ids.contains(&ProviderId::Claude)).await; @@ -205,13 +206,18 @@ async fn bounded_fetch( } } -/// Local cost data for the two scanned providers, computed off the async +/// Local cost data for the scanned providers, computed off the async /// runtime so a large corpus cannot stall dashboard builds. -async fn collect_costs() -> HashMap { - let result = tokio::task::spawn_blocking(|| { - let scanner = CostScanner::new(30).with_options(CostScanOptions::app_driven()); +async fn collect_costs(pi_selected: bool) -> HashMap { + let result = tokio::task::spawn_blocking(move || { + let mut scan_options = CostScanOptions::app_driven(); + scan_options.include_pi_sessions = !pi_selected; + let scanner = CostScanner::new(30).with_options(scan_options); let codex = scanner.scan_codex_with_cancel(None); - let claude = scanner.scan_claude_with_cancel(None); + let claude = scanner.scan_claude_with_cancel_and_pi_sessions(None, !pi_selected); + let pi = scanner.scan_pi_with_cancel(None); + let pi_contract = + build_local_spend_contract_from_summary("pi", 30, false, false, false, pi); let today = Local::now().date_naive().format("%Y-%m-%d").to_string(); let today_of = |provider: &str| { cost_scanner::get_daily_cost_history(provider, 30) @@ -234,6 +240,13 @@ async fn collect_costs() -> HashMap { last_30_days_usd: Some(claude.total_cost_usd), }, ); + costs.insert( + "pi".to_string(), + RawCostPayload { + today_usd: today_of("pi"), + last_30_days_usd: pi_contract.known_cost_usd, + }, + ); costs }) .await; diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index 37130eba8d..d1841db6bc 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -88,6 +88,7 @@ pub async fn cost_response(provider: Option<&str>) -> String { let (supported, summary) = match provider_id { ProviderId::Codex => (true, scanner.scan_codex()), ProviderId::Claude => (true, scanner.scan_claude()), + ProviderId::Pi => (true, scanner.scan_pi()), _ => (false, Default::default()), }; if supported { diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index f33373b90c..2b1dfff0af 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -25,7 +25,7 @@ pub(super) enum UsageOutput { Toon(Vec), } -pub const PROVIDER_ARG_HELP: &str = "Provider to query (for example: codex, claude, gemini, antigravity/agy, nanogpt, deepseek, codebuff, windsurf, all, both)"; +pub const PROVIDER_ARG_HELP: &str = "Provider to query (for example: codex, claude, pi, gemini, antigravity/agy, nanogpt, deepseek, codebuff, windsurf, all, both)"; /// Arguments for the usage command #[derive(Args, Debug, Default)] diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index feedadedcd..d6a38df551 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -14,6 +14,7 @@ use super::provider_state::ProviderStateKind; pub enum ProviderId { Codex, Claude, + Pi, Cursor, Factory, Gemini, @@ -97,6 +98,7 @@ impl ProviderId { &[ ProviderId::Codex, ProviderId::Claude, + ProviderId::Pi, ProviderId::Cursor, ProviderId::Factory, ProviderId::Gemini, @@ -179,6 +181,7 @@ impl ProviderId { match self { ProviderId::Codex => "codex", ProviderId::Claude => "claude", + ProviderId::Pi => "pi", ProviderId::Cursor => "cursor", ProviderId::Factory => "factory", ProviderId::Gemini => "gemini", @@ -261,6 +264,7 @@ impl ProviderId { match self { ProviderId::Codex => "Codex", ProviderId::Claude => "Claude", + ProviderId::Pi => "Pi", ProviderId::Cursor => "Cursor", ProviderId::Factory => "Factory", ProviderId::Gemini => "Gemini", @@ -349,6 +353,7 @@ impl ProviderId { ProviderId::Cursor => Some("cursor.com"), ProviderId::Factory => Some("app.factory.ai"), ProviderId::Codex => Some("chatgpt.com"), + ProviderId::Pi => None, ProviderId::Gemini => Some("aistudio.google.com"), ProviderId::Kiro => Some("kiro.dev"), ProviderId::Kimi => Some("kimi.moonshot.cn"), @@ -432,6 +437,7 @@ impl ProviderId { match name.to_lowercase().as_str() { "codex" | "openai" => Some(ProviderId::Codex), "claude" | "anthropic" => Some(ProviderId::Claude), + "pi" | "pi-mono" => Some(ProviderId::Pi), "cursor" => Some(ProviderId::Cursor), "factory" | "droid" => Some(ProviderId::Factory), "gemini" | "google" => Some(ProviderId::Gemini), @@ -1001,6 +1007,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { match id { ProviderId::Codex => "#49A3B0", ProviderId::Claude => "#CC7C5E", + ProviderId::Pi => "#7C3AED", ProviderId::Cursor => "#00BFA5", ProviderId::Factory => "#FF6B35", ProviderId::Gemini => "#AB87EA", @@ -1089,9 +1096,10 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 76); + assert_eq!(all.len(), 77); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); + assert!(all.contains(&ProviderId::Pi)); assert!(all.contains(&ProviderId::Fireworks)); assert!(all.contains(&ProviderId::Kimi)); assert!(all.contains(&ProviderId::KimiK2)); diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index b8fb6f7f7e..aff4b6f31d 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -18,10 +18,10 @@ use crate::providers::{ LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, NousProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, - PerplexityProvider, PoeProvider, QoderProvider, QwenCloudProvider, ReplicateProvider, - SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, - VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, - ZedProvider, ZenMuxProvider, ZoomMateProvider, + PerplexityProvider, PiProvider, PoeProvider, QoderProvider, QwenCloudProvider, + ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, + VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, + XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -32,6 +32,7 @@ pub fn instantiate(id: ProviderId) -> Box { match id { ProviderId::Claude => Box::new(ClaudeProvider::new()), ProviderId::Codex => Box::new(CodexProvider::new()), + ProviderId::Pi => Box::new(PiProvider::new()), ProviderId::Cursor => Box::new(CursorProvider::new()), ProviderId::Gemini => Box::new(GeminiProvider::new()), ProviderId::Copilot => Box::new(CopilotProvider::new()), diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 0e1b6c076f..a8f3cc71ae 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -332,6 +332,7 @@ impl TokenAccountSupport { }), // These providers don't support token accounts ProviderId::Codex + | ProviderId::Pi | ProviderId::Gemini | ProviderId::Antigravity | ProviderId::Kiro diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index e8b1dd03cd..619f9ba7b0 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -523,6 +523,38 @@ impl CostScanner { self } + /// Scan standalone Pi and OMP local history. + /// + /// Pi session rows can represent either Codex or Claude models. They are + /// priced with the mapped provider's table but owned by the standalone Pi + /// source, so this path never adds native Codex/Claude transcripts. + pub fn scan_pi(&self) -> CostSummary { + self.scan_pi_with_cancel(None) + } + + pub fn scan_pi_with_cancel(&self, cancel: Option<&AtomicBool>) -> CostSummary { + let today = Utc::now().date_naive(); + let mut summary = CostSummary { + period_start: Some(today - Duration::days(self.days as i64)), + period_end: Some(today), + ..CostSummary::default() + }; + let mut seen_entries = HashSet::new(); + let evidence = crate::pi_session_cost::scan_pi_into( + &mut summary, + self.days, + cancel, + &mut seen_entries, + ); + summary.history_coverage_established = evidence.complete && !is_cancelled(cancel); + summary.known_zero = summary.history_coverage_established + && summary.sessions_count == 0 + && summary.input_tokens == 0 + && summary.output_tokens == 0 + && summary.cached_tokens == 0; + summary + } + /// Scan Codex local logs pub fn scan_claude(&self) -> CostSummary { self.scan_claude_with_cancel(None) @@ -530,6 +562,19 @@ impl CostScanner { /// Scan Claude local logs, stopping early when the caller cancels the scan. pub fn scan_claude_with_cancel(&self, cancel: Option<&AtomicBool>) -> CostSummary { + self.scan_claude_with_cancel_and_pi_sessions(cancel, true) + } + + /// Scan Claude local logs with optional Pi/OMP-compatible history. + /// + /// The default scanner remains inclusive for backwards compatibility. A + /// combined Codex/Claude/Pi selection can turn this off so the standalone + /// Pi row owns those mirrored events exactly once. + pub fn scan_claude_with_cancel_and_pi_sessions( + &self, + cancel: Option<&AtomicBool>, + include_pi_sessions: bool, + ) -> CostSummary { let projects_dir = self.get_claude_projects_dir(); let mut summary = CostSummary::default(); let today = Utc::now().date_naive(); @@ -565,14 +610,16 @@ impl CostScanner { } // OMP / pi-compatible anthropic rows, deduped across shared files. - let mut seen_pi = HashSet::new(); - crate::pi_session_cost::scan_pi_compatible_into( - &mut summary, - crate::pi_session_cost::PiMappedProvider::Claude, - self.days, - cancel, - &mut seen_pi, - ); + if include_pi_sessions { + let mut seen_pi = HashSet::new(); + crate::pi_session_cost::scan_pi_compatible_into( + &mut summary, + crate::pi_session_cost::PiMappedProvider::Claude, + self.days, + cancel, + &mut seen_pi, + ); + } // Claude has no persisted provider cost-report cache in the Windows // port. Rebuilding from the transcript inventory on every scan makes @@ -1092,7 +1139,7 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< let date_str = date.format("%Y-%m-%d").to_string(); daily_costs.insert( date_str, - (provider != "codex" && provider != "claude").then_some(0.0), + (provider != "codex" && provider != "claude" && provider != "pi").then_some(0.0), ); } @@ -1173,6 +1220,21 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< } } } + "pi" => { + let scan = crate::pi_session_cost::scan_pi_daily(days, None); + for (day_key, cost) in &scan.costs { + if let Some(slot) = daily_costs.get_mut(day_key) { + *slot = (!scan.unpriced_days.contains(day_key)).then_some(*cost); + } + } + if scan.history_coverage_established { + for (day_key, slot) in &mut daily_costs { + if slot.is_none() && !scan.unpriced_days.contains(day_key) { + *slot = Some(0.0); + } + } + } + } _ => {} } @@ -1247,6 +1309,17 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); } } + "pi" => { + let scan = crate::pi_session_cost::scan_pi_daily(days, None); + for (day_key, tokens) in scan.tokens { + if let Some(slot) = daily_tokens.get_mut(&day_key) { + *slot = tokens; + } + } + if scan.history_coverage_established { + covered_days.extend(daily_tokens.keys().cloned()); + } + } _ => {} } @@ -1257,12 +1330,18 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> // Codex only: the bounded catch-up may not have reached the requested // depth yet. Incomplete = history exists but the oldest quarter of the // window has no scanned day. - let incomplete = provider == "codex" - && !covered_days.is_empty() - && covered_days.len() < days as usize - && result[..(result.len() / 4).max(1)] - .iter() - .any(|(date, _)| !covered_days.contains(date)); + let incomplete = if provider == "pi" { + // Pi scans are bounded filesystem walks, so a complete parse covers + // the requested window even when the roots contain no sessions. + covered_days.is_empty() + } else { + provider == "codex" + && !covered_days.is_empty() + && covered_days.len() < days as usize + && result[..(result.len() / 4).max(1)] + .iter() + .any(|(date, _)| !covered_days.contains(date)) + }; (result, incomplete) } diff --git a/rust/src/pi_session_cost.rs b/rust/src/pi_session_cost.rs index e372bec322..a9f8f1b82c 100644 --- a/rust/src/pi_session_cost.rs +++ b/rust/src/pi_session_cost.rs @@ -1,10 +1,10 @@ //! Pi-compatible + OMP agent session cost scan (upstream #2269). //! -//! Walks `~/.pi/agent/sessions/**/*.jsonl` and `~/.omp/agent/sessions/**/*.jsonl` -//! and attributes openai-codex / anthropic assistant rows into cost summaries -//! without double-counting the same entry id across shared files. +//! Resolves Pi-family session roots and walks their JSONL files, attributing +//! openai-codex / anthropic assistant rows into cost summaries without +//! double-counting the same entry id across shared files. -use chrono::{DateTime, Duration, Utc}; +use chrono::{DateTime, Duration, Local, Utc}; use serde_json::Value; use std::collections::HashSet; use std::fs::File; @@ -12,8 +12,12 @@ use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; use std::sync::atomic::AtomicBool; +use crate::agent_sessions::pi_family::roots::{ + EnvMap, PiProfile, omp_all_profile_roots, omp_default_profile_root, omp_named_profile_root, + omp_profile_selector, pi_settings_session_directory, +}; use crate::core::CostUsagePricing; -use crate::cost_scanner::CostSummary; +use crate::cost_scanner::{CostSummary, ModelPricingCompleteness}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PiMappedProvider { @@ -21,59 +25,254 @@ pub enum PiMappedProvider { Claude, } +#[derive(Debug, Default)] +pub struct PiDailyScan { + pub costs: std::collections::HashMap, + pub tokens: std::collections::HashMap, + pub unpriced_days: HashSet, + pub history_coverage_established: bool, +} + +#[derive(Debug, Clone, Copy)] +pub struct PiScanEvidence { + pub complete: bool, +} + +impl Default for PiScanEvidence { + fn default() -> Self { + Self { complete: true } + } +} + /// Session roots to scan: `.pi` and `.omp` under the user home. pub fn pi_compatible_session_roots(home: Option) -> Vec { let Some(home) = home else { return Vec::new(); }; - [".pi", ".omp"] + let cwd = std::env::current_dir().unwrap_or_else(|_| home.clone()); + let environment: EnvMap = std::env::vars().collect(); + pi_compatible_session_roots_for(&home, &cwd, &environment) +} + +fn pi_compatible_session_roots_for(home: &Path, cwd: &Path, environment: &EnvMap) -> Vec { + let mut roots = Vec::new(); + let pi_root = environment + .get("PI_CODING_AGENT_SESSION_DIR") + .and_then(|value| resolve_environment_path(value, cwd)) + .or_else(|| { + environment + .get("PI_CODING_AGENT_DIR") + .and_then(|value| resolve_environment_path(value, cwd)) + .map(|root| root.join("sessions")) + }) + .or_else(|| pi_settings_session_directory(cwd, home)) + .unwrap_or_else(|| home.join(".pi").join("agent").join("sessions")); + roots.push(pi_root); + + match omp_profile_selector(environment) { + PiProfile::Invalid => {} + PiProfile::Named(profile) => { + if let Some(root) = omp_named_profile_root(&profile, environment, cwd, home) { + roots.push(root); + } + } + PiProfile::Default => { + if let Some(root) = omp_default_profile_root(environment, cwd, home) { + roots.push(root); + } + roots.extend(omp_all_profile_roots(environment, home)); + } + } + + let mut seen = HashSet::new(); + roots .into_iter() - .map(|dir| home.join(dir).join("agent").join("sessions")) + .filter(|root| { + let key = std::fs::canonicalize(root) + .unwrap_or_else(|_| root.clone()) + .to_string_lossy() + .to_ascii_lowercase(); + seen.insert(key) + }) .collect() } +fn resolve_environment_path(value: &str, cwd: &Path) -> Option { + let value = value.trim(); + if value.is_empty() { + return None; + } + let path = PathBuf::from(value); + Some(if path.is_absolute() { + path + } else { + cwd.join(path) + }) +} + pub fn scan_pi_compatible_into( summary: &mut CostSummary, target: PiMappedProvider, days: u32, cancel: Option<&AtomicBool>, seen_entries: &mut HashSet, -) { +) -> PiScanEvidence { + scan_roots_into( + summary, + days, + cancel, + seen_entries, + pi_compatible_session_roots(dirs::home_dir()), + Some(target), + ) +} + +/// Scan Pi and OMP history as one standalone provider-owned source. +/// +/// The compatible Codex/Claude paths above intentionally project these rows +/// into their native summaries. This path keeps the same parser and pricing, +/// but accepts both mapped providers and shares one deduplication set across +/// both roots so standalone Pi history is not double-counted. +pub fn scan_pi_into( + summary: &mut CostSummary, + days: u32, + cancel: Option<&AtomicBool>, + seen_entries: &mut HashSet, +) -> PiScanEvidence { + scan_roots_into( + summary, + days, + cancel, + seen_entries, + pi_compatible_session_roots(dirs::home_dir()), + None, + ) +} + +/// Scan standalone Pi/OMP history into daily cost and token buckets. +pub fn scan_pi_daily(days: u32, cancel: Option<&AtomicBool>) -> PiDailyScan { + let cutoff = Utc::now() - Duration::days(days as i64); + scan_pi_daily_from_roots( + cutoff, + cancel, + pi_compatible_session_roots(dirs::home_dir()), + ) +} + +fn scan_pi_daily_from_roots( + cutoff: DateTime, + cancel: Option<&AtomicBool>, + roots: Vec, +) -> PiDailyScan { + let mut result = PiDailyScan::default(); + let mut seen_entries = HashSet::new(); + let mut missing_timestamp = false; + let mut evidence = PiScanEvidence::default(); + for root in roots { + if cancelled(cancel) { + break; + } + if !root.is_dir() { + continue; + } + if !walk_jsonl(&root, cancel, &mut |path| { + if cancelled(cancel) { + return false; + } + let file = for_each_pi_entry(path, cutoff, None, &mut seen_entries, |entry| { + let Some(timestamp) = entry.timestamp else { + missing_timestamp = true; + return; + }; + let day = timestamp + .with_timezone(&Local) + .date_naive() + .format("%Y-%m-%d") + .to_string(); + if !entry.pricing_known { + result.unpriced_days.insert(day.clone()); + } + *result.costs.entry(day.clone()).or_insert(0.0) += entry.cost; + let tokens = entry.input.saturating_add(entry.output); + let total = result.tokens.entry(day.clone()).or_insert(0); + *total = total.saturating_add(tokens); + }); + file.complete + }) { + evidence.complete = false; + } + } + result.history_coverage_established = + evidence.complete && !cancelled(cancel) && !missing_timestamp; + result +} + +fn scan_roots_into( + summary: &mut CostSummary, + days: u32, + cancel: Option<&AtomicBool>, + seen_entries: &mut HashSet, + roots: Vec, + target: Option, +) -> PiScanEvidence { let cutoff = Utc::now() - Duration::days(days as i64); let mut sessions = 0u32; - for root in pi_compatible_session_roots(dirs::home_dir()) { + let mut evidence = PiScanEvidence::default(); + for root in roots { if cancelled(cancel) { break; } if !root.is_dir() { continue; } - walk_jsonl(&root, cancel, &mut |path| { + if !walk_jsonl(&root, cancel, &mut |path| { if cancelled(cancel) { - return; + return false; } let before = seen_entries.len(); - let counted = for_each_pi_entry(path, cutoff, target, seen_entries, |entry| { + let file = for_each_pi_entry(path, cutoff, target, seen_entries, |entry| { apply_entry(summary, &entry); }); - if counted > 0 || seen_entries.len() > before { + if file.counted > 0 || seen_entries.len() > before { sessions += 1; } - }); + file.complete + }) { + evidence.complete = false; + } } summary.sessions_count = summary.sessions_count.saturating_add(sessions); + evidence } struct PiEntry { + timestamp: Option>, + provider: PiMappedProvider, model: String, input: u64, output: u64, cache_read: u64, cache_create: u64, cost: f64, + pricing_known: bool, } fn apply_entry(summary: &mut CostSummary, entry: &PiEntry) { + if !entry.pricing_known { + summary.unknown_models.insert(entry.model.clone()); + match &mut summary.model_pricing_completeness { + ModelPricingCompleteness::Complete => { + summary.model_pricing_completeness = ModelPricingCompleteness::Partial { + unpriced_models: vec![entry.model.clone()], + }; + } + ModelPricingCompleteness::Partial { unpriced_models } => { + if !unpriced_models.contains(&entry.model) { + unpriced_models.push(entry.model.clone()); + } + } + } + } summary.input_tokens += entry.input; summary.output_tokens += entry.output; summary.cached_tokens += entry.cache_read + entry.cache_create; @@ -92,62 +291,130 @@ fn cancelled(cancel: Option<&AtomicBool>) -> bool { cancel.is_some_and(|f| f.load(std::sync::atomic::Ordering::Relaxed)) } -fn walk_jsonl(root: &Path, cancel: Option<&AtomicBool>, on_file: &mut dyn FnMut(&Path)) { +fn walk_jsonl( + root: &Path, + cancel: Option<&AtomicBool>, + on_file: &mut dyn FnMut(&Path) -> bool, +) -> bool { let Ok(entries) = std::fs::read_dir(root) else { - return; + return false; }; - for entry in entries.flatten() { + let mut complete = true; + for entry in entries { if cancelled(cancel) { - return; + return false; } + let Ok(entry) = entry else { + complete = false; + continue; + }; let path = entry.path(); if path.is_dir() { - walk_jsonl(&path, cancel, on_file); + complete = walk_jsonl(&path, cancel, on_file) && complete; } else if path .extension() .and_then(|e| e.to_str()) .is_some_and(|e| e.eq_ignore_ascii_case("jsonl")) { - on_file(&path); + complete = on_file(&path) && complete; } } + complete +} + +#[derive(Debug, Clone, Copy)] +struct PiFileScanResult { + counted: u32, + complete: bool, } fn for_each_pi_entry( path: &Path, cutoff: DateTime, - target: PiMappedProvider, + target: Option, seen: &mut HashSet, mut on_entry: impl FnMut(PiEntry), -) -> u32 { +) -> PiFileScanResult { let Ok(file) = File::open(path) else { - return 0; + return PiFileScanResult { + counted: 0, + complete: false, + }; }; let mut counted = 0u32; + let mut complete = true; + let mut session_id = None; let reader = BufReader::new(file); - for line in reader.lines().map_while(Result::ok) { + for (ordinal, line_result) in reader.lines().enumerate() { + let Ok(line) = line_result else { + complete = false; + continue; + }; + if line.trim().is_empty() { + continue; + } let Ok(value) = serde_json::from_str::(&line) else { + complete = false; continue; }; - let Some(entry) = parse_pi_assistant_entry(&value, target) else { + if session_id.is_none() { + session_id = value + .get("id") + .and_then(|id| id.as_str()) + .filter(|_| value.get("type").and_then(Value::as_str) == Some("session")) + .map(str::to_string); + } + let Some(entry) = parse_pi_assistant_entry_any(&value) else { + if looks_like_usage_candidate(&value) { + complete = false; + } continue; }; + if entry.timestamp.is_none() { + complete = false; + } if let Some(ts) = entry_timestamp(&value) && ts < cutoff { continue; } - let entry_id = entry_dedup_key(&value, path, counted); + if target.is_some_and(|target| entry.provider != target) { + continue; + } + let entry_id = entry_dedup_key(&value, path, ordinal, session_id.as_deref()); if !seen.insert(entry_id) { continue; } on_entry(entry); counted += 1; } - counted + PiFileScanResult { counted, complete } } -fn entry_dedup_key(value: &Value, path: &Path, ordinal: u32) -> String { +fn looks_like_usage_candidate(value: &Value) -> bool { + let message = value.get("message").unwrap_or(value); + let role = message + .get("role") + .or_else(|| value.get("role")) + .and_then(Value::as_str) + .unwrap_or(""); + let typ = value.get("type").and_then(Value::as_str).unwrap_or(""); + let assistant_shape = role.eq_ignore_ascii_case("assistant") + || typ.eq_ignore_ascii_case("assistant") + || typ.eq_ignore_ascii_case("message"); + assistant_shape + && message + .get("usage") + .or_else(|| value.get("usage")) + .is_some() +} + +fn entry_dedup_key(value: &Value, path: &Path, ordinal: usize, session_id: Option<&str>) -> String { + // Pi/OMP migrations can mirror the same event into both roots, so a + // stable event id wins within the logical session. Scope it with the + // session header when available: message IDs can be reused by separate + // sessions. The file stem and physical line ordinal cover legacy rows + // without a session header or stable message id. if let Some(id) = value .get("id") .or_else(|| value.get("messageId")) @@ -156,7 +423,14 @@ fn entry_dedup_key(value: &Value, path: &Path, ordinal: u32) -> String { .map(str::trim) .filter(|s| !s.is_empty()) { - return id.to_string(); + let scope = session_id + .map(str::to_string) + .or_else(|| { + path.file_stem() + .map(|stem| stem.to_string_lossy().into_owned()) + }) + .unwrap_or_else(|| path.display().to_string()); + return format!("{scope}#{id}"); } format!("{}#{ordinal}", path.display()) } @@ -190,7 +464,13 @@ fn map_provider(raw: &str) -> Option { None } +#[cfg(test)] fn parse_pi_assistant_entry(value: &Value, target: PiMappedProvider) -> Option { + let entry = parse_pi_assistant_entry_any(value)?; + (entry.provider == target).then_some(entry) +} + +fn parse_pi_assistant_entry_any(value: &Value) -> Option { // Accept either flat or nested { message: {...} } pi-compatible rows. let message = value.get("message").unwrap_or(value); let role = message @@ -216,10 +496,6 @@ fn parse_pi_assistant_entry(value: &Value, target: PiMappedProvider) -> Option

Option

CostUsagePricing::codex_cost_usd_with_cache_write( + let (cost, pricing_known) = match mapped { + PiMappedProvider::Codex => match CostUsagePricing::codex_cost_usd_with_cache_write( &model, input, cache_read, cache_create, output, - ) - .unwrap_or(0.0), + ) { + Some(cost) => (cost, true), + None => (0.0, false), + }, PiMappedProvider::Claude => { // Token counts come from API usage records and fit within i32; // the canonical Claude pricing table takes i32 per-token counts. @@ -297,27 +575,37 @@ fn parse_pi_assistant_entry(value: &Value, target: PiMappedProvider) -> Option

(), 132); + } + + #[test] + fn malformed_usage_input_keeps_valid_rows_but_marks_source_incomplete() { + let dir = tempdir().unwrap(); + let sessions = dir.path().join("agent").join("sessions"); + std::fs::create_dir_all(&sessions).unwrap(); + let valid = r#"{"id":"valid","role":"assistant","provider":"openai-codex","model":"gpt-5","timestamp":"2026-07-20T12:00:00Z","usage":{"input":11,"output":3}}"#; + std::fs::write( + sessions.join("mixed.jsonl"), + format!("{valid}\n{{\"role\":\"assistant\",\"usage\":\n"), + ) + .unwrap(); + + let scan = scan_pi_daily_from_roots( + DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc), + None, + vec![sessions.clone()], + ); + assert!(!scan.history_coverage_established); + assert_eq!(scan.tokens.values().sum::(), 14); + + let mut summary = CostSummary::default(); + let mut seen = HashSet::new(); + let evidence = scan_roots_into(&mut summary, 365, None, &mut seen, vec![sessions], None); + assert!(!evidence.complete); + assert_eq!(summary.input_tokens, 11); + } + + #[test] + fn standalone_daily_scan_dedupes_pi_and_omp_roots() { + let dir = tempdir().unwrap(); + let pi_sessions = dir.path().join(".pi").join("agent").join("sessions"); + let omp_sessions = dir.path().join(".omp").join("agent").join("sessions"); + std::fs::create_dir_all(&pi_sessions).unwrap(); + std::fs::create_dir_all(&omp_sessions).unwrap(); + let codex = r#"{"id":"shared","role":"assistant","provider":"openai-codex","model":"gpt-5","timestamp":"2026-07-20T12:00:00Z","usage":{"input":50,"output":5}}"#; + let claude = r#"{"id":"claude-only","role":"assistant","provider":"anthropic","model":"claude-sonnet-4-6","timestamp":"2026-07-20T13:00:00Z","usage":{"input":70,"output":7}}"#; + std::fs::write( + pi_sessions.join("one.jsonl"), + format!("{codex}\n{claude}\n"), + ) + .unwrap(); + std::fs::write(omp_sessions.join("one.jsonl"), format!("{codex}\n")).unwrap(); + + let scan = scan_pi_daily_from_roots( + DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc), + None, + vec![pi_sessions, omp_sessions], + ); + assert!(scan.history_coverage_established); + assert_eq!(scan.tokens.values().sum::(), 132); + assert_eq!(scan.tokens.len(), 1); + } + #[test] fn session_roots_include_pi_and_omp() { - let roots = pi_compatible_session_roots(Some(PathBuf::from("/home/user"))); + let home = PathBuf::from("/home/user"); + let roots = pi_compatible_session_roots_for(&home, &home, &EnvMap::new()); assert!( roots .iter() diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 6ac5e0f1a8..ad697ad321 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -65,6 +65,7 @@ pub mod opencode; pub mod opencodego; pub mod openrouter; pub mod perplexity; +pub mod pi; pub mod poe; pub mod qoder; pub mod qwencloud; @@ -143,6 +144,7 @@ pub use opencode::OpenCodeProvider; pub use opencodego::OpenCodeGoProvider; pub use openrouter::OpenRouterProvider; pub use perplexity::PerplexityProvider; +pub use pi::PiProvider; pub use poe::PoeProvider; pub use qoder::QoderProvider; pub use qwencloud::QwenCloudProvider; diff --git a/rust/src/providers/pi.rs b/rust/src/providers/pi.rs new file mode 100644 index 0000000000..862f069abc --- /dev/null +++ b/rust/src/providers/pi.rs @@ -0,0 +1,87 @@ +//! Local Pi provider. +//! +//! Pi has no remote quota endpoint in the upstream provider model. Its usage +//! and token-cost history come from local Pi/OMP session JSONL files; the +//! ordinary provider refresh therefore exposes an informational local row and +//! leaves cost history to the dedicated scanner path. + +use async_trait::async_trait; + +use crate::core::{ + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, + RateWindow, SourceMode, UsageSnapshot, +}; + +pub struct PiProvider { + metadata: ProviderMetadata, +} + +impl PiProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::Pi, + display_name: "Pi", + session_label: "Session", + weekly_label: "Weekly", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://github.com/badlogic/pi-mono"), + status_page_url: None, + tertiary_label_key: None, + }, + } + } +} + +impl Default for PiProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for PiProvider { + fn id(&self) -> ProviderId { + ProviderId::Pi + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + if ctx.source_mode != SourceMode::Auto { + return Err(ProviderError::UnsupportedSource(ctx.source_mode)); + } + + Ok(ProviderFetchResult::new( + UsageSnapshot::new(RateWindow::informational("Local Pi history")), + "local", + ) + .with_non_authoritative_pace()) + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto] + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exposes_local_only_metadata() { + let provider = PiProvider::new(); + assert_eq!(provider.id(), ProviderId::Pi); + assert_eq!(provider.metadata().display_name, "Pi"); + assert!(!provider.metadata().default_enabled); + assert_eq!(provider.available_sources(), vec![SourceMode::Auto]); + assert!(!provider.supports_oauth()); + assert!(!provider.supports_web()); + assert!(!provider.supports_cli()); + } +} diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index 9a2e928b34..af9cdfef34 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -346,6 +346,7 @@ pub fn build_local_spend_contract( let summary = match provider_id { "codex" => scanner.scan_codex(), "claude" => scanner.scan_claude(), + "pi" => scanner.scan_pi(), "opencodego" => scanner.scan_opencodego_with_cancel(None), _ => CostSummary::default(), }; From fc54a8c81a5191959d85656c10a01598220c90dd Mon Sep 17 00:00:00 2001 From: gleaming9 Date: Tue, 22 Sep 2026 15:51:12 +0900 Subject: [PATCH 05/62] Keep shared-workspace accounts distinct by user --- .../codex_accounts/account_manager/tests.rs | 50 ++++++++- rust/src/codex_accounts/models.rs | 103 ++++++++++++++++++ 2 files changed, 152 insertions(+), 1 deletion(-) diff --git a/rust/src/codex_accounts/account_manager/tests.rs b/rust/src/codex_accounts/account_manager/tests.rs index bfb8584f6c..aefd6f9718 100644 --- a/rust/src/codex_accounts/account_manager/tests.rs +++ b/rust/src/codex_accounts/account_manager/tests.rs @@ -5,9 +5,13 @@ mod tests { /// Write an auth.json carrying a JWT identity for the given account id. fn write_auth(home_path: &Path, email: &str, account_id: &str) { + write_user_auth(home_path, email, account_id, &format!("auth0|{account_id}")); + } + + fn write_user_auth(home_path: &Path, email: &str, account_id: &str, subject: &str) { let payload = serde_json::json!({ "email": email, - "sub": format!("auth0|{account_id}"), + "sub": subject, "https://api.openai.com/auth": { "chatgpt_plan_type": "team", "chatgpt_account_id": account_id, @@ -88,6 +92,50 @@ mod tests { super::super::file_locations::clear_app_support_directory_override(); } + #[test] + fn shared_team_users_keep_separate_discovery_and_managed_homes() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + super::super::file_locations::with_app_support_directory(root.to_path_buf()); + let first_home = root.join("managed-homes").join("first"); + let second_home = root.join("managed-homes").join("second"); + let ambient_home = root.join("ambient"); + for home in [&first_home, &second_home, &ambient_home] { + std::fs::create_dir_all(home).unwrap(); + } + write_user_auth(&first_home, "a@x.test", "shared-team", "user-a"); + write_user_auth(&second_home, "b@x.test", "shared-team", "user-b"); + write_user_auth(&ambient_home, "a@x.test", "shared-team", "user-a"); + let second_auth = std::fs::read(second_home.join("auth.json")).unwrap(); + let mut first = make_account(first_home.clone(), "a@x.test", "shared-team"); + first.auth_subject = Some("user-a".into()); + let manager = CodexAccountManager::new(); + let discovered = manager + .discover_managed_accounts(std::slice::from_ref(&first)) + .unwrap(); + assert_eq!(discovered.len(), 2); + let second = discovered + .iter() + .find(|a| a.codex_home_path == second_home) + .unwrap(); + assert_ne!(second.id, first.id); + assert!(!first.matches(second)); + + let mut ambient = first.clone(); + ambient.codex_home_path = ambient_home; + manager.remove_managed_files_if_owned(&first).unwrap(); + assert!(!first_home.exists()); + assert!(second_home.exists()); + let materialized = manager.materialize_as_managed(&ambient).unwrap(); + assert_ne!(materialized.codex_home_path, second_home); + assert_eq!( + std::fs::read(second_home.join("auth.json")).unwrap(), + second_auth + ); + assert_eq!(managed_home_count(root), 2); + super::super::file_locations::clear_app_support_directory_override(); + } + /// Write an auth.json whose credentials carry an explicit refresh time. fn write_auth_refreshed_at( home_path: &Path, diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index 40c29abb60..53757fb844 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -61,6 +61,25 @@ fn normalize_identifier(value: Option<&str>) -> Option { .filter(|v| !v.is_empty()) } +/// A shared workspace or home cannot override conflicting user evidence. +fn user_identity_conflicts( + subject: Option<&str>, + other_subject: Option<&str>, + email: Option<&str>, + other_email: Option<&str>, +) -> bool { + if let (Some(a), Some(b)) = ( + normalize_identifier(subject), + normalize_identifier(other_subject), + ) { + return a != b; + } + matches!( + (normalize_identifier(email), normalize_identifier(other_email)), + (Some(a), Some(b)) if a != b + ) +} + /// Where an account's `CODEX_HOME` lives. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -206,6 +225,14 @@ impl CodexAccount { /// Whether two accounts refer to the same identity. pub fn matches(&self, other: &CodexAccount) -> bool { + if user_identity_conflicts( + self.auth_subject.as_deref(), + other.auth_subject.as_deref(), + self.email_hint.as_deref(), + other.email_hint.as_deref(), + ) { + return false; + } if let (Some(a), Some(b)) = ( self.effective_workspace_account_id(), other.effective_workspace_account_id(), @@ -323,6 +350,14 @@ impl RemovedAccountIdentity { } pub fn matches(&self, account: &CodexAccount) -> bool { + if user_identity_conflicts( + self.auth_subject.as_deref(), + account.auth_subject.as_deref(), + self.email_hint.as_deref(), + account.email_hint.as_deref(), + ) { + return false; + } if self.standardized_home_path() == account.standardized_home_path() { return true; } @@ -573,6 +608,74 @@ mod tests { assert!(a.matches(&b)); } + #[test] + fn different_subjects_do_not_match_in_a_shared_workspace_or_home() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/shared", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.auth_subject = Some("user-a".into()); + a.email_hint = Some("same@x.test".into()); + let mut b = a.clone(); + b.id = Uuid::new_v4(); + b.auth_subject = Some("user-b".into()); + + assert!(!a.matches(&b)); + assert!(!b.matches(&a)); + assert!(!RemovedAccountIdentity::from_account(&a).matches(&b)); + } + + #[test] + fn different_emails_do_not_match_in_a_shared_workspace_without_subjects() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.email_hint = Some("user-a@x.test".into()); + let mut b = a.clone(); + b.id = Uuid::new_v4(); + b.codex_home_path = PathBuf::from("/managed/b"); + b.email_hint = Some("user-b@x.test".into()); + + assert!(!a.matches(&b)); + assert!(!RemovedAccountIdentity::from_account(&a).matches(&b)); + } + + #[test] + fn matching_subject_remains_authoritative_when_email_changes() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.auth_subject = Some("user-a".into()); + a.email_hint = Some("old@x.test".into()); + let mut b = a.clone(); + b.auth_subject = Some("USER-A".into()); + b.email_hint = Some("new@x.test".into()); + assert!(a.matches(&b)); + } + + #[test] + fn same_user_in_different_workspaces_remains_separate() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("team-a"), + ); + a.auth_subject = Some("same-user".into()); + let mut b = a.clone(); + b.provider_account_id = Some("team-b".into()); + assert!(!a.matches(&b)); + assert!(!b.matches(&a)); + } + #[test] fn disambiguates_different_provider_ids() { let a = account( From 5fcc5e88d1013febd0b40d2f8e41c62a71e77c2e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:28:42 +0700 Subject: [PATCH 06/62] Honor Kimi manual cookie policy --- rust/src/providers/kimi/web.rs | 54 +++++++++++++++++----------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 0cdfad9cfa..4f5289e097 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -26,16 +26,17 @@ pub(crate) fn cookie_source() -> String { .to_string() } -/// Upstream `KimiBrowserImportPolicy.allowsImport`: everything but `off`. +/// Upstream `KimiBrowserImportPolicy.allowsImport`: automatic discovery is +/// allowed only when the user selected the automatic source. fn browser_import_allowed(cookie_source: &str) -> bool { - !cookie_source.eq_ignore_ascii_case("off") + cookie_source.eq_ignore_ascii_case("auto") || cookie_source.eq_ignore_ascii_case("browser") } /// Web auth token chain for both the web fetch and the Code-API enrichment /// (upstream `KimiWebEnrichmentTokenResolver.resolve`): /// 1. Manual cookie header (its `kimi-auth`/auth cookie), source-independent. -/// 2. Kimi Desktop session token (skipped when cookie source is `off`). -/// 3. Browser cookie import (skipped when cookie source is `off`). +/// 2. Kimi Desktop session token (automatic source only). +/// 3. Browser cookie import (automatic source only). pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { resolve_web_tokens(WebTokenInput { manual_header, @@ -369,22 +370,17 @@ mod tests { } #[test] - fn cookie_source_off_blocks_desktop_and_browser_but_not_manual() { - assert_eq!( - resolve_web_tokens(input(None, "off", static_desktop, static_browser)), - Vec::new() - ); - assert_eq!( - resolve_web_tokens(input(None, "off", no_token, static_browser)), - Vec::new() - ); - assert_eq!( - resolve_web_tokens(input(Some("kimi-auth=manual"), "off", no_token, no_token)), - vec![WebTokenCandidate { - token: "manual".to_string(), - source: WebTokenSource::Manual, - }] - ); + fn off_and_manual_sources_block_automatic_discovery() { + for source in ["off", "manual"] { + assert_eq!( + resolve_web_tokens(input(None, source, static_desktop, static_browser)), + Vec::new() + ); + assert_eq!( + resolve_web_tokens(input(Some("not-a-token"), source, no_token, static_browser)), + Vec::new() + ); + } } #[test] @@ -400,15 +396,18 @@ mod tests { } #[test] - fn manual_default_source_still_allows_desktop_token() { - // Upstream: desktop-session token applies for any non-off source; - // the local default ("manual") must keep desktop sessions working. - let candidates = resolve_web_tokens(input(None, "manual", static_desktop, no_token)); + fn explicit_manual_token_stays_authoritative() { + let candidates = resolve_web_tokens(input( + Some("kimi-auth=manual-token"), + "manual", + static_desktop, + static_browser, + )); assert_eq!( candidates, vec![WebTokenCandidate { - token: "desktop-token".to_string(), - source: WebTokenSource::Desktop, + token: "manual-token".to_string(), + source: WebTokenSource::Manual, }] ); } @@ -429,7 +428,8 @@ mod tests { fn browser_import_gate_is_case_insensitive() { assert!(!browser_import_allowed("OFF")); assert!(browser_import_allowed("browser")); - assert!(browser_import_allowed("manual")); + assert!(browser_import_allowed("AUTO")); + assert!(!browser_import_allowed("manual")); } #[test] From 7ca9714cd3fa6b535f9ecec542eca638d125ad75 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:49:04 +0700 Subject: [PATCH 07/62] Default Kimi cookie discovery to automatic --- rust/src/settings.rs | 9 +++++++-- rust/src/settings/tests.rs | 7 +++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..a46445ea1e 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -933,12 +933,17 @@ impl Settings { self.provider_configs.entry(id).or_default() } - /// Cookie source for `id`, or the default `"manual"` if unset. + /// Cookie source for `id`. Kimi follows upstream's automatic default; + /// providers with no specific default retain the legacy manual default. pub fn cookie_source(&self, id: ProviderId) -> &str { self.provider_configs .get(&id) .and_then(|c| c.cookie_source.as_deref()) - .unwrap_or(DEFAULT_COOKIE_SOURCE) + .unwrap_or(if id == ProviderId::Kimi { + "auto" + } else { + DEFAULT_COOKIE_SOURCE + }) } pub fn set_cookie_source(&mut self, id: ProviderId, source: impl Into) { diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..b109b1f88b 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -31,6 +31,13 @@ fn test_settings_default() { ); } +#[test] +fn kimi_cookie_source_defaults_to_automatic_discovery() { + let settings = Settings::default(); + assert_eq!(settings.cookie_source(ProviderId::Kimi), "auto"); + assert_eq!(settings.cookie_source(ProviderId::Claude), "manual"); +} + #[test] fn overview_layout_defaults_to_compact_and_round_trips() { let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) From bc004906135a03eaf1dba06bb906dd8a8be3ea5e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:49:51 +0700 Subject: [PATCH 08/62] Route Kimi through selected region --- .../src-tauri/src/commands/provider_detail.rs | 6 + .../src/commands/provider_settings.rs | 13 +++ .../src-tauri/src/commands/system.rs | 10 ++ .../src-tauri/src/commands/tests.rs | 21 ++++ .../providers/sections/RegionSection.tsx | 2 +- rust/src/providers/kimi/code_api.rs | 30 ++--- rust/src/providers/kimi/desktop_token.rs | 55 ++++++--- rust/src/providers/kimi/mod.rs | 33 ++++-- rust/src/providers/kimi/region.rs | 106 ++++++++++++++++++ rust/src/providers/kimi/web.rs | 84 ++++++++------ rust/src/providers/mod.rs | 2 +- 11 files changed, 281 insertions(+), 81 deletions(-) create mode 100644 rust/src/providers/kimi/region.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs index 418bd5c146..a061856da7 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs @@ -79,6 +79,12 @@ pub(crate) fn build_provider_detail( settings.api_region(id), )), ) + } else if id == codexbar::core::ProviderId::Kimi { + Some( + codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) + .console_url() + .to_string(), + ) } else { metadata.dashboard_url.map(|s| s.to_string()) }; diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs index 7d26ba8fdc..ad1eb069ba 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -270,6 +270,7 @@ fn region_provider(provider_id: &str) -> Option { "alibabatokenplan" => ProviderId::AlibabaTokenPlan, "zai" => ProviderId::Zai, "minimax" => ProviderId::MiniMax, + "kimi" => ProviderId::Kimi, _ => return None, }) } @@ -282,6 +283,10 @@ pub(crate) fn provider_region_lookup(settings: &Settings, provider_id: &str) -> )) .settings_value() .to_string() + } else if id == codexbar::core::ProviderId::Kimi { + codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) + .settings_value() + .to_string() } else { settings.api_region(id).to_string() } @@ -784,6 +789,14 @@ pub fn region_options_for(provider_id: &str) -> Vec { .to_string(), }, ], + "kimi" => codexbar::providers::KimiRegion::ALL + .iter() + .copied() + .map(|region| RegionOption { + value: region.settings_value().to_string(), + label: region.display_name().to_string(), + }) + .collect(), "alibabatokenplan" => codexbar::providers::AlibabaTokenPlanRegion::ALL .iter() .copied() diff --git a/apps/desktop-tauri/src-tauri/src/commands/system.rs b/apps/desktop-tauri/src-tauri/src/commands/system.rs index 9ae17f57ce..104018628e 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/system.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/system.rs @@ -205,6 +205,16 @@ pub fn quit_app(app: tauri::AppHandle) { } fn dashboard_url_for_provider(provider_id: &str) -> Option { + if provider_id == ProviderId::Kimi.cli_name() { + let settings = Settings::load(); + return Some( + codexbar::providers::KimiRegion::from_settings(Some( + settings.api_region(ProviderId::Kimi), + )) + .console_url() + .to_string(), + ); + } if provider_id == ProviderId::MiniMax.cli_name() { let settings = Settings::load(); return Some( diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index d25a0a064a..f9f615463b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -263,6 +263,20 @@ fn minimax_region_lookup_normalizes_legacy_china_value() { assert_eq!(provider_region_lookup(&s, "minimax").as_deref(), Some("cn")); } +#[test] +fn kimi_region_lookup_defaults_to_china_and_roundtrips_international() { + let mut settings = Settings::default(); + assert_eq!( + provider_region_lookup(&settings, "kimi").as_deref(), + Some("china") + ); + super::provider_region_set(&mut settings, "kimi", "international".to_string()).unwrap(); + assert_eq!( + provider_region_lookup(&settings, "kimi").as_deref(), + Some("international") + ); +} + #[test] fn minimax_cookie_domain_follows_selected_region() { let mut s = Settings::default(); @@ -1824,6 +1838,13 @@ fn minimax_region_options_match_upstream_hosts() { ); } +#[test] +fn kimi_region_options_match_regional_hosts() { + let opts = super::region_options_for("kimi"); + let values: Vec<_> = opts.iter().map(|option| option.value.as_str()).collect(); + assert_eq!(values, vec!["china", "international"]); +} + #[test] fn region_options_empty_for_non_regional_provider() { assert!(super::region_options_for("claude").is_empty()); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx index 77ed218f79..e1a95c64aa 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx @@ -12,7 +12,7 @@ interface Props { } /** - * API-region dropdown for Alibaba / Z.ai / MiniMax. + * API-region dropdown for providers with regional endpoints, including Kimi. * * Port of the region ComboBox rows in * `rust/src/native_ui/preferences.rs::render_provider_detail_panel`. diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 7e488594e8..2a80f73347 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -9,12 +9,11 @@ use std::path::{Path, PathBuf}; use super::web; use super::{ - FetchContext, KimiCodeApiUsageResponse, KimiProvider, KimiRatioPool, KimiUsageDetail, - ProviderError, UsageSnapshot, ascii_header_value, cleaned_env, cleaned_owned, + FetchContext, KimiCodeApiUsageResponse, KimiProvider, KimiRatioPool, KimiRegion, + KimiUsageDetail, ProviderError, UsageSnapshot, ascii_header_value, cleaned_env, cleaned_owned, kimi_window_minutes, }; -const KIMI_CODE_API_BASE: &str = "https://api.kimi.com"; const KIMI_CODE_API_KEY_ENV: &str = "KIMI_CODE_API_KEY"; const KIMI_CODE_BASE_URL_ENV: &str = "KIMI_CODE_BASE_URL"; const KIMI_CODE_HOME_ENV: &str = "KIMI_CODE_HOME"; @@ -43,12 +42,13 @@ struct KimiCodeCredentialFile { /// to the un-enriched snapshot. pub(crate) async fn fetch_via_code_api( ctx: &FetchContext, + region: KimiRegion, api_key_override: Option<&str>, identity_headers_override: Option<&[(&str, String)]>, login_method: &str, ) -> Result { let api_key = code_api_key(api_key_override.or(ctx.api_key.as_deref()))?; - let base_url = code_api_base_url()?; + let base_url = code_api_base_url(region)?; let endpoint = code_api_usage_endpoint(&base_url)?; let client = crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(30)) @@ -89,14 +89,15 @@ pub(crate) async fn fetch_via_code_api( // Upstream #2622: enrich Code API + CLI usage with the monthly membership // pool from a signed-in Kimi Desktop (or browser/manual) session. - for web_token in web::web_auth_tokens(ctx.manual_cookie_header.as_deref()) { - match web::fetch_subscription_for_enrichment_result(&client, &web_token).await { + for web_token in web::web_auth_tokens(ctx.manual_cookie_header.as_deref(), region) { + match web::fetch_subscription_for_enrichment_result(&client, &web_token, region).await { Ok(subscription) => { if let Some(subscription) = subscription { snapshot = super::apply_subscription_windows(snapshot, &subscription); } if !has_plan_name - && let Some(plan) = web::fetch_subscription_plan(&client, &web_token).await + && let Some(plan) = + web::fetch_subscription_plan(&client, &web_token, region).await { snapshot.login_method = Some(plan); } @@ -247,8 +248,9 @@ pub(crate) fn code_api_key(explicit: Option<&str>) -> Result Result { - let raw = cleaned_env(KIMI_CODE_BASE_URL_ENV).unwrap_or_else(|| KIMI_CODE_API_BASE.to_string()); +fn code_api_base_url(region: KimiRegion) -> Result { + let raw = cleaned_env(KIMI_CODE_BASE_URL_ENV) + .unwrap_or_else(|| region.code_api_base_url().to_string()); crate::providers::validated_https_url(&raw, "Kimi Code API base") } @@ -285,8 +287,8 @@ pub(crate) fn kimi_code_home() -> Option { /// /// Never refreshes or rewrites CLI-owned `credentials/kimi-code.json`. /// Skips when `KIMI_CODE_BASE_URL` / OAuth host overrides are set. -pub(crate) fn kimi_code_cli_access_token(now_unix: f64) -> Option { - if has_code_endpoint_override() { +pub(crate) fn kimi_code_cli_access_token(region: KimiRegion, now_unix: f64) -> Option { + if region != KimiRegion::China || has_code_endpoint_override() { return None; } let home = kimi_code_home()?; @@ -417,7 +419,7 @@ mod tests { std::env::set_var(KIMI_CODE_HOME_ENV, home.path()); } - let token = kimi_code_cli_access_token(now); + let token = kimi_code_cli_access_token(KimiRegion::China, now); assert_eq!(token.as_deref(), Some("oauth-token")); let after = std::fs::read(&cred_path).unwrap(); @@ -466,7 +468,7 @@ mod tests { std::env::set_var(KIMI_CODE_BASE_URL_ENV, "https://proxy.example.com/kimi"); } assert!(has_code_endpoint_override()); - assert!(kimi_code_cli_access_token(now).is_none()); + assert!(kimi_code_cli_access_token(KimiRegion::China, now).is_none()); // SAFETY: still under the same env_lock() guard; swapping which // override keys are present between assertions. @@ -474,7 +476,7 @@ mod tests { std::env::remove_var(KIMI_CODE_BASE_URL_ENV); std::env::set_var(KIMI_CODE_OAUTH_HOST_ENV, "https://oauth.example.com"); } - assert!(kimi_code_cli_access_token(now).is_none()); + assert!(kimi_code_cli_access_token(KimiRegion::China, now).is_none()); // SAFETY: final cleanup while the env_lock() guard is still alive. unsafe { diff --git a/rust/src/providers/kimi/desktop_token.rs b/rust/src/providers/kimi/desktop_token.rs index 5c2d6bae42..ac40f793ba 100644 --- a/rust/src/providers/kimi/desktop_token.rs +++ b/rust/src/providers/kimi/desktop_token.rs @@ -29,8 +29,6 @@ const DESKTOP_APP_DIR: &str = "kimi-desktop"; const COOKIES_FILE: &str = "Cookies"; const LOCAL_STATE_FILE: &str = "Local State"; const AUTH_COOKIE_NAME: &str = "kimi-auth"; -const AUTH_COOKIE_HOSTS: [&str; 4] = ["www.kimi.com", ".www.kimi.com", ".kimi.com", "kimi.com"]; - impl KimiDesktopAuthToken { /// Cookies database inside a caller-provided `data_root` (upstream /// `cookiesDatabaseURL(homeDirectory:)` shape for test injection). @@ -47,12 +45,20 @@ impl KimiDesktopAuthToken { /// Desktop session, or `None` when the app/database/cookie is absent or /// unreadable. Production entry point. pub fn load() -> Option { + Self::load_for_region(super::KimiRegion::China) + } + + pub fn load_for_region(region: super::KimiRegion) -> Option { let data_root = dirs::data_dir()?; - Self::load_from(&data_root) + Self::load_from_region(&data_root, region) } /// Read from an explicit `data_root` (Electron `userData` parent). pub fn load_from(data_root: &Path) -> Option { + Self::load_from_region(data_root, super::KimiRegion::China) + } + + pub fn load_from_region(data_root: &Path, region: super::KimiRegion) -> Option { let aes_key = crate::browser::cookies::CookieExtractor::get_chromium_encryption_key( &Self::local_state_path(data_root), ) @@ -63,13 +69,21 @@ impl KimiDesktopAuthToken { ); }) .ok(); - Self::load_token(&Self::cookies_database_path(data_root), aes_key.as_deref()) + Self::load_token( + &Self::cookies_database_path(data_root), + aes_key.as_deref(), + region.desktop_cookie_hosts(), + ) } /// Core read (upstream `read(databaseURL:immutable:)`): WAL-safe /// read-only open → newest `kimi-auth` row → decode. `aes_key` is the /// Chromium app cookie key; `None` restricts reads to plaintext rows. - fn load_token(database_path: &Path, aes_key: Option<&[u8]>) -> Option { + fn load_token( + database_path: &Path, + aes_key: Option<&[u8]>, + hosts: &[&str; 4], + ) -> Option { if !database_path.is_file() { return None; } @@ -81,7 +95,7 @@ impl KimiDesktopAuthToken { tracing::debug!(error = %err, "Kimi Desktop Cookies open failed"); }) .ok()?; - read_newest_auth_cookie(&conn) + read_newest_auth_cookie(&conn, hosts) .inspect_err(|err| { tracing::debug!(error = %err, "Kimi Desktop cookies read failed"); }) @@ -145,7 +159,10 @@ fn decode_cookie_value(row: (String, Vec), aes_key: Option<&[u8]>) -> Option .filter(|plain| !plain.is_empty()) } -fn read_newest_auth_cookie(conn: &rusqlite::Connection) -> rusqlite::Result<(String, Vec)> { +fn read_newest_auth_cookie( + conn: &rusqlite::Connection, + hosts: &[&str; 4], +) -> rusqlite::Result<(String, Vec)> { // Upstream query verbatim: newest `kimi-auth` across the registered // kimi.com cookie scopes by last access. let mut statement = conn.prepare( @@ -157,13 +174,7 @@ fn read_newest_auth_cookie(conn: &rusqlite::Connection) -> rusqlite::Result<(Str LIMIT 1", )?; statement.query_row( - rusqlite::params![ - AUTH_COOKIE_NAME, - AUTH_COOKIE_HOSTS[0], - AUTH_COOKIE_HOSTS[1], - AUTH_COOKIE_HOSTS[2], - AUTH_COOKIE_HOSTS[3], - ], + rusqlite::params![AUTH_COOKIE_NAME, hosts[0], hosts[1], hosts[2], hosts[3],], |row| Ok((row.get::<_, String>(0)?, row.get::<_, Vec>(1)?)), ) } @@ -357,12 +368,24 @@ mod tests { insert_cookie_row(&conn, "www.kimi.com", "", &encrypted, 1); assert_eq!( - KimiDesktopAuthToken::load_token(&database, Some(key.as_slice())).as_deref(), + KimiDesktopAuthToken::load_token( + &database, + Some(key.as_slice()), + crate::providers::KimiRegion::China.desktop_cookie_hosts(), + ) + .as_deref(), Some("encrypted-kimi-token") ); // Without a key the encrypted row cannot be used. - assert_eq!(KimiDesktopAuthToken::load_token(&database, None), None); + assert_eq!( + KimiDesktopAuthToken::load_token( + &database, + None, + crate::providers::KimiRegion::China.desktop_cookie_hosts(), + ), + None + ); // `load_from` without a usable `Local State` reads plaintext only and // yields nothing (no panic, no secret in logs). assert_eq!(KimiDesktopAuthToken::load_from(root.path()), None); diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index d96d3a36fc..e326afa3d6 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -14,8 +14,11 @@ mod code_api; pub mod desktop_token; +mod region; mod web; +pub use region::KimiRegion; + use async_trait::async_trait; use chrono::{DateTime, Utc}; use reqwest::Client; @@ -27,13 +30,11 @@ use crate::core::{ RateWindow, SourceMode, UsageSnapshot, }; -const KIMI_WEB_USAGE_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages"; -const KIMI_SUBSCRIPTION_STATS_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.membership.v2.MembershipService/GetSubscriptionStats"; -const KIMI_SUBSCRIPTION_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.membership.v2.MembershipService/GetSubscription"; -const KIMI_COOKIE_DOMAINS: [&str; 2] = ["www.kimi.com", "kimi.moonshot.cn"]; +const KIMI_WEB_USAGE_SERVICE: &str = "kimi.gateway.billing.v1.BillingService/GetUsages"; +const KIMI_SUBSCRIPTION_STATS_SERVICE: &str = + "kimi.gateway.membership.v2.MembershipService/GetSubscriptionStats"; +const KIMI_SUBSCRIPTION_SERVICE: &str = + "kimi.gateway.membership.v2.MembershipService/GetSubscription"; #[derive(Debug, Deserialize)] struct KimiCodeApiUsageResponse { @@ -330,11 +331,12 @@ impl Provider for KimiProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching Kimi usage"); + let region = KimiRegion::from_settings(ctx.api_region.as_deref()); match ctx.source_mode { SourceMode::Auto => { if code_api::code_api_key(ctx.api_key.as_deref()).is_ok() { - match code_api::fetch_via_code_api(ctx, None, None, "Code API").await { + match code_api::fetch_via_code_api(ctx, region, None, None, "Code API").await { Ok(usage) => { return Ok(ProviderFetchResult::new(usage, "code-api")); } @@ -347,11 +349,14 @@ impl Provider for KimiProvider { } } - if let Some(cli_token) = code_api::kimi_code_cli_access_token(unix_now_secs()) { + if let Some(cli_token) = + code_api::kimi_code_cli_access_token(region, unix_now_secs()) + { let home = code_api::kimi_code_home().unwrap_or_default(); let headers = code_api::kimi_code_cli_identity_headers(&home); match code_api::fetch_via_code_api( ctx, + region, Some(&cli_token), Some(&headers), "Kimi Code CLI", @@ -370,15 +375,16 @@ impl Provider for KimiProvider { } } - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref()).await?; + let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::OAuth => { - let usage = code_api::fetch_via_code_api(ctx, None, None, "Code API").await?; + let usage = + code_api::fetch_via_code_api(ctx, region, None, None, "Code API").await?; Ok(ProviderFetchResult::new(usage, "code-api")) } SourceMode::Web => { - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref()).await?; + let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::Cli => Err(ProviderError::UnsupportedSource(SourceMode::Cli)), @@ -489,6 +495,7 @@ fn is_equivalent_to_weekly_window(window: &RateWindow, weekly: &RateWindow) -> b async fn kimi_web_post( client: &Client, url: &str, + region: KimiRegion, token: &str, body: serde_json::Value, ) -> Result { @@ -498,6 +505,8 @@ async fn kimi_web_post( .header("Cookie", format!("kimi-auth={token}")) .header("Accept", "application/json") .header("Content-Type", "application/json") + .header("Origin", region.web_base_url()) + .header("Referer", region.console_url()) .header( "User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", diff --git a/rust/src/providers/kimi/region.rs b/rust/src/providers/kimi/region.rs new file mode 100644 index 0000000000..2043ce7a93 --- /dev/null +++ b/rust/src/providers/kimi/region.rs @@ -0,0 +1,106 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum KimiRegion { + China, + International, +} + +impl KimiRegion { + pub const ALL: [Self; 2] = [Self::China, Self::International]; + + pub fn from_settings(value: Option<&str>) -> Self { + match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() { + Some("international" | "intl" | "global") => Self::International, + _ => Self::China, + } + } + + pub const fn settings_value(self) -> &'static str { + match self { + Self::China => "china", + Self::International => "international", + } + } + + pub const fn display_name(self) -> &'static str { + match self { + Self::China => "China (kimi.com)", + Self::International => "International (kimi.ai)", + } + } + + pub const fn code_api_base_url(self) -> &'static str { + match self { + Self::China => "https://api.kimi.com", + Self::International => "https://api.kimi.ai", + } + } + + pub const fn web_base_url(self) -> &'static str { + match self { + Self::China => "https://www.kimi.com", + Self::International => "https://www.kimi.ai", + } + } + + pub const fn console_url(self) -> &'static str { + match self { + Self::China => "https://www.kimi.com/code/console", + Self::International => "https://www.kimi.ai/code/console", + } + } + + pub const fn cookie_domains(self) -> &'static [&'static str] { + match self { + Self::China => &["www.kimi.com", "kimi.com"], + Self::International => &["www.kimi.ai", "kimi.ai"], + } + } + + pub const fn desktop_cookie_hosts(self) -> &'static [&'static str; 4] { + match self { + Self::China => &["www.kimi.com", ".www.kimi.com", ".kimi.com", "kimi.com"], + Self::International => &["www.kimi.ai", ".www.kimi.ai", ".kimi.ai", "kimi.ai"], + } + } + + pub fn web_api_url(self, service: &str) -> String { + format!("{}/apiv2/{service}", self.web_base_url()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn unknown_and_missing_settings_preserve_china_default() { + assert_eq!(KimiRegion::from_settings(None), KimiRegion::China); + assert_eq!( + KimiRegion::from_settings(Some("unknown")), + KimiRegion::China + ); + assert_eq!( + KimiRegion::from_settings(Some("international")), + KimiRegion::International + ); + } + + #[test] + fn regional_hosts_remain_coherent() { + for region in KimiRegion::ALL { + let suffix = match region { + KimiRegion::China => "kimi.com", + KimiRegion::International => "kimi.ai", + }; + assert!(region.code_api_base_url().ends_with(suffix)); + assert!(region.web_base_url().ends_with(suffix)); + assert!(region.console_url().ends_with("/code/console")); + assert!( + region + .cookie_domains() + .iter() + .all(|host| host.ends_with(suffix)) + ); + } + } +} diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 4f5289e097..020854b6f4 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -11,9 +11,9 @@ use reqwest::Client; use super::desktop_token::KimiDesktopAuthToken; use super::{ - KIMI_COOKIE_DOMAINS, KIMI_SUBSCRIPTION_STATS_URL, KIMI_SUBSCRIPTION_URL, KIMI_WEB_USAGE_URL, - KimiProvider, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, KimiWebUsageResponse, - apply_subscription_windows, kimi_web_post, + KIMI_SUBSCRIPTION_SERVICE, KIMI_SUBSCRIPTION_STATS_SERVICE, KIMI_WEB_USAGE_SERVICE, + KimiProvider, KimiRegion, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, + KimiWebUsageResponse, apply_subscription_windows, kimi_web_post, }; use crate::browser::cookies::get_cookie_header; use crate::core::{ProviderError, ProviderId, UsageSnapshot}; @@ -37,11 +37,12 @@ fn browser_import_allowed(cookie_source: &str) -> bool { /// 1. Manual cookie header (its `kimi-auth`/auth cookie), source-independent. /// 2. Kimi Desktop session token (automatic source only). /// 3. Browser cookie import (automatic source only). -pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { +pub(crate) fn web_auth_tokens(manual_header: Option<&str>, region: KimiRegion) -> Vec { resolve_web_tokens(WebTokenInput { manual_header, cookie_source: &cookie_source(), - desktop_token: KimiDesktopAuthToken::load, + region, + desktop_token: KimiDesktopAuthToken::load_for_region, browser_token: browser_auth_token, }) .into_iter() @@ -52,8 +53,9 @@ pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { struct WebTokenInput<'a> { manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: fn() -> Option, - browser_token: fn() -> Option, + region: KimiRegion, + desktop_token: fn(KimiRegion) -> Option, + browser_token: fn(KimiRegion) -> Option, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -84,7 +86,7 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { let mut candidates = Vec::new(); let mut seen = std::collections::HashSet::new(); - if let Some(token) = (input.desktop_token)() + if let Some(token) = (input.desktop_token)(input.region) && seen.insert(token.clone()) { candidates.push(WebTokenCandidate { @@ -92,7 +94,7 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { source: WebTokenSource::Desktop, }); } - if let Some(token) = (input.browser_token)() + if let Some(token) = (input.browser_token)(input.region) && seen.insert(token.clone()) { candidates.push(WebTokenCandidate { @@ -105,8 +107,9 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { /// Browser import only: the first usable `kimi-auth`-class token from any of /// the registered Kimi cookie domains. -fn browser_auth_token() -> Option { - KIMI_COOKIE_DOMAINS +fn browser_auth_token(region: KimiRegion) -> Option { + region + .cookie_domains() .iter() .find_map(|domain| { get_cookie_header(domain) @@ -119,6 +122,7 @@ fn browser_auth_token() -> Option { /// Fetch usage via Kimi web API (weekly quota + rate limit + subscription). pub(crate) async fn fetch_via_web( cookie_header: Option<&str>, + region: KimiRegion, ) -> Result { let source = cookie_source(); if let Some(token) = @@ -127,7 +131,7 @@ pub(crate) async fn fetch_via_web( // An explicit manual credential is authoritative. A rejected manual // token must not silently switch accounts underneath the user. let client = client()?; - return fetch_via_web_token(&client, &token).await; + return fetch_via_web_token(&client, &token, region).await; } if !browser_import_allowed(&source) { @@ -143,20 +147,20 @@ pub(crate) async fn fetch_via_web( // Read and try the desktop session first. Browser cookies are intentionally // read only after the server rejects this automatic session, so a healthy // desktop account never causes another credential store to be touched. - if let Some(token) = KimiDesktopAuthToken::load() + if let Some(token) = KimiDesktopAuthToken::load_for_region(region) && seen.insert(token.clone()) { - match fetch_via_web_token(&client, &token).await { + match fetch_via_web_token(&client, &token, region).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(error), } } - if let Some(token) = browser_auth_token() + if let Some(token) = browser_auth_token(region) && seen.insert(token.clone()) { - match fetch_via_web_token(&client, &token).await { + match fetch_via_web_token(&client, &token, region).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(error), @@ -176,10 +180,13 @@ fn client() -> Result { async fn fetch_via_web_token( client: &reqwest::Client, token: &str, + region: KimiRegion, ) -> Result { + let usage_url = region.web_api_url(KIMI_WEB_USAGE_SERVICE); let resp = kimi_web_post( client, - KIMI_WEB_USAGE_URL, + &usage_url, + region, token, serde_json::json!({ "scope": ["FEATURE_CODING"] }), ) @@ -198,7 +205,7 @@ async fn fetch_via_web_token( .await .map_err(|e| ProviderError::Parse(e.to_string()))?; - let (subscription, plan_name) = fetch_subscription_details(client, token).await; + let (subscription, plan_name) = fetch_subscription_details(client, token, region).await; snapshot_from_web_usage_response_with_plan(usage, subscription, plan_name) } @@ -208,16 +215,17 @@ const SUBSCRIPTION_ENRICHMENT_TIMEOUT: std::time::Duration = std::time::Duration async fn fetch_subscription_details( client: &reqwest::Client, token: &str, + region: KimiRegion, ) -> (Option, Option) { // The quota statistics and the optional title are independent. Keep a // completed statistics response when the plan endpoint is slow or absent. let stats = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_for_enrichment(client, token), + fetch_subscription_for_enrichment(client, token, region), ); let plan = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_plan(client, token), + fetch_subscription_plan(client, token, region), ); let (stats, plan) = tokio::join!(stats, plan); (stats.ok().flatten(), plan.ok().flatten()) @@ -260,8 +268,13 @@ fn snapshot_from_web_usage_response_with_plan( Ok(usage) } -pub(super) async fn fetch_subscription_plan(client: &Client, token: &str) -> Option { - match kimi_web_post(client, KIMI_SUBSCRIPTION_URL, token, serde_json::json!({})).await { +pub(super) async fn fetch_subscription_plan( + client: &Client, + token: &str, + region: KimiRegion, +) -> Option { + let url = region.web_api_url(KIMI_SUBSCRIPTION_SERVICE); + match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json::() .await @@ -276,8 +289,9 @@ pub(super) async fn fetch_subscription_plan(client: &Client, token: &str) -> Opt pub(super) async fn fetch_subscription_for_enrichment( client: &Client, token: &str, + region: KimiRegion, ) -> Option { - fetch_subscription_for_enrichment_result(client, token) + fetch_subscription_for_enrichment_result(client, token, region) .await .ok() .flatten() @@ -286,15 +300,10 @@ pub(super) async fn fetch_subscription_for_enrichment( pub(super) async fn fetch_subscription_for_enrichment_result( client: &Client, token: &str, + region: KimiRegion, ) -> Result, ProviderError> { - match kimi_web_post( - client, - KIMI_SUBSCRIPTION_STATS_URL, - token, - serde_json::json!({}), - ) - .await - { + let url = region.web_api_url(KIMI_SUBSCRIPTION_STATS_SERVICE); + match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json() .await @@ -312,33 +321,34 @@ pub(super) async fn fetch_subscription_for_enrichment_result( mod tests { use super::*; - fn static_desktop() -> Option { + fn static_desktop(_: KimiRegion) -> Option { Some("desktop-token".to_string()) } - fn static_browser() -> Option { + fn static_browser(_: KimiRegion) -> Option { Some("browser-token".to_string()) } - fn no_token() -> Option { + fn no_token(_: KimiRegion) -> Option { None } fn input<'a>( manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: fn() -> Option, - browser_token: fn() -> Option, + desktop_token: fn(KimiRegion) -> Option, + browser_token: fn(KimiRegion) -> Option, ) -> WebTokenInput<'a> { WebTokenInput { manual_header, cookie_source, + region: KimiRegion::China, desktop_token, browser_token, } } - fn duplicate_browser() -> Option { + fn duplicate_browser(_: KimiRegion) -> Option { Some("desktop-token".to_string()) } diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 6ac5e0f1a8..a133e1f1f3 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -121,7 +121,7 @@ pub use huggingface::HuggingFaceProvider; pub use infini::InfiniProvider; pub use jetbrains::JetBrainsProvider; pub use kilo::KiloProvider; -pub use kimi::KimiProvider; +pub use kimi::{KimiProvider, KimiRegion}; pub use kimik2::KimiK2Provider; pub use kiro::KiroProvider; pub use litellm::LiteLLMProvider; From b1a619cd0671afd4ba53adb942575753ffadc673 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:56:48 +0700 Subject: [PATCH 09/62] Contain oversized Claude history values --- rust/src/cost_scanner.rs | 128 ++++++++++++++++++------ rust/src/cost_scanner/claude_pricing.rs | 114 ++++++++++++++++++--- rust/src/cost_scanner/tests.rs | 96 +++++++++++++++++- 3 files changed, 291 insertions(+), 47 deletions(-) diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 619f9ba7b0..2b730cdc01 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -435,7 +435,9 @@ struct ClaudeUsageRecord { output: u64, cache_create: u64, cache_read: u64, - cost: f64, + /// `None` means pricing was unavailable or produced a non-finite value. + /// Keep that distinct from a real zero-dollar row. + cost: Option, } /// Exact Claude request rows retained for quota-window projection. @@ -468,6 +470,7 @@ struct ClaudeFileScanResult { malformed_lines: u32, incomplete_requests: u32, read_failures: u32, + aggregation_failures: u32, } impl ClaudeFileScanResult { @@ -478,10 +481,16 @@ impl ClaudeFileScanResult { .incomplete_requests .saturating_add(other.incomplete_requests); self.read_failures = self.read_failures.saturating_add(other.read_failures); + self.aggregation_failures = self + .aggregation_failures + .saturating_add(other.aggregation_failures); } fn is_complete(self) -> bool { - self.malformed_lines == 0 && self.incomplete_requests == 0 && self.read_failures == 0 + self.malformed_lines == 0 + && self.incomplete_requests == 0 + && self.read_failures == 0 + && self.aggregation_failures == 0 } } @@ -591,16 +600,21 @@ impl CostScanner { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let mut handle_file = |path: &Path| { - let file_result = scan_claude_file_with_pricing( + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( path, &cutoff, &mut seen, cancel, &mut pricing, |record| { - add_claude_record_to_summary(&mut summary, record); + aggregation_complete &= add_claude_record_to_summary(&mut summary, record); }, ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } if file_result.counted > 0 { summary.sessions_count += 1; } @@ -667,7 +681,8 @@ impl CostScanner { let mut pricing = ClaudeScanPricingResolver::default(); self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut |path| { let mut file_has_usage = false; - let file_result = scan_claude_file_with_pricing( + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( path, &cutoff, &mut seen, @@ -675,9 +690,11 @@ impl CostScanner { &mut pricing, |record| { file_has_usage = true; - add_claude_record_to_summary(&mut summary, record); - add_claude_record_to_daily_costs(&mut daily_cost, record); - add_claude_record_to_daily_tokens(&mut daily_tokens, record); + aggregation_complete &= add_claude_record_to_summary(&mut summary, record); + aggregation_complete &= + add_claude_record_to_daily_costs(&mut daily_cost, record); + aggregation_complete &= + add_claude_record_to_daily_tokens(&mut daily_tokens, record); if let Some(quota_record) = quota_history_record_from_usage(record) { quota_records.push(quota_record); } else { @@ -685,6 +702,10 @@ impl CostScanner { } }, ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } if file_has_usage { summary.sessions_count += 1; } @@ -1007,7 +1028,7 @@ fn claude_usage_record_from_event_with_pricing( let cache_create_1h = usage.one_hour_cache_creation_tokens(cache_create); let pricing_known = pricing.is_known(model); - let cost = pricing.cost_usd_with_cache_ttl( + let computed_cost = pricing.cost_usd_with_cache_ttl( model, input, cache_create, @@ -1015,6 +1036,7 @@ fn claude_usage_record_from_event_with_pricing( cache_read, output, ); + let cost = computed_cost.is_finite().then_some(computed_cost); Some(ClaudeUsageRecord { model: model.to_string(), @@ -1033,25 +1055,52 @@ fn claude_usage_record_from_event_with_pricing( }) } -fn add_claude_record_to_summary(summary: &mut CostSummary, record: &ClaudeUsageRecord) { +fn add_claude_record_to_summary(summary: &mut CostSummary, record: &ClaudeUsageRecord) -> bool { if !record.pricing_known { summary.unknown_models.insert(record.model.clone()); } - summary.input_tokens += record.input; - summary.output_tokens += record.output; - summary.cached_tokens += record.cache_create + record.cache_read; - summary.total_cost_usd += record.cost; + let mut complete = checked_add_assign(&mut summary.input_tokens, record.input); + complete &= checked_add_assign(&mut summary.output_tokens, record.output); + let cached = record.cache_create.checked_add(record.cache_read); + complete &= cached.is_some_and(|value| checked_add_assign(&mut summary.cached_tokens, value)); - *summary.by_model.entry(record.model.clone()).or_insert(0.0) += record.cost; + if let Some(cost) = record.cost { + complete &= checked_add_finite(&mut summary.total_cost_usd, cost); + complete &= checked_add_finite( + summary.by_model.entry(record.model.clone()).or_insert(0.0), + cost, + ); + } else { + complete = false; + } let model_tokens = summary .by_model_tokens .entry(record.model.clone()) .or_default(); - model_tokens.input_tokens += record.input; - model_tokens.output_tokens += record.output; - model_tokens.cached_tokens += record.cache_create + record.cache_read; + complete &= checked_add_assign(&mut model_tokens.input_tokens, record.input); + complete &= checked_add_assign(&mut model_tokens.output_tokens, record.output); + complete &= + cached.is_some_and(|value| checked_add_assign(&mut model_tokens.cached_tokens, value)); + complete +} + +fn checked_add_assign(total: &mut u64, value: u64) -> bool { + let Some(sum) = total.checked_add(value) else { + return false; + }; + *total = sum; + true +} + +fn checked_add_finite(total: &mut f64, value: f64) -> bool { + let sum = *total + value; + if !value.is_finite() || !sum.is_finite() { + return false; + } + *total = sum; + true } fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option { @@ -1080,9 +1129,9 @@ fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option= 0.0, + cost_is_complete: record.pricing_known && record.cost.is_some_and(|cost| cost >= 0.0), dedup_key, attribution: ClaudeHistoryAttribution::Unavailable, }) @@ -1094,9 +1143,9 @@ fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option>, record: &ClaudeUsageRecord, -) { +) -> bool { let Some(timestamp) = record.timestamp else { - return; + return true; }; let date_str = timestamp .with_timezone(&Local) @@ -1104,8 +1153,18 @@ fn add_claude_record_to_daily_costs( .format("%Y-%m-%d") .to_string(); if let Some(cost) = daily_costs.get_mut(&date_str) { - *cost = Some(cost.unwrap_or(0.0) + record.cost); + let Some(record_cost) = record.cost else { + *cost = None; + return false; + }; + let sum = cost.unwrap_or(0.0) + record_cost; + if !sum.is_finite() { + *cost = None; + return false; + } + *cost = Some(sum); } + true } /// Check if any cost usage sources are available @@ -1294,6 +1353,7 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); + let mut aggregation_complete = true; let mut handle_file = |path: &Path| { for_each_claude_usage_record_with_pricing( path, @@ -1302,11 +1362,17 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> None, &mut pricing, |record| { - add_claude_record_to_daily_tokens(&mut daily_tokens, record); + aggregation_complete &= + add_claude_record_to_daily_tokens(&mut daily_tokens, record); }, ); }; scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + if !aggregation_complete { + covered_days.clear(); + } else { + covered_days.extend(daily_tokens.keys().cloned()); + } } } "pi" => { @@ -1330,7 +1396,9 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> // Codex only: the bounded catch-up may not have reached the requested // depth yet. Incomplete = history exists but the oldest quarter of the // window has no scanned day. - let incomplete = if provider == "pi" { + let incomplete = if provider == "claude" { + covered_days.is_empty() + } else if provider == "pi" { // Pi scans are bounded filesystem walks, so a complete parse covers // the requested window even when the roots contain no sessions. covered_days.is_empty() @@ -1349,9 +1417,9 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> fn add_claude_record_to_daily_tokens( daily_tokens: &mut HashMap, record: &ClaudeUsageRecord, -) { +) -> bool { let Some(timestamp) = record.timestamp else { - return; + return true; }; let date_str = timestamp .with_timezone(&Local) @@ -1359,6 +1427,10 @@ fn add_claude_record_to_daily_tokens( .format("%Y-%m-%d") .to_string(); if let Some(slot) = daily_tokens.get_mut(&date_str) { - *slot += record.input + record.output; + let Some(tokens) = record.input.checked_add(record.output) else { + return false; + }; + return checked_add_assign(slot, tokens); } + true } diff --git a/rust/src/cost_scanner/claude_pricing.rs b/rust/src/cost_scanner/claude_pricing.rs index 89bcbe138a..e8f3e66973 100644 --- a/rust/src/cost_scanner/claude_pricing.rs +++ b/rust/src/cost_scanner/claude_pricing.rs @@ -139,24 +139,10 @@ impl ClaudeScanPricingResolver { let cache_create_1h = cache_create_1h.min(cache_create); let cache_create_5m = cache_create.saturating_sub(cache_create_1h); - #[allow( - clippy::cast_possible_truncation, - reason = "clamped to i32::MAX before casting" - )] - let clamp = |value: u64| value.min(i32::MAX as u64) as i32; - let resolved = self.resolve(model); let billable = resolved.or_else(|| self.resolve(FALLBACK_CLAUDE_MODEL)); let base = billable - .map(|pricing| { - CostUsagePricing::claude_cost_usd_from_resolution( - pricing, - clamp(input), - clamp(cache_read), - clamp(cache_create_5m), - clamp(output), - ) - }) + .map(|pricing| claude_cost_usd_u64(pricing, input, cache_read, cache_create_5m, output)) .unwrap_or(0.0); let input_rate = billable .map(CostUsagePricing::claude_input_cost_per_token_from_resolution) @@ -165,3 +151,101 @@ impl ClaudeScanPricingResolver { base + (cache_create_1h as f64) * input_rate * 2.0 } } + +/// Price transcript counters without narrowing them to `i32`. Local history is +/// untrusted input and can contain values far above the API's ordinary range; +/// narrowing those values silently understates spend before aggregation gets a +/// chance to mark non-finite results unavailable. +fn claude_cost_usd_u64( + resolution: ClaudePricingResolution, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> f64 { + match resolution { + ClaudePricingResolution::BuiltIn(pricing) => { + let tiered = |tokens: u64, base: f64, above: Option| { + let Some(threshold) = pricing.threshold_tokens.map(|value| value.max(0) as u64) + else { + return (tokens as f64) * base; + }; + let Some(above) = above else { + return (tokens as f64) * base; + }; + let below = tokens.min(threshold); + let over = tokens.saturating_sub(threshold); + (below as f64) * base + (over as f64) * above + }; + + tiered( + input, + pricing.input_cost_per_token, + pricing.input_cost_per_token_above_threshold, + ) + tiered( + cache_read, + pricing.cache_read_input_cost_per_token, + pricing.cache_read_input_cost_per_token_above_threshold, + ) + tiered( + cache_write, + pricing.cache_creation_input_cost_per_token, + pricing.cache_creation_input_cost_per_token_above_threshold, + ) + tiered( + output, + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + ) + } + ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens, + } => { + let use_tier = threshold_tokens.is_some_and(|threshold| { + input + .checked_add(cache_read) + .and_then(|value| value.checked_add(cache_write)) + .is_none_or(|total| total > threshold) + }); + let pick = |base: f64, above: Option| { + if use_tier { + above.unwrap_or(base) + } else { + base + } + }; + let input_rate = pick( + pricing.input_cost_per_token, + pricing.input_cost_per_token_above_threshold, + ); + let cache_read_rate = if use_tier { + pricing + .cache_read_input_cost_per_token_above_threshold + .or(pricing.cache_read_input_cost_per_token) + .unwrap_or(input_rate) + } else { + pricing + .cache_read_input_cost_per_token + .unwrap_or(input_rate) + }; + let cache_write_rate = if use_tier { + pricing + .cache_write_input_cost_per_token_above_threshold + .or(pricing.cache_write_input_cost_per_token) + .unwrap_or(input_rate) + } else { + pricing + .cache_write_input_cost_per_token + .unwrap_or(input_rate) + }; + let output_rate = pick( + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + ); + + (input as f64) * input_rate + + (cache_read as f64) * cache_read_rate + + (cache_write as f64) * cache_write_rate + + (output as f64) * output_rate + } + } +} diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 066e8e2d98..d20e770b26 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -395,7 +395,7 @@ fn counts_claude_usage_once_across_duplicate_records() { assert_eq!(record.output, 50); assert_eq!(record.cache_create, 10); assert_eq!(record.cache_read, 20); - assert!(record.cost > 0.0); + assert!(record.cost.is_some_and(|cost| cost > 0.0)); let cutoff = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") .unwrap() @@ -552,9 +552,17 @@ fn shared_claude_reader_excludes_vertex_rows_but_keeps_anthropic_usage() { let anthropic = format!( r#"{{"type":"assistant","timestamp":"{timestamp}","requestId":"req_anthropic","message":{{"id":"msg_anthropic","model":"claude-sonnet-4-6","usage":{{"input_tokens":10,"output_tokens":5}}}}}}"# ); - let vertex = format!( - r#"{{"type":"assistant","timestamp":"{timestamp}","requestId":"req_vrtx_123","message":{{"id":"msg_vrtx_123","model":"claude-sonnet-4-6","usage":{{"input_tokens":1000,"output_tokens":500}}}}}}"# - ); + let vertex = serde_json::json!({ + "type": "assistant", + "timestamp": timestamp, + "requestId": "req_vrtx_123", + "message": { + "id": "msg_vrtx_123", + "model": "claude-sonnet-4-6", + "usage": {"input_tokens": u64::MAX, "output_tokens": u64::MAX} + } + }) + .to_string(); std::fs::write(&path, format!("{anthropic}\n{vertex}\n")).unwrap(); let cutoff = Utc::now() - Duration::days(30); @@ -569,6 +577,86 @@ fn shared_claude_reader_excludes_vertex_rows_but_keeps_anthropic_usage() { let _removed = std::fs::remove_file(&path); } +#[test] +fn oversized_claude_history_preserves_independent_components_and_fails_closed() { + let first: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_overflow_1","message":{{"id":"msg_overflow_1","model":"claude-sonnet-4-6","usage":{{"input_tokens":{},"output_tokens":2}}}}}}"#, + u64::MAX + )) + .unwrap(); + let second: ClaudeEvent = serde_json::from_str( + r#"{"type":"assistant","timestamp":"2026-09-20T12:01:00Z","requestId":"req_overflow_2","message":{"id":"msg_overflow_2","model":"claude-sonnet-4-6","usage":{"input_tokens":1,"output_tokens":3}}}"#, + ) + .unwrap(); + let first = claude_usage_record_from_event(&first).expect("first usage row"); + let second = claude_usage_record_from_event(&second).expect("second usage row"); + let mut summary = CostSummary::default(); + + assert!(add_claude_record_to_summary(&mut summary, &first)); + assert!(!add_claude_record_to_summary(&mut summary, &second)); + assert_eq!(summary.input_tokens, u64::MAX); + assert_eq!(summary.output_tokens, 5); + assert!(summary.total_cost_usd.is_finite()); + + finalize_claude_summary( + &mut summary, + true, + ClaudeFileScanResult { + counted: 2, + aggregation_failures: 1, + ..ClaudeFileScanResult::default() + }, + false, + ); + assert!(!summary.history_coverage_established); + assert!(!summary.known_zero); +} + +#[test] +fn oversized_single_claude_row_keeps_cost_but_marks_combined_quota_tokens_unknown() { + let event: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_combined_overflow","message":{{"id":"msg_combined_overflow","model":"claude-sonnet-4-6","usage":{{"input_tokens":{},"output_tokens":1}}}}}}"#, + u64::MAX + )) + .unwrap(); + let record = claude_usage_record_from_event(&event).expect("usage row"); + let quota = quota_history_record_from_usage(&record).expect("timestamped quota row"); + + assert!(record.cost.is_some_and(f64::is_finite)); + assert_eq!(quota.tokens, None); + assert!(!quota.tokens_are_complete); + assert!(quota.cost_usd.is_some_and(f64::is_finite)); + assert!(quota.cost_is_complete); +} + +#[test] +fn nonfinite_claude_price_is_unknown_instead_of_zero() { + let snapshot = crate::core::ModelsDevPricingSnapshot::from_catalog_json_for_tests( + r#"{ + "anthropic": {"models": {"claude-test-extreme-price": { + "id": "claude-test-extreme-price", "cost": {"input": 1e308, "output": 1} + }}} + }"#, + ) + .expect("pricing fixture"); + let mut pricing = ClaudeScanPricingResolver::with_snapshot(snapshot); + let event: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_nonfinite","message":{{"id":"msg_nonfinite","model":"claude-test-extreme-price","usage":{{"input_tokens":{},"output_tokens":1}}}}}}"#, + u64::MAX + )) + .unwrap(); + let record = + claude_usage_record_from_event_with_pricing(&event, &mut pricing).expect("usage row"); + let mut summary = CostSummary::default(); + + assert_eq!(record.cost, None); + assert!(!add_claude_record_to_summary(&mut summary, &record)); + assert_eq!(summary.input_tokens, u64::MAX); + assert_eq!(summary.output_tokens, 1); + assert_eq!(summary.total_cost_usd, 0.0); + assert!(!summary.known_zero); +} + fn claude_transcript_line( timestamp: &str, request_key: &str, From 6d5763ace04c95fb44bf5cf6c922bd6644f239e6 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:44:35 +0700 Subject: [PATCH 10/62] Centralize unsigned Claude pricing --- rust/src/core/cost_pricing/claude.rs | 80 +++++++++++++++--- rust/src/cost_scanner/claude_pricing.rs | 108 ++---------------------- 2 files changed, 78 insertions(+), 110 deletions(-) diff --git a/rust/src/core/cost_pricing/claude.rs b/rust/src/core/cost_pricing/claude.rs index 8cdab4f34d..05429d57a9 100644 --- a/rust/src/core/cost_pricing/claude.rs +++ b/rust/src/core/cost_pricing/claude.rs @@ -86,20 +86,38 @@ impl CostUsagePricing { cache_read_input_tokens: i32, cache_creation_input_tokens: i32, output_tokens: i32, + ) -> f64 { + Self::claude_cost_usd_u64_from_resolution( + resolution, + u64::try_from(input_tokens).unwrap_or(0), + u64::try_from(cache_read_input_tokens).unwrap_or(0), + u64::try_from(cache_creation_input_tokens).unwrap_or(0), + u64::try_from(output_tokens).unwrap_or(0), + ) + } + + /// Calculate cost from a resolved Claude pricing source without narrowing + /// untrusted local-history counters to the API-oriented signed type. + pub(crate) fn claude_cost_usd_u64_from_resolution( + resolution: ClaudePricingResolution, + input_tokens: u64, + cache_read_input_tokens: u64, + cache_creation_input_tokens: u64, + output_tokens: u64, ) -> f64 { match resolution { ClaudePricingResolution::BuiltIn(pricing) => { fn tiered( - tokens: i32, + tokens: u64, base: f64, above: Option, threshold: Option, ) -> f64 { - let tokens = tokens.max(0); match (threshold, above) { (Some(thresh), Some(above_rate)) => { + let thresh = u64::try_from(thresh).unwrap_or(0); let below = tokens.min(thresh); - let over = (tokens - thresh).max(0); + let over = tokens.saturating_sub(thresh); (below as f64) * base + (over as f64) * above_rate } _ => (tokens as f64) * base, @@ -131,14 +149,54 @@ impl CostUsagePricing { ClaudePricingResolution::ModelsDev { pricing, threshold_tokens, - } => claude_routed_pricing::cost_usd_from_pricing_with_threshold( - pricing, - threshold_tokens, - input_tokens, - cache_read_input_tokens, - cache_creation_input_tokens, - output_tokens, - ), + } => { + let use_tier = threshold_tokens.is_some_and(|threshold| { + input_tokens + .checked_add(cache_read_input_tokens) + .and_then(|value| value.checked_add(cache_creation_input_tokens)) + .is_none_or(|total| total > threshold) + }); + let pick = |base: f64, above: Option| { + if use_tier { + above.unwrap_or(base) + } else { + base + } + }; + let input_rate = pick( + pricing.input_cost_per_token, + pricing.input_cost_per_token_above_threshold, + ); + let cache_read_rate = if use_tier { + pricing + .cache_read_input_cost_per_token_above_threshold + .or(pricing.cache_read_input_cost_per_token) + .unwrap_or(input_rate) + } else { + pricing + .cache_read_input_cost_per_token + .unwrap_or(input_rate) + }; + let cache_write_rate = if use_tier { + pricing + .cache_write_input_cost_per_token_above_threshold + .or(pricing.cache_write_input_cost_per_token) + .unwrap_or(input_rate) + } else { + pricing + .cache_write_input_cost_per_token + .unwrap_or(input_rate) + }; + let output_rate = pick( + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + ); + + (input_tokens as f64) * input_rate + + (cache_read_input_tokens as f64) * cache_read_rate + + (cache_creation_input_tokens as f64) * cache_write_rate + + (output_tokens as f64) * output_rate + } } } diff --git a/rust/src/cost_scanner/claude_pricing.rs b/rust/src/cost_scanner/claude_pricing.rs index e8f3e66973..5d3d251036 100644 --- a/rust/src/cost_scanner/claude_pricing.rs +++ b/rust/src/cost_scanner/claude_pricing.rs @@ -142,7 +142,15 @@ impl ClaudeScanPricingResolver { let resolved = self.resolve(model); let billable = resolved.or_else(|| self.resolve(FALLBACK_CLAUDE_MODEL)); let base = billable - .map(|pricing| claude_cost_usd_u64(pricing, input, cache_read, cache_create_5m, output)) + .map(|pricing| { + CostUsagePricing::claude_cost_usd_u64_from_resolution( + pricing, + input, + cache_read, + cache_create_5m, + output, + ) + }) .unwrap_or(0.0); let input_rate = billable .map(CostUsagePricing::claude_input_cost_per_token_from_resolution) @@ -151,101 +159,3 @@ impl ClaudeScanPricingResolver { base + (cache_create_1h as f64) * input_rate * 2.0 } } - -/// Price transcript counters without narrowing them to `i32`. Local history is -/// untrusted input and can contain values far above the API's ordinary range; -/// narrowing those values silently understates spend before aggregation gets a -/// chance to mark non-finite results unavailable. -fn claude_cost_usd_u64( - resolution: ClaudePricingResolution, - input: u64, - cache_read: u64, - cache_write: u64, - output: u64, -) -> f64 { - match resolution { - ClaudePricingResolution::BuiltIn(pricing) => { - let tiered = |tokens: u64, base: f64, above: Option| { - let Some(threshold) = pricing.threshold_tokens.map(|value| value.max(0) as u64) - else { - return (tokens as f64) * base; - }; - let Some(above) = above else { - return (tokens as f64) * base; - }; - let below = tokens.min(threshold); - let over = tokens.saturating_sub(threshold); - (below as f64) * base + (over as f64) * above - }; - - tiered( - input, - pricing.input_cost_per_token, - pricing.input_cost_per_token_above_threshold, - ) + tiered( - cache_read, - pricing.cache_read_input_cost_per_token, - pricing.cache_read_input_cost_per_token_above_threshold, - ) + tiered( - cache_write, - pricing.cache_creation_input_cost_per_token, - pricing.cache_creation_input_cost_per_token_above_threshold, - ) + tiered( - output, - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - ) - } - ClaudePricingResolution::ModelsDev { - pricing, - threshold_tokens, - } => { - let use_tier = threshold_tokens.is_some_and(|threshold| { - input - .checked_add(cache_read) - .and_then(|value| value.checked_add(cache_write)) - .is_none_or(|total| total > threshold) - }); - let pick = |base: f64, above: Option| { - if use_tier { - above.unwrap_or(base) - } else { - base - } - }; - let input_rate = pick( - pricing.input_cost_per_token, - pricing.input_cost_per_token_above_threshold, - ); - let cache_read_rate = if use_tier { - pricing - .cache_read_input_cost_per_token_above_threshold - .or(pricing.cache_read_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_read_input_cost_per_token - .unwrap_or(input_rate) - }; - let cache_write_rate = if use_tier { - pricing - .cache_write_input_cost_per_token_above_threshold - .or(pricing.cache_write_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_write_input_cost_per_token - .unwrap_or(input_rate) - }; - let output_rate = pick( - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - ); - - (input as f64) * input_rate - + (cache_read as f64) * cache_read_rate - + (cache_write as f64) * cache_write_rate - + (output as f64) * output_rate - } - } -} From f1b291843793bd37d0694f2ef6cd134e52b29b3e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:21:37 +0700 Subject: [PATCH 11/62] Centralize provider dashboard routing --- .../src-tauri/src/commands/mod.rs | 19 +++++++++ .../src-tauri/src/commands/provider_detail.rs | 16 +------- .../src-tauri/src/commands/system.rs | 41 +------------------ .../src-tauri/src/commands/tests.rs | 16 ++++++++ 4 files changed, 38 insertions(+), 54 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index b7d71cf937..14472b82cd 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -98,6 +98,25 @@ fn canonical_provider_arg(provider_id: &str) -> Result { Ok(parse_provider_arg(provider_id)?.cli_name().to_string()) } +fn provider_dashboard_url(id: ProviderId, settings: &Settings) -> Option { + match id { + ProviderId::MiniMax => Some( + codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( + settings.api_region(id), + )), + ), + ProviderId::Kimi => Some( + codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) + .console_url() + .to_string(), + ), + _ => instantiate_provider(id) + .metadata() + .dashboard_url + .map(str::to_string), + } +} + fn validate_single_line_secret(value: &str, field: &str, max_len: usize) -> Result<(), String> { let trimmed = value.trim(); if trimmed.is_empty() { diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs index a061856da7..6feeea4868 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs @@ -73,21 +73,7 @@ pub(crate) fn build_provider_detail( let provider = instantiate_provider(id); let metadata = provider.metadata(); let resume_supported = auto_resume_supported(id); - let dashboard_url = if id == codexbar::core::ProviderId::MiniMax { - Some( - codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( - settings.api_region(id), - )), - ) - } else if id == codexbar::core::ProviderId::Kimi { - Some( - codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) - .console_url() - .to_string(), - ) - } else { - metadata.dashboard_url.map(|s| s.to_string()) - }; + let dashboard_url = provider_dashboard_url(id, &settings); let detail = ProviderDetail { id: id.cli_name().to_string(), diff --git a/apps/desktop-tauri/src-tauri/src/commands/system.rs b/apps/desktop-tauri/src-tauri/src/commands/system.rs index 104018628e..d7f56b124d 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/system.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/system.rs @@ -205,46 +205,9 @@ pub fn quit_app(app: tauri::AppHandle) { } fn dashboard_url_for_provider(provider_id: &str) -> Option { - if provider_id == ProviderId::Kimi.cli_name() { - let settings = Settings::load(); - return Some( - codexbar::providers::KimiRegion::from_settings(Some( - settings.api_region(ProviderId::Kimi), - )) - .console_url() - .to_string(), - ); - } - if provider_id == ProviderId::MiniMax.cli_name() { - let settings = Settings::load(); - return Some( - codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( - settings.api_region(ProviderId::MiniMax), - )), - ); - } - - // OpenRouter's Usage Dashboard is the Activity page. Resolve it from the - // provider metadata before the legacy API-key catalog entry, which still - // points at the credits settings page. - if provider_id == ProviderId::OpenRouter.cli_name() { - return instantiate_provider(ProviderId::OpenRouter) - .metadata() - .dashboard_url - .map(|s| s.to_string()); - } - - if let Some(url) = codexbar::settings::get_api_key_providers() - .into_iter() - .find(|p| p.id.cli_name() == provider_id) - .and_then(|p| p.dashboard_url.map(|s| s.to_string())) - { - return Some(url); - } - let id = ProviderId::from_cli_name(provider_id)?; - let provider = instantiate_provider(id); - provider.metadata().dashboard_url.map(|s| s.to_string()) + let settings = Settings::load(); + provider_dashboard_url(id, &settings) } fn status_page_url_for_provider(provider_id: &str) -> Option { diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index f9f615463b..685a2b2cb7 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -704,6 +704,22 @@ fn fetch_context_includes_minimax_region() { assert_eq!(ctx.api_region.as_deref(), Some("cn")); } +#[test] +fn provider_dashboard_url_uses_selected_regional_console() { + let mut settings = Settings::default(); + settings.set_api_region(ProviderId::MiniMax, "cn"); + settings.set_api_region(ProviderId::Kimi, "international"); + + assert_eq!( + super::provider_dashboard_url(ProviderId::MiniMax, &settings).as_deref(), + Some("https://platform.minimaxi.com/user-center/payment/coding-plan?cycle_type=3") + ); + assert_eq!( + super::provider_dashboard_url(ProviderId::Kimi, &settings).as_deref(), + Some("https://www.kimi.ai/code/console") + ); +} + #[test] fn fetch_context_token_account_uses_web_cookie_header() { let settings = Settings::default(); From 78c0c9556c2123b1d6e26f3ec94530af2ef3193e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:29:31 +0700 Subject: [PATCH 12/62] Add stacked provider tray mode --- .../src-tauri/src/commands/bridge.rs | 5 + .../src-tauri/src/commands/settings.rs | 46 +++++ .../src-tauri/src/tray_bridge.rs | 180 ++++++++++++++++-- apps/desktop-tauri/src/i18n/keys.ts | 3 + apps/desktop-tauri/src/styles.css | 25 --- apps/desktop-tauri/src/surfaces/Settings.tsx | 8 +- .../sections/AccentColorSection.test.tsx | 21 ++ .../providers/sections/AccentColorSection.tsx | 10 - .../settings/tabs/DisplayTab.test.tsx | 28 +++ .../src/surfaces/settings/tabs/DisplayTab.tsx | 55 +++++- apps/desktop-tauri/src/types/bridge.ts | 6 +- rust/src/locale.rs | 3 + rust/src/locale/en-US.ftl | 3 + rust/src/locale/es-MX.ftl | 3 + rust/src/locale/ja-JP.ftl | 3 + rust/src/locale/ko-KR.ftl | 3 + rust/src/locale/ru-RU.ftl | 3 + rust/src/locale/tr-TR.ftl | 3 + rust/src/locale/zh-CN.ftl | 3 + rust/src/locale/zh-TW.ftl | 3 + rust/src/settings.rs | 12 ++ rust/src/settings/raw.rs | 6 + rust/src/settings/tests.rs | 21 ++ rust/src/settings/types.rs | 4 + rust/src/tray/mod.rs | 4 +- rust/src/tray/render.rs | 81 ++++++++ 26 files changed, 490 insertions(+), 52 deletions(-) create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ab5efbef8a..96955cf740 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -661,6 +661,8 @@ pub struct SettingsSnapshot { predictive_pace_warning_enabled: bool, show_pace: bool, tray_icon_mode: &'static str, + stacked_tray_top_provider: Option, + stacked_tray_bottom_provider: Option, switcher_shows_icons: bool, menu_bar_shows_highest_usage: bool, menu_bar_shows_percent: bool, @@ -783,6 +785,8 @@ impl From for SettingsSnapshot { predictive_pace_warning_enabled: settings.predictive_pace_warning_enabled, show_pace: settings.show_pace, tray_icon_mode: tray_icon_mode_label(settings.tray_icon_mode), + stacked_tray_top_provider: settings.stacked_tray_top_provider, + stacked_tray_bottom_provider: settings.stacked_tray_bottom_provider, switcher_shows_icons: settings.switcher_shows_icons, menu_bar_shows_highest_usage: settings.menu_bar_shows_highest_usage, menu_bar_shows_percent: settings.menu_bar_shows_percent, @@ -876,6 +880,7 @@ fn tray_icon_mode_label(mode: TrayIconMode) -> &'static str { match mode { TrayIconMode::Single => "single", TrayIconMode::PerProvider => "perProvider", + TrayIconMode::Stacked => "stacked", } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..b687c9fd48 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -26,6 +26,8 @@ pub struct SettingsUpdate { pub predictive_pace_warning_enabled: Option, pub show_pace: Option, pub tray_icon_mode: Option, + pub stacked_tray_top_provider: Option, + pub stacked_tray_bottom_provider: Option, pub switcher_shows_icons: Option, pub menu_bar_shows_highest_usage: Option, pub menu_bar_shows_percent: Option, @@ -121,6 +123,8 @@ impl SettingsUpdate { fn refreshes_tray_presentation(&self) -> bool { self.tray_icon_mode.is_some() + || self.stacked_tray_top_provider.is_some() + || self.stacked_tray_bottom_provider.is_some() || self.switcher_shows_icons.is_some() || self.menu_bar_shows_highest_usage.is_some() || self.menu_bar_shows_percent.is_some() @@ -190,6 +194,12 @@ impl SettingsUpdate { { settings.tray_icon_mode = mode; } + if let Some(provider) = self.stacked_tray_top_provider.clone() { + settings.stacked_tray_top_provider = normalize_optional_provider_id(provider); + } + if let Some(provider) = self.stacked_tray_bottom_provider.clone() { + settings.stacked_tray_bottom_provider = normalize_optional_provider_id(provider); + } if let Some(v) = self.provider_metrics.clone() { apply_provider_metrics(settings, v); } @@ -482,10 +492,16 @@ fn parse_tray_icon_mode(s: &str) -> Option { match s { "single" => Some(TrayIconMode::Single), "perProvider" => Some(TrayIconMode::PerProvider), + "stacked" => Some(TrayIconMode::Stacked), _ => None, } } +fn normalize_optional_provider_id(value: String) -> Option { + let value = value.trim(); + (!value.is_empty()).then(|| value.to_string()) +} + fn parse_update_channel(s: &str) -> Option { match s { "stable" => Some(UpdateChannel::Stable), @@ -700,6 +716,36 @@ mod tests { } .refreshes_tray_presentation() ); + assert!( + SettingsUpdate { + stacked_tray_top_provider: Some("claude".to_string()), + ..Default::default() + } + .refreshes_tray_presentation() + ); + } + + #[test] + fn stacked_tray_update_accepts_mode_and_clears_automatic_provider() { + let mut settings = Settings { + stacked_tray_top_provider: Some("codex".to_string()), + ..Settings::default() + }; + + SettingsUpdate { + tray_icon_mode: Some("stacked".to_string()), + stacked_tray_top_provider: Some(String::new()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Default::default() + } + .apply_provider_settings(&mut settings); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); + assert_eq!(settings.stacked_tray_top_provider, None); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("claude") + ); } #[test] diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index aebb7c2115..d5e667b68e 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -12,7 +12,9 @@ use tauri::menu::{CheckMenuItemBuilder, IsMenuItem, Menu, MenuItem, PredefinedMe use tauri::tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}; use tauri::{AppHandle, Manager}; -use codexbar::tray::{render_bar_icon_rgba, render_percent_icon_rgba}; +use codexbar::tray::{ + render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, +}; use crate::shell; use crate::state::{AppState, TrayAnchor}; @@ -480,15 +482,27 @@ pub fn update_tray_icon_and_tooltip( let picked = pick_tray_provider(&ok_snapshots, prefer_highest); - let (session_pct, weekly_pct) = match picked { - Some(s) => selected_tray_percents(s, &settings), - None => ( - ok_snapshots - .iter() - .map(|s| selected_tray_percents(s, &settings).0) - .fold(0.0_f64, f64::max), - None, - ), + let (session_pct, weekly_pct) = if settings.tray_icon_mode == TrayIconMode::Stacked { + pick_stacked_tray_providers(&ok_snapshots, &settings) + .map(|(top, bottom)| { + ( + selected_tray_percents(top, &settings).0, + Some(selected_tray_percents(bottom, &settings).0), + ) + }) + .or_else(|| picked.map(|snapshot| selected_tray_percents(snapshot, &settings))) + .unwrap_or((0.0, None)) + } else { + match picked { + Some(s) => selected_tray_percents(s, &settings), + None => ( + ok_snapshots + .iter() + .map(|s| selected_tray_percents(s, &settings).0) + .fold(0.0_f64, f64::max), + None, + ), + } }; let (rgba, w, h) = render_tray_icon_for_settings(&settings, session_pct, weekly_pct, all_error); @@ -517,6 +531,23 @@ fn status_labels_for_settings( .collect::>(); } + if settings.tray_icon_mode == TrayIconMode::Stacked { + return pick_stacked_tray_providers(&healthy, settings) + .map(|(top, bottom)| { + vec![ + provider_status_label(top, lang), + provider_status_label(bottom, lang), + ] + }) + .unwrap_or_else(|| { + healthy + .first() + .map(|s| provider_status_label(s, lang)) + .into_iter() + .collect() + }); + } + let Some(selected) = pick_tray_provider( &healthy, settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal", @@ -629,13 +660,58 @@ fn render_tray_icon_for_settings( weekly_pct: Option, all_error: bool, ) -> (Vec, u32, u32) { - if settings.menu_bar_shows_percent { + if settings.tray_icon_mode == TrayIconMode::Stacked + && let Some(bottom_pct) = weekly_pct + { + render_stacked_bar_icon_rgba(session_pct, bottom_pct, all_error) + } else if settings.menu_bar_shows_percent { render_percent_icon_rgba(session_pct, all_error) } else { render_bar_icon_rgba(session_pct, weekly_pct, all_error) } } +/// Resolve a stable top/bottom pair while retaining stale saved preferences. +/// Eligible provider order is the user's provider display order. An invalid, +/// disabled, or duplicate preference falls back without rewriting settings. +fn pick_stacked_tray_providers<'a>( + ok_snapshots: &'a [&'a crate::commands::ProviderUsageSnapshot], + settings: &Settings, +) -> Option<( + &'a crate::commands::ProviderUsageSnapshot, + &'a crate::commands::ProviderUsageSnapshot, +)> { + if ok_snapshots.len() < 2 { + return None; + } + + let preferred = |provider_id: Option<&str>| { + provider_id.and_then(|id| { + ok_snapshots + .iter() + .copied() + .find(|snapshot| snapshot.provider_id == id) + }) + }; + let preferred_bottom = preferred(settings.stacked_tray_bottom_provider.as_deref()); + let top = preferred(settings.stacked_tray_top_provider.as_deref()).or_else(|| { + ok_snapshots.iter().copied().find(|snapshot| { + preferred_bottom.map(|bottom| bottom.provider_id.as_str()) + != Some(snapshot.provider_id.as_str()) + }) + })?; + let bottom = preferred_bottom + .filter(|snapshot| snapshot.provider_id != top.provider_id) + .or_else(|| { + ok_snapshots + .iter() + .copied() + .find(|snapshot| snapshot.provider_id != top.provider_id) + })?; + + Some((top, bottom)) +} + /// Pick the provider whose usage the tray icon should render. /// /// Exposed so that the unit tests can exercise both `highest` and `first` @@ -1215,6 +1291,71 @@ mod tests { ); } + #[test] + fn stacked_mode_resolves_distinct_preferred_providers() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let codex = fake_snapshot("codex", "Codex", 30.0); + let claude = fake_snapshot("claude", "Claude", 72.0); + let gemini = fake_snapshot("gemini", "Gemini", 44.0); + let snapshots = vec![&codex, &claude, &gemini]; + + let pair = pick_stacked_tray_providers(&snapshots, &settings).unwrap(); + + assert_eq!(pair.0.provider_id, "claude"); + assert_eq!(pair.1.provider_id, "codex"); + } + + #[test] + fn stacked_mode_falls_back_around_stale_and_duplicate_preferences() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("missing".to_string()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Settings::default() + }; + let codex = fake_snapshot("codex", "Codex", 30.0); + let claude = fake_snapshot("claude", "Claude", 72.0); + let snapshots = vec![&codex, &claude]; + + let pair = pick_stacked_tray_providers(&snapshots, &settings).unwrap(); + + assert_eq!(pair.0.provider_id, "codex"); + assert_eq!(pair.1.provider_id, "claude"); + } + + #[test] + fn stacked_mode_lists_both_provider_statuses() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let labels = status_labels_for_settings( + &settings, + &snapshots, + codexbar::settings::Language::English, + ); + + assert_eq!( + labels, + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); + } + #[test] fn tray_icon_renderer_uses_percent_mode_when_enabled() { let bar_settings = Settings { @@ -1235,6 +1376,23 @@ mod tests { assert_ne!(bar, percent); } + #[test] + fn tray_icon_renderer_uses_stacked_rows_for_two_providers() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + menu_bar_shows_percent: true, + ..Settings::default() + }; + + let (stacked, width, height) = + render_tray_icon_for_settings(&settings, 72.0, Some(40.0), false); + let (expected, expected_width, expected_height) = + render_stacked_bar_icon_rgba(72.0, 40.0, false); + + assert_eq!((width, height), (expected_width, expected_height)); + assert_eq!(stacked, expected); + } + #[test] fn tooltip_uses_compact_status_labels() { let mut claude = fake_snapshot("claude", "Claude", 13.0); diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..1c07040562 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -541,6 +541,9 @@ export const ALL_LOCALE_KEYS = [ "TrayIconModeHelper", "TrayIconModeSingle", "TrayIconModePerProvider", + "TrayIconModeStacked", + "StackedTrayTopProvider", + "StackedTrayBottomProvider", "ShowProviderIcons", "ShowProviderIconsHelper", "PreferHighestUsage", diff --git a/apps/desktop-tauri/src/styles.css b/apps/desktop-tauri/src/styles.css index a41405732c..bd97b85345 100644 --- a/apps/desktop-tauri/src/styles.css +++ b/apps/desktop-tauri/src/styles.css @@ -6035,31 +6035,6 @@ html:has(.menu-surface--tray) { color: inherit; } -.accent-color-swatch-row { - display: flex; - align-items: center; - gap: 8px; - margin-top: 8px; - font-size: 12px; -} - -.accent-color-swatch { - display: inline-block; - width: 16px; - height: 16px; - border-radius: 4px; - border: 1px solid var(--border-color); -} - -.accent-color-swatch-label { - color: var(--text-secondary); -} - -.accent-color-swatch-value { - font-family: var(--font-mono, monospace); - color: var(--text-secondary); -} - /* ── Mistral monthly spend row (#2821, #2947) ──────────────────── */ .menu-card__monthly-spend { margin-top: 4px; diff --git a/apps/desktop-tauri/src/surfaces/Settings.tsx b/apps/desktop-tauri/src/surfaces/Settings.tsx index 21b5828dda..f3fb995d44 100644 --- a/apps/desktop-tauri/src/surfaces/Settings.tsx +++ b/apps/desktop-tauri/src/surfaces/Settings.tsx @@ -253,7 +253,13 @@ export default function Settings({ state, initialTab: propTab }: { state: Bootst )} {activeTab === "menuBar" && ( - + )} {activeTab === "menu" && ( diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx new file mode 100644 index 0000000000..6a5a385660 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx @@ -0,0 +1,21 @@ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { AccentColorSection } from "./AccentColorSection"; + +describe("AccentColorSection", () => { + it("uses the native color input as the only color preview", () => { + const { container } = render( + key} + onChange={vi.fn()} + />, + ); + + expect(container.querySelector('input[type="color"]')).toHaveValue("#123456"); + expect(container.querySelector(".accent-color-swatch-row")).toBeNull(); + expect(container.querySelector(".accent-color-swatch")).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx index 9e0f8dca41..112e09f856 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx @@ -94,16 +94,6 @@ export function AccentColorSection({ {t("ProviderAccentColorReset")} -

- - {t("ProviderAccentColor")} - - - {effective} -
{error &&

{error}

} ); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx index eec4c7093b..9780fd0311 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx @@ -14,7 +14,10 @@ import DisplayTab from "./DisplayTab"; import type { SettingsSnapshot } from "../../../types/bridge"; const baseSettings = { + enabledProviders: ["codex", "claude"], trayIconMode: "single", + stackedTrayTopProvider: null, + stackedTrayBottomProvider: null, trayPanelAlwaysOnTop: false, switcherShowsIcons: false, menuBarShowsHighestUsage: false, @@ -98,3 +101,28 @@ describe("DisplayTab window scale", () => { expect(set).toHaveBeenCalledWith({ trayPanelAlwaysOnTop: true }); }); }); + +describe("DisplayTab stacked tray providers", () => { + it("persists explicit top and bottom provider choices", () => { + const set = vi.fn(); + render( + , + ); + const selects = screen.getAllByRole("combobox"); + + fireEvent.change(selects[1], { target: { value: "claude" } }); + fireEvent.change(selects[2], { target: { value: "codex" } }); + + expect(set).toHaveBeenCalledWith({ stackedTrayTopProvider: "claude" }); + expect(set).toHaveBeenCalledWith({ stackedTrayBottomProvider: "codex" }); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx index aed13a723b..5823c4a0a9 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx @@ -4,6 +4,7 @@ import { Field, Select, Toggle } from "../../../components/FormControls"; import type { MenuBarDisplayMode, OverviewLayout, + ProviderCatalogEntry, TrayIconMode, TrayVisibilityStatusDto, } from "../../../types/bridge"; @@ -20,7 +21,11 @@ export default function DisplayTab({ settings, set, saving, -}: TabProps & { mode?: "menuBar" | "menu" }) { + providers = [], +}: TabProps & { + mode?: "menuBar" | "menu"; + providers?: ProviderCatalogEntry[]; +}) { const { t } = useLocale(); const [windowScaleDraft, setWindowScaleDraft] = useState(() => clampWindowScalePercent(settings.windowScalePercent), @@ -43,6 +48,13 @@ export default function DisplayTab({ set({ windowScalePercent: next }); } }, [set, settings.windowScalePercent, windowScaleDraft]); + const providerName = new Map( + providers.map((provider) => [provider.id, provider.displayName]), + ); + const stackedProviderOptions = settings.enabledProviders.map((providerId) => ({ + value: providerId, + label: providerName.get(providerId) ?? providerId, + })); return ( <> {/* ── Menu bar ─────────────────────────────────────────────── */} @@ -59,10 +71,47 @@ export default function DisplayTab({ options={[ { value: "single", label: t("TrayIconModeSingle") }, { value: "perProvider", label: t("TrayIconModePerProvider") }, + { value: "stacked", label: t("TrayIconModeStacked") }, ]} onChange={(v) => set({ trayIconMode: v as TrayIconMode })} /> + {settings.trayIconMode === "stacked" && ( + <> + + + provider.value !== settings.stackedTrayTopProvider, + ), + ]} + onChange={(provider) => + set({ stackedTrayBottomProvider: provider }) + } + /> + + + )} set({ menuBarShowsHighestUsage: v })} /> @@ -92,7 +141,7 @@ export default function DisplayTab({ > set({ menuBarShowsPercent: v })} /> diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 5aa0a40f27..46b629f03f 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -12,7 +12,7 @@ export type SettingsTabId = // ── Narrowed string-literal unions (persisted settings enums) ───────── -export type TrayIconMode = "single" | "perProvider"; +export type TrayIconMode = "single" | "perProvider" | "stacked"; export type NotificationSoundTheme = "windows" | "codexBar"; @@ -190,6 +190,8 @@ export interface SettingsSnapshot { predictivePaceWarningEnabled: boolean; showPace?: boolean; trayIconMode: TrayIconMode; + stackedTrayTopProvider?: string | null; + stackedTrayBottomProvider?: string | null; switcherShowsIcons: boolean; menuBarShowsHighestUsage: boolean; menuBarShowsPercent: boolean; @@ -299,6 +301,8 @@ export interface SettingsUpdate { predictivePaceWarningEnabled?: boolean; showPace?: boolean; trayIconMode?: TrayIconMode; + stackedTrayTopProvider?: string; + stackedTrayBottomProvider?: string; switcherShowsIcons?: boolean; menuBarShowsHighestUsage?: boolean; menuBarShowsPercent?: boolean; diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..276a31160e 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -807,6 +807,9 @@ locale_keys! { TrayIconModeHelper, TrayIconModeSingle, TrayIconModePerProvider, + TrayIconModeStacked, + StackedTrayTopProvider, + StackedTrayBottomProvider, ShowProviderIcons, ShowProviderIconsHelper, PreferHighestUsage, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..53b4453f6d 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -478,6 +478,9 @@ TrayIconModeLabel = Tray icon mode TrayIconModeHelper = Single unified icon or one icon per enabled provider. TrayIconModeSingle = Single TrayIconModePerProvider = Per provider +TrayIconModeStacked = Stacked providers +StackedTrayTopProvider = Top provider +StackedTrayBottomProvider = Bottom provider ShowProviderIcons = Show provider icons ShowProviderIconsHelper = Display provider icons in the tray switcher. PreferHighestUsage = Prefer highest usage diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index 38dc0b05ae..bae2c814c3 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -434,6 +434,9 @@ TrayIconModeLabel = Modo de ícono de bandeja TrayIconModeHelper = Ícono único combinado o un ícono por cada proveedor habilitado. TrayIconModeSingle = Único TrayIconModePerProvider = Por proveedor +TrayIconModeStacked = Proveedores apilados +StackedTrayTopProvider = Proveedor superior +StackedTrayBottomProvider = Proveedor inferior ShowProviderIcons = Mostrar íconos de proveedores ShowProviderIconsHelper = Mostrar íconos de proveedores en el selector de bandeja. PreferHighestUsage = Preferir uso más alto diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index b6c2dc0fc8..92618f72f7 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -416,6 +416,9 @@ TrayIconModeLabel = トレイアイコンモード TrayIconModeHelper = 単一の統合アイコンか、有効なプロバイダーごとのアイコン。 TrayIconModeSingle = 単一 TrayIconModePerProvider = プロバイダー別 +TrayIconModeStacked = プロバイダーを積み重ねる +StackedTrayTopProvider = 上のプロバイダー +StackedTrayBottomProvider = 下のプロバイダー ShowProviderIcons = プロバイダーアイコンを表示 ShowProviderIconsHelper = トレイスイッチャーにプロバイダーアイコンを表示。 PreferHighestUsage = 最も使用量が多いものを優先 diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 299dcf1a54..7f59c6fc29 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -421,6 +421,9 @@ TrayIconModeLabel = 트레이 아이콘 모드 TrayIconModeHelper = 단일 통합 아이콘 또는 활성화된 제공업체당 하나의 아이콘. TrayIconModeSingle = 단일 TrayIconModePerProvider = 제공업체별 +TrayIconModeStacked = 제공업체 쌓기 +StackedTrayTopProvider = 위쪽 제공업체 +StackedTrayBottomProvider = 아래쪽 제공업체 ShowProviderIcons = 제공업체 아이콘 표시 ShowProviderIconsHelper = 트레이 메뉴에 제공업체 아이콘을 표시합니다. PreferHighestUsage = 가장 높은 사용량 우선 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index f978d540e7..855f7a9eb6 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -400,6 +400,9 @@ TrayIconModeLabel = Режим значков в трее TrayIconModeHelper = Один унифицированный значок или один значок для каждого включенного провайдера. TrayIconModeSingle = Одинокий TrayIconModePerProvider = За провайдера +TrayIconModeStacked = Провайдеры стопкой +StackedTrayTopProvider = Верхний провайдер +StackedTrayBottomProvider = Нижний провайдер ShowProviderIcons = Показать значки провайдеров ShowProviderIconsHelper = Отображать значки провайдеров в переключателе трея. PreferHighestUsage = Предпочитаю максимальное использование diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index 08de23e2e3..ca4fabe791 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -437,6 +437,9 @@ TrayIconModeLabel = Tepsi simgesi modu TrayIconModeHelper = Tek birleşik simge veya etkin sağlayıcı başına bir simge. TrayIconModeSingle = Tek TrayIconModePerProvider = Sağlayıcı başına +TrayIconModeStacked = Yığılmış sağlayıcılar +StackedTrayTopProvider = Üst sağlayıcı +StackedTrayBottomProvider = Alt sağlayıcı ShowProviderIcons = Sağlayıcı simgelerini göster ShowProviderIconsHelper = Tepsi değiştiricisinde sağlayıcı simgelerini göster. PreferHighestUsage = En yüksek kullanımı tercih et diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index daf8e3c9af..98dd870db8 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -415,6 +415,9 @@ TrayIconModeLabel = 托盘图标模式 TrayIconModeHelper = 使用单一合并图标，或为每个已启用服务商显示独立图标。 TrayIconModeSingle = 合并 TrayIconModePerProvider = 按服务商 +TrayIconModeStacked = 堆叠服务商 +StackedTrayTopProvider = 上方服务商 +StackedTrayBottomProvider = 下方服务商 ShowProviderIcons = 显示服务商图标 ShowProviderIconsHelper = 在托盘切换器中显示服务商图标。 PreferHighestUsage = 优先显示最高用量 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 31bb3bcf02..f893d5b423 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -415,6 +415,9 @@ TrayIconModeLabel = 系統匣圖示模式 TrayIconModeHelper = 使用單一合併圖示，或為每個已啟用提供者顯示獨立圖示。 TrayIconModeSingle = 合併 TrayIconModePerProvider = 按提供者 +TrayIconModeStacked = 堆疊提供者 +StackedTrayTopProvider = 上方提供者 +StackedTrayBottomProvider = 下方提供者 ShowProviderIcons = 顯示提供者圖示 ShowProviderIconsHelper = 在系統匣切換器中顯示提供者圖示。 PreferHighestUsage = 優先顯示最高用量 diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..1dc2ad560f 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -136,6 +136,16 @@ pub struct Settings { #[serde(default)] pub tray_icon_mode: TrayIconMode, + /// Optional preferred provider for the upper row of a stacked tray icon. + /// Stale or disabled values are retained and ignored until eligible again. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stacked_tray_top_provider: Option, + + /// Optional preferred provider for the lower row of a stacked tray icon. + /// Stale or duplicate values fall back to the next eligible provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stacked_tray_bottom_provider: Option, + /// Show provider icons in the merged switcher UI #[serde(default = "default_true")] pub switcher_shows_icons: bool, @@ -541,6 +551,8 @@ impl Default for Settings { provider_usage_thresholds: HashMap::new(), merge_tray_icons: false, // Show single provider by default tray_icon_mode: TrayIconMode::default(), // Single icon by default + stacked_tray_top_provider: None, + stacked_tray_bottom_provider: None, switcher_shows_icons: true, menu_bar_shows_highest_usage: false, menu_bar_shows_percent: false, diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..4ff614af7d 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -35,6 +35,8 @@ pub(super) struct RawSettings { provider_usage_thresholds: HashMap, merge_tray_icons: bool, tray_icon_mode: TrayIconMode, + stacked_tray_top_provider: Option, + stacked_tray_bottom_provider: Option, #[serde(default = "default_true")] switcher_shows_icons: bool, menu_bar_shows_highest_usage: bool, @@ -206,6 +208,8 @@ impl Default for RawSettings { provider_usage_thresholds: HashMap::new(), merge_tray_icons: s.merge_tray_icons, tray_icon_mode: s.tray_icon_mode, + stacked_tray_top_provider: s.stacked_tray_top_provider, + stacked_tray_bottom_provider: s.stacked_tray_bottom_provider, switcher_shows_icons: s.switcher_shows_icons, menu_bar_shows_highest_usage: s.menu_bar_shows_highest_usage, menu_bar_shows_percent: s.menu_bar_shows_percent, @@ -533,6 +537,8 @@ impl From for Settings { ), merge_tray_icons: raw.merge_tray_icons, tray_icon_mode: raw.tray_icon_mode, + stacked_tray_top_provider: raw.stacked_tray_top_provider, + stacked_tray_bottom_provider: raw.stacked_tray_bottom_provider, switcher_shows_icons: raw.switcher_shows_icons, menu_bar_shows_highest_usage: raw.menu_bar_shows_highest_usage, menu_bar_shows_percent: raw.menu_bar_shows_percent, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..ed3f4d28c2 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -771,6 +771,27 @@ fn test_settings_with_utf8_bom_parses_perprovider_tray_mode() { assert_eq!(settings.tray_icon_mode, TrayIconMode::PerProvider); } +#[test] +fn stacked_tray_mode_preserves_provider_preferences() { + let json = r#"{ + "tray_icon_mode": "stacked", + "stacked_tray_top_provider": "claude", + "stacked_tray_bottom_provider": "codex" + }"#; + + let settings: Settings = serde_json::from_str(json).unwrap(); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); + assert_eq!( + settings.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); +} + #[test] fn test_language_serde_serialization() { // Test that Language serializes to lowercase string diff --git a/rust/src/settings/types.rs b/rust/src/settings/types.rs index 6113969b15..062338076e 100644 --- a/rust/src/settings/types.rs +++ b/rust/src/settings/types.rs @@ -247,6 +247,8 @@ pub enum TrayIconMode { Single, /// One tray icon per enabled provider PerProvider, + /// One tray icon with the selected metrics for two providers stacked vertically + Stacked, } impl TrayIconMode { @@ -255,6 +257,7 @@ impl TrayIconMode { match self { TrayIconMode::Single => "Single Icon", TrayIconMode::PerProvider => "Per Provider", + TrayIconMode::Stacked => "Stacked Providers", } } @@ -263,6 +266,7 @@ impl TrayIconMode { match self { TrayIconMode::Single => "Show one tray icon for all providers", TrayIconMode::PerProvider => "Show a separate tray icon for each enabled provider", + TrayIconMode::Stacked => "Show two providers as stacked usage meters", } } } diff --git a/rust/src/tray/mod.rs b/rust/src/tray/mod.rs index bf4c52ef88..0e2e80afc7 100755 --- a/rust/src/tray/mod.rs +++ b/rust/src/tray/mod.rs @@ -6,4 +6,6 @@ pub mod icon; pub mod render; pub use icon::LoadingPattern; -pub use render::{TRAY_ICON_SIZE, render_bar_icon_rgba, render_percent_icon_rgba}; +pub use render::{ + TRAY_ICON_SIZE, render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, +}; diff --git a/rust/src/tray/render.rs b/rust/src/tray/render.rs index 5b486a5b6b..1787f60f3c 100644 --- a/rust/src/tray/render.rs +++ b/rust/src/tray/render.rs @@ -93,6 +93,67 @@ pub fn render_bar_icon_rgba( (img.into_raw(), SZ, SZ) } +/// Render two providers as equally prominent stacked usage meters. +/// +/// Unlike [`render_bar_icon_rgba`], both rows represent the selected metric +/// for separate providers. The upper and lower rows therefore use equal +/// height so neither provider is presented as a secondary quota window. +pub fn render_stacked_bar_icon_rgba( + top_percent: f64, + bottom_percent: f64, + has_error: bool, +) -> (Vec, u32, u32) { + const SZ: u32 = TRAY_ICON_SIZE; + let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); + + for pixel in img.pixels_mut() { + *pixel = Rgba([0, 0, 0, 0]); + } + + let bg_alpha = if has_error { 180 } else { 255 }; + for y in 2..SZ - 2 { + for x in 2..SZ - 2 { + img.put_pixel(x, y, Rgba([60, 60, 70, bg_alpha])); + } + } + + let bar_left = 4u32; + let bar_right = SZ - 4; + let bar_width = bar_right - bar_left; + let mut draw_provider = |y_start: u32, y_end: u32, percent: f64| { + let (r, g, b) = UsageLevel::from_percent(percent).color(); + let color = if has_error { + #[allow( + clippy::cast_possible_truncation, + reason = "mean of three u8 channels is bounded to 0..=255" + )] + let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; + Rgba([gray, gray, gray, 255]) + } else { + Rgba([r, g, b, 255]) + }; + #[allow( + clippy::cast_possible_truncation, + reason = "percent is clamped to 0..=100 and scaled to a 24-pixel meter" + )] + let fill = ((percent.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; + let fill_end = (bar_left + fill).min(bar_right); + + for y in y_start..y_end { + for x in bar_left..bar_right { + img.put_pixel(x, y, Rgba([80, 80, 90, 255])); + } + for x in bar_left..fill_end { + img.put_pixel(x, y, color); + } + } + }; + + draw_provider(6, 14, top_percent); + draw_provider(18, 26, bottom_percent); + (img.into_raw(), SZ, SZ) +} + /// Render a compact numeric percent tray icon as raw RGBA bytes. pub fn render_percent_icon_rgba(percent: f64, has_error: bool) -> (Vec, u32, u32) { const SZ: u32 = TRAY_ICON_SIZE; @@ -304,4 +365,24 @@ mod tests { let (rgba, w, h) = render_percent_icon_rgba(125.0, false); assert_eq!(u32::try_from(rgba.len()).unwrap(), w * h * 4); } + + #[test] + fn stacked_provider_icon_uses_equal_separate_rows() { + let (rgba, width, height) = render_stacked_bar_icon_rgba(100.0, 0.0, false); + assert_eq!((width, height), (TRAY_ICON_SIZE, TRAY_ICON_SIZE)); + + let pixel = |x: u32, y: u32| { + let index = ((y * width + x) * 4) as usize; + [ + rgba[index], + rgba[index + 1], + rgba[index + 2], + rgba[index + 3], + ] + }; + let (r, g, b) = UsageLevel::Critical.color(); + assert_eq!(pixel(8, 8), [r, g, b, 255]); + assert_eq!(pixel(8, 20), [80, 80, 90, 255]); + assert_eq!(pixel(8, 15), [60, 60, 70, 255]); + } } From a311d933e2996eae7573e4dea05786c9aaaf72c1 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:35:55 +0700 Subject: [PATCH 13/62] Centralize tray presentation planning --- apps/desktop-tauri/src-tauri/src/main.rs | 1 + .../src-tauri/src/tray_bridge.rs | 836 +-------------- .../src-tauri/src/tray_presentation.rs | 990 ++++++++++++++++++ 3 files changed, 1009 insertions(+), 818 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/tray_presentation.rs diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 83720c9d6b..30458cb9ea 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -19,6 +19,7 @@ mod surface_target; mod tray_accounts; mod tray_bridge; mod tray_menu; +mod tray_presentation; mod tray_visibility; mod usage_metric; mod window_positioner; diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index d5e667b68e..eb277cf2b0 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -3,19 +3,12 @@ use std::sync::Mutex; use crate::commands::ProviderCatalogEntry; -#[cfg(test)] -use codexbar::core::ProviderId; -use codexbar::settings::MetricPreference; -use codexbar::settings::{Settings, TrayIconMode}; +use codexbar::settings::Settings; use tauri::image::Image; use tauri::menu::{CheckMenuItemBuilder, IsMenuItem, Menu, MenuItem, PredefinedMenuItem, Submenu}; use tauri::tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}; use tauri::{AppHandle, Manager}; -use codexbar::tray::{ - render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, -}; - use crate::shell; use crate::state::{AppState, TrayAnchor}; use crate::surface::SurfaceMode; @@ -23,6 +16,7 @@ use crate::surface_target::SurfaceTarget; #[cfg(test)] use crate::tray_menu::build_tray_menu; use crate::tray_menu::{TrayMenuEntry, build_tray_menu_with}; +use crate::tray_presentation::{TrayPresentationPlan, headline_window}; #[derive(Debug, Clone, Copy)] struct MonitorScaleInfo { @@ -410,7 +404,8 @@ pub(crate) fn rebuild_tray_menu(app: &AppHandle) { let settings = Settings::load(); let status_labels = if let Some(st) = app.try_state::>() { let guard = st.lock().unwrap(); - status_labels_for_settings(&settings, &guard.provider_cache, settings.ui_language) + TrayPresentationPlan::resolve(&settings, &guard.provider_cache) + .status_labels(settings.ui_language) } else { vec![] }; @@ -428,7 +423,8 @@ pub fn update_tray_status_items( ) { let catalog = crate::commands::get_provider_catalog(); let settings = Settings::load(); - let status_labels = status_labels_for_settings(&settings, snapshots, settings.ui_language); + let status_labels = + TrayPresentationPlan::resolve(&settings, snapshots).status_labels(settings.ui_language); if let Ok(menu) = build_native_tray_menu(app, &catalog, &status_labels) && let Some(tray) = app.tray_by_id("codexbar-main") @@ -466,302 +462,16 @@ pub fn update_tray_icon_and_tooltip( return; }; - // ── Icon ───────────────────────────────────────────────────────────── let settings = Settings::load(); - let snapshots = snapshots.to_vec(); - let ordered_snapshots = ordered_snapshot_refs(&settings, &snapshots); - let ok_snapshots: Vec<_> = ordered_snapshots - .iter() - .copied() - .filter(|s| s.error.is_none()) - .collect(); - let all_error = ok_snapshots.is_empty() && !snapshots.is_empty(); - - let prefer_highest = settings.menu_bar_shows_highest_usage - || settings.menu_bar_display_mode.as_str() == "minimal"; - - let picked = pick_tray_provider(&ok_snapshots, prefer_highest); - - let (session_pct, weekly_pct) = if settings.tray_icon_mode == TrayIconMode::Stacked { - pick_stacked_tray_providers(&ok_snapshots, &settings) - .map(|(top, bottom)| { - ( - selected_tray_percents(top, &settings).0, - Some(selected_tray_percents(bottom, &settings).0), - ) - }) - .or_else(|| picked.map(|snapshot| selected_tray_percents(snapshot, &settings))) - .unwrap_or((0.0, None)) - } else { - match picked { - Some(s) => selected_tray_percents(s, &settings), - None => ( - ok_snapshots - .iter() - .map(|s| selected_tray_percents(s, &settings).0) - .fold(0.0_f64, f64::max), - None, - ), - } - }; - - let (rgba, w, h) = render_tray_icon_for_settings(&settings, session_pct, weekly_pct, all_error); + let plan = TrayPresentationPlan::resolve(&settings, snapshots); + let (rgba, w, h) = plan.render_icon(); let icon = Image::new_owned(rgba, w, h); let _ = tray.set_icon(Some(icon)); - // ── Tooltip ─────────────────────────────────────────────────────────── - let tooltip = build_tooltip(&snapshots, settings.ui_language); + let tooltip = build_tooltip(snapshots, settings.ui_language); let _ = tray.set_tooltip(Some(tooltip)); } -fn status_labels_for_settings( - settings: &Settings, - snapshots: &[crate::commands::ProviderUsageSnapshot], - lang: codexbar::settings::Language, -) -> Vec<(String, String)> { - let ordered_snapshots = ordered_snapshot_refs(settings, snapshots); - let healthy: Vec<_> = ordered_snapshots - .into_iter() - .filter(|s| s.error.is_none()) - .collect(); - if settings.tray_icon_mode == TrayIconMode::PerProvider { - return healthy - .into_iter() - .map(|s| provider_status_label(s, lang)) - .collect::>(); - } - - if settings.tray_icon_mode == TrayIconMode::Stacked { - return pick_stacked_tray_providers(&healthy, settings) - .map(|(top, bottom)| { - vec![ - provider_status_label(top, lang), - provider_status_label(bottom, lang), - ] - }) - .unwrap_or_else(|| { - healthy - .first() - .map(|s| provider_status_label(s, lang)) - .into_iter() - .collect() - }); - } - - let Some(selected) = pick_tray_provider( - &healthy, - settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal", - ) else { - return vec![]; - }; - - let (_, label) = provider_status_label(selected, lang); - vec![("status_summary".to_string(), label)] -} - -fn ordered_snapshot_refs<'a>( - settings: &Settings, - snapshots: &'a [crate::commands::ProviderUsageSnapshot], -) -> Vec<&'a crate::commands::ProviderUsageSnapshot> { - let order = settings - .provider_display_order_names() - .into_iter() - .enumerate() - .map(|(index, provider_id)| (provider_id, index)) - .collect::>(); - let mut ordered = snapshots.iter().collect::>(); - ordered.sort_by(|a, b| { - let a_order = order.get(&a.provider_id); - let b_order = order.get(&b.provider_id); - match (a_order, b_order) { - (Some(a_order), Some(b_order)) if a_order != b_order => a_order.cmp(b_order), - (Some(_), None) => std::cmp::Ordering::Less, - (None, Some(_)) => std::cmp::Ordering::Greater, - _ => a.display_name.cmp(&b.display_name), - } - }); - ordered -} - -fn provider_status_label( - snapshot: &crate::commands::ProviderUsageSnapshot, - lang: codexbar::settings::Language, -) -> (String, String) { - // MonthlyPlan metric (PAYG spend, e.g. Mistral): show formatted cost. - let provider = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id); - let preference = provider - .map(|id| Settings::load().get_provider_metric(id)) - .unwrap_or_default(); - if preference == MetricPreference::MonthlyPlan - && let Some(cost) = snapshot.cost.as_ref() - { - let amount = if !cost.formatted_used.is_empty() { - cost.formatted_used.clone() - } else { - crate::commands::format_cost_amount(cost) - }; - return ( - snapshot.provider_id.clone(), - format!("{} {}", snapshot.display_name, amount), - ); - } - - let label = crate::commands::compact_tray_status_label(headline_window(snapshot), lang); - ( - snapshot.provider_id.clone(), - format!("{} {}", snapshot.display_name, label), - ) -} - -/// Window that headline tray surfaces should label for a provider. -/// -/// F5 (upstream 0.48.0): for Codex, prefer the first non-informational lane so -/// a monthly-only plan shows the monthly window with its reset countdown -/// instead of the informational "No active 5h session" placeholder. -/// -/// Shared by the tray menu rows (`provider_status_label`) and the tray tooltip -/// (`build_tooltip`) so the two cannot drift apart. -fn headline_window( - snapshot: &crate::commands::ProviderUsageSnapshot, -) -> &crate::commands::RateWindowSnapshot { - if snapshot.provider_id == "codex" { - codex_lane_headline_window(snapshot) - } else { - &snapshot.primary - } -} - -/// F5 (upstream 0.48.0): pick the first non-informational Codex lane in -/// session → weekly → monthly order. When all lanes are informational -/// (no active session at all), fall back to the primary for the -/// "No active 5h session" placeholder. -pub(crate) fn codex_lane_headline_window( - snapshot: &crate::commands::ProviderUsageSnapshot, -) -> &crate::commands::RateWindowSnapshot { - if !snapshot.primary.is_informational { - return &snapshot.primary; - } - if let Some(ref secondary) = snapshot.secondary - && !secondary.is_informational - { - return secondary; - } - if let Some(ref tertiary) = snapshot.tertiary - && !tertiary.is_informational - { - return tertiary; - } - &snapshot.primary -} - -fn render_tray_icon_for_settings( - settings: &Settings, - session_pct: f64, - weekly_pct: Option, - all_error: bool, -) -> (Vec, u32, u32) { - if settings.tray_icon_mode == TrayIconMode::Stacked - && let Some(bottom_pct) = weekly_pct - { - render_stacked_bar_icon_rgba(session_pct, bottom_pct, all_error) - } else if settings.menu_bar_shows_percent { - render_percent_icon_rgba(session_pct, all_error) - } else { - render_bar_icon_rgba(session_pct, weekly_pct, all_error) - } -} - -/// Resolve a stable top/bottom pair while retaining stale saved preferences. -/// Eligible provider order is the user's provider display order. An invalid, -/// disabled, or duplicate preference falls back without rewriting settings. -fn pick_stacked_tray_providers<'a>( - ok_snapshots: &'a [&'a crate::commands::ProviderUsageSnapshot], - settings: &Settings, -) -> Option<( - &'a crate::commands::ProviderUsageSnapshot, - &'a crate::commands::ProviderUsageSnapshot, -)> { - if ok_snapshots.len() < 2 { - return None; - } - - let preferred = |provider_id: Option<&str>| { - provider_id.and_then(|id| { - ok_snapshots - .iter() - .copied() - .find(|snapshot| snapshot.provider_id == id) - }) - }; - let preferred_bottom = preferred(settings.stacked_tray_bottom_provider.as_deref()); - let top = preferred(settings.stacked_tray_top_provider.as_deref()).or_else(|| { - ok_snapshots.iter().copied().find(|snapshot| { - preferred_bottom.map(|bottom| bottom.provider_id.as_str()) - != Some(snapshot.provider_id.as_str()) - }) - })?; - let bottom = preferred_bottom - .filter(|snapshot| snapshot.provider_id != top.provider_id) - .or_else(|| { - ok_snapshots - .iter() - .copied() - .find(|snapshot| snapshot.provider_id != top.provider_id) - })?; - - Some((top, bottom)) -} - -/// Pick the provider whose usage the tray icon should render. -/// -/// Exposed so that the unit tests can exercise both `highest` and `first` -/// paths without needing a live Tauri app handle. -fn pick_tray_provider<'a>( - ok_snapshots: &'a [&'a crate::commands::ProviderUsageSnapshot], - prefer_highest: bool, -) -> Option<&'a crate::commands::ProviderUsageSnapshot> { - if ok_snapshots.is_empty() { - return None; - } - if prefer_highest { - ok_snapshots.iter().copied().max_by(|a, b| { - a.primary - .used_percent - .partial_cmp(&b.primary.used_percent) - .unwrap_or(std::cmp::Ordering::Equal) - }) - } else { - Some(ok_snapshots[0]) - } -} - -fn selected_tray_percents( - snapshot: &crate::commands::ProviderUsageSnapshot, - settings: &Settings, -) -> (f64, Option) { - let (selected, companion) = - crate::usage_metric::selected_usage_icon_windows(snapshot, settings); - ( - display_metric_percent(&selected, settings.show_as_used), - companion - .as_ref() - .map(|window| display_metric_percent(window, settings.show_as_used)), - ) -} - -fn display_metric_percent(window: &crate::commands::RateWindowSnapshot, show_as_used: bool) -> f64 { - if window.is_informational { - return 0.0; - } - if window.is_exhausted || window.used_percent >= 100.0 { - return if show_as_used { 100.0 } else { 0.0 }; - } - - let used_percent = window.used_percent; - let used = used_percent.clamp(0.0, 100.0); - if show_as_used { used } else { 100.0 - used } -} - /// Build a compact multi-line tooltip string from provider snapshots. fn build_tooltip( snapshots: &[crate::commands::ProviderUsageSnapshot], @@ -1186,213 +896,6 @@ mod tests { fake_snapshot_with(id, display, used_percent, None, None, None) } - fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { - crate::commands::NamedRateWindowSnapshot { - id: "additional_budget".to_string(), - title: "Additional Budget".to_string(), - fallback_lane: false, - window: crate::commands::RateWindowSnapshot { - used_percent: percent, - remaining_percent: 100.0 - percent, - window_minutes: None, - resets_at: None, - reset_description: None, - is_exhausted: false, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }, - } - } - - #[test] - fn pick_tray_provider_highest_picks_max_primary() { - let a = fake_snapshot("codex", "Codex", 30.0); - let b = fake_snapshot("claude", "Claude", 72.5); - let c = fake_snapshot("gemini", "Gemini", 50.0); - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![&a, &b, &c]; - - let picked = pick_tray_provider(&refs, /* prefer_highest = */ true) - .expect("highest mode should pick a provider"); - assert_eq!(picked.provider_id, "claude"); - } - - #[test] - fn pick_tray_provider_first_preserves_catalog_order() { - let a = fake_snapshot("codex", "Codex", 30.0); - let b = fake_snapshot("claude", "Claude", 72.5); - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![&a, &b]; - - let picked = pick_tray_provider(&refs, /* prefer_highest = */ false) - .expect("non-highest mode should still pick the first entry"); - assert_eq!(picked.provider_id, "codex"); - } - - #[test] - fn pick_tray_provider_none_when_empty() { - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![]; - assert!(pick_tray_provider(&refs, true).is_none()); - assert!(pick_tray_provider(&refs, false).is_none()); - } - - #[test] - fn status_labels_per_provider_mode_lists_each_healthy_provider() { - let settings = Settings { - tray_icon_mode: TrayIconMode::PerProvider, - provider_order: codexbar::settings::normalize_provider_order(&[ - "claude".to_string(), - "codex".to_string(), - ]), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = status_labels_for_settings( - &settings, - &snapshots, - codexbar::settings::Language::English, - ); - - assert_eq!( - labels, - vec![ - ("claude".to_string(), "Claude 72%".to_string()), - ("codex".to_string(), "Codex 30%".to_string()), - ] - ); - } - - #[test] - fn status_labels_single_mode_collapses_to_selected_provider() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Single, - menu_bar_shows_highest_usage: true, - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = status_labels_for_settings( - &settings, - &snapshots, - codexbar::settings::Language::English, - ); - - assert_eq!( - labels, - vec![("status_summary".to_string(), "Claude 72%".to_string())] - ); - } - - #[test] - fn stacked_mode_resolves_distinct_preferred_providers() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("claude".to_string()), - stacked_tray_bottom_provider: Some("codex".to_string()), - ..Settings::default() - }; - let codex = fake_snapshot("codex", "Codex", 30.0); - let claude = fake_snapshot("claude", "Claude", 72.0); - let gemini = fake_snapshot("gemini", "Gemini", 44.0); - let snapshots = vec![&codex, &claude, &gemini]; - - let pair = pick_stacked_tray_providers(&snapshots, &settings).unwrap(); - - assert_eq!(pair.0.provider_id, "claude"); - assert_eq!(pair.1.provider_id, "codex"); - } - - #[test] - fn stacked_mode_falls_back_around_stale_and_duplicate_preferences() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("missing".to_string()), - stacked_tray_bottom_provider: Some("claude".to_string()), - ..Settings::default() - }; - let codex = fake_snapshot("codex", "Codex", 30.0); - let claude = fake_snapshot("claude", "Claude", 72.0); - let snapshots = vec![&codex, &claude]; - - let pair = pick_stacked_tray_providers(&snapshots, &settings).unwrap(); - - assert_eq!(pair.0.provider_id, "codex"); - assert_eq!(pair.1.provider_id, "claude"); - } - - #[test] - fn stacked_mode_lists_both_provider_statuses() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("claude".to_string()), - stacked_tray_bottom_provider: Some("codex".to_string()), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = status_labels_for_settings( - &settings, - &snapshots, - codexbar::settings::Language::English, - ); - - assert_eq!( - labels, - vec![ - ("claude".to_string(), "Claude 72%".to_string()), - ("codex".to_string(), "Codex 30%".to_string()), - ] - ); - } - - #[test] - fn tray_icon_renderer_uses_percent_mode_when_enabled() { - let bar_settings = Settings { - menu_bar_shows_percent: false, - ..Settings::default() - }; - let percent_settings = Settings { - menu_bar_shows_percent: true, - ..Settings::default() - }; - - let (bar, bar_w, bar_h) = - render_tray_icon_for_settings(&bar_settings, 72.0, Some(40.0), false); - let (percent, pct_w, pct_h) = - render_tray_icon_for_settings(&percent_settings, 72.0, Some(40.0), false); - - assert_eq!((bar_w, bar_h), (pct_w, pct_h)); - assert_ne!(bar, percent); - } - - #[test] - fn tray_icon_renderer_uses_stacked_rows_for_two_providers() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - menu_bar_shows_percent: true, - ..Settings::default() - }; - - let (stacked, width, height) = - render_tray_icon_for_settings(&settings, 72.0, Some(40.0), false); - let (expected, expected_width, expected_height) = - render_stacked_bar_icon_rgba(72.0, 40.0, false); - - assert_eq!((width, height), (expected_width, expected_height)); - assert_eq!(stacked, expected); - } - #[test] fn tooltip_uses_compact_status_labels() { let mut claude = fake_snapshot("claude", "Claude", 13.0); @@ -1469,319 +972,16 @@ mod tests { "{japanese_tooltip}" ); - let (_, english_label) = - provider_status_label(&claude, codexbar::settings::Language::English); - let (_, japanese_label) = - provider_status_label(&claude, codexbar::settings::Language::Japanese); + let settings = Settings::default(); + let snapshots = vec![claude]; + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + let english_label = plan.status_labels(codexbar::settings::Language::English)[0] + .1 + .clone(); + let japanese_label = plan.status_labels(codexbar::settings::Language::Japanese)[0] + .1 + .clone(); assert!(english_label.contains("Resets in"), "{english_label}"); assert!(japanese_label.contains("リセットまで"), "{japanese_label}"); } - - #[test] - fn selected_tray_percent_uses_cursor_extra_usage_cost() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(20.0)); - } - - #[test] - fn selected_tray_percent_tracks_extra_rate_window() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn copilot_automatic_tracks_highest_extra_rate_window() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - } - - #[test] - fn selected_tray_percent_respects_remaining_display_mode() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 85.0); - assert_eq!(secondary, Some(80.0)); - } - - #[test] - fn exhausted_automatic_window_never_renders_as_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(40.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(100.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = false; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn missing_automatic_window_does_not_look_like_available_remaining_progress() { - let settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); - snapshot.primary.is_informational = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(remaining, 0.0); - } - - #[test] - fn selected_tray_percent_falls_back_when_extra_usage_missing() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 72.0); - } - - #[test] - fn single_meaningful_secondary_quota_uses_full_single_meter() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { - let settings = Settings::default(); - let mut snapshot = - fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, Some(30.0)); - } - - #[test] - fn two_meaningful_quotas_keep_two_meter_layout() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); - let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(40.0)); - } - - #[test] - fn informational_primary_skips_session_and_automatic_phantom_zero() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - // Session preference must not paint the synthetic 0% primary; - // it falls through to Automatic which prefers weekly (42%). - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - assert_ne!(primary, 0.0); - - // Automatic also prefers weekly over informational primary. - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - } - - #[test] - fn claude_automatic_prefers_weekly_when_model_exhausted() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot_with("claude", "Claude", 40.0, Some(22.0), None, None); - snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { - used_percent: 100.0, - remaining_percent: 0.0, - window_minutes: Some(10080), - resets_at: None, - reset_description: None, - is_exhausted: true, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 22.0); - - // Explicit model override is untouched. - let mut overridden = settings.clone(); - overridden.set_provider_metric(ProviderId::Claude, MetricPreference::Model); - let (primary, _) = selected_tray_percents(&snapshot, &overridden); - assert_eq!(primary, 100.0); - } - - #[test] - fn automatic_prefers_exhausted_weekly_over_low_session() { - let settings = Settings::default(); - let snapshot = fake_snapshot_with("codex", "Codex", 20.0, Some(100.0), None, None); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 100.0); - - // Explicit session override still wins. - let mut overridden = settings.clone(); - overridden.set_provider_metric(ProviderId::Codex, MetricPreference::Session); - let (primary, _) = selected_tray_percents(&snapshot, &overridden); - assert_eq!(primary, 20.0); - } - - #[test] - fn automatic_picks_highest_among_model_and_extra_windows() { - let settings = Settings::default(); - let mut snapshot = - fake_snapshot_with("gemini", "Gemini", 10.0, Some(30.0), Some(40.0), None); - snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { - used_percent: 55.0, - remaining_percent: 45.0, - window_minutes: None, - resets_at: None, - reset_description: None, - is_exhausted: false, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }); - snapshot.extra_rate_windows.push(fake_extra_window(90.0)); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 90.0); - } - - #[test] - fn f5_headline_prefers_non_informational_primary() { - let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_secondary_when_primary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); - snapshot.primary.is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); - snapshot.primary.is_informational = true; - snapshot.secondary.as_mut().unwrap().is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_returns_primary_when_all_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); - snapshot.primary.is_informational = true; - if let Some(sec) = &mut snapshot.secondary { - sec.is_informational = true; - } - if let Some(ter) = &mut snapshot.tertiary { - ter.is_informational = true; - } - let headline = codex_lane_headline_window(&snapshot); - // Falls back to primary (the placeholder) when all are informational. - assert!(headline.is_informational); - } } diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs new file mode 100644 index 0000000000..e577f70312 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs @@ -0,0 +1,990 @@ +//! Pure tray presentation policy shared by the native tray surfaces. + +use crate::commands::{ProviderUsageSnapshot, RateWindowSnapshot}; +use codexbar::settings::{Language, MetricPreference, Settings, TrayIconMode}; +use codexbar::tray::{ + render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, +}; + +#[derive(Debug, Clone, Copy, PartialEq)] +enum TrayIconPlan { + Bars { + primary_percent: f64, + secondary_percent: Option, + has_error: bool, + }, + Percent { + percent: f64, + has_error: bool, + }, + Stacked { + top_percent: f64, + bottom_percent: f64, + has_error: bool, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TrayStatusKey { + Summary, + Provider, +} + +#[derive(Debug, Clone, Copy)] +struct TrayStatusRow<'a> { + key: TrayStatusKey, + snapshot: &'a ProviderUsageSnapshot, +} + +/// Fully resolved tray presentation, independent of Tauri and operating-system state. +/// +/// The plan is the single policy boundary for provider ordering, mode-specific +/// selection, metric selection, status rows, and icon renderer choice. +pub(crate) struct TrayPresentationPlan<'a> { + settings: &'a Settings, + icon: TrayIconPlan, + status_rows: Vec>, +} + +impl<'a> TrayPresentationPlan<'a> { + pub(crate) fn resolve(settings: &'a Settings, snapshots: &'a [ProviderUsageSnapshot]) -> Self { + let ordered = ordered_snapshot_refs(settings, snapshots); + let healthy = ordered + .into_iter() + .filter(|snapshot| snapshot.error.is_none()) + .collect::>(); + let has_error = healthy.is_empty() && !snapshots.is_empty(); + let prefer_highest = + settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal"; + let selected = pick_tray_provider(&healthy, prefer_highest); + + let (primary_percent, secondary_percent, status_rows) = match settings.tray_icon_mode { + TrayIconMode::Stacked => { + if let Some((top, bottom)) = pick_stacked_tray_providers(&healthy, settings) { + ( + selected_tray_percents(top, settings).0, + Some(selected_tray_percents(bottom, settings).0), + vec![ + TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot: top, + }, + TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot: bottom, + }, + ], + ) + } else { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or((0.0, None)); + let rows = healthy + .first() + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot, + }) + .into_iter() + .collect(); + (percents.0, percents.1, rows) + } + } + TrayIconMode::PerProvider => { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or_else(|| fallback_percents(&healthy, settings)); + let rows = healthy + .iter() + .copied() + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot, + }) + .collect(); + (percents.0, percents.1, rows) + } + TrayIconMode::Single => { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or_else(|| fallback_percents(&healthy, settings)); + let rows = selected + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Summary, + snapshot, + }) + .into_iter() + .collect(); + (percents.0, percents.1, rows) + } + }; + + let icon = resolve_icon_plan(settings, primary_percent, secondary_percent, has_error); + + Self { + settings, + icon, + status_rows, + } + } + + pub(crate) fn render_icon(&self) -> (Vec, u32, u32) { + match self.icon { + TrayIconPlan::Bars { + primary_percent, + secondary_percent, + has_error, + } => render_bar_icon_rgba(primary_percent, secondary_percent, has_error), + TrayIconPlan::Percent { percent, has_error } => { + render_percent_icon_rgba(percent, has_error) + } + TrayIconPlan::Stacked { + top_percent, + bottom_percent, + has_error, + } => render_stacked_bar_icon_rgba(top_percent, bottom_percent, has_error), + } + } + + pub(crate) fn status_labels(&self, language: Language) -> Vec<(String, String)> { + self.status_rows + .iter() + .map(|row| { + let (_, label) = provider_status_label(row.snapshot, self.settings, language); + let key = match row.key { + TrayStatusKey::Summary => "status_summary".to_string(), + TrayStatusKey::Provider => row.snapshot.provider_id.clone(), + }; + (key, label) + }) + .collect() + } +} + +fn resolve_icon_plan( + settings: &Settings, + primary_percent: f64, + secondary_percent: Option, + has_error: bool, +) -> TrayIconPlan { + if settings.tray_icon_mode == TrayIconMode::Stacked + && let Some(bottom_percent) = secondary_percent + { + TrayIconPlan::Stacked { + top_percent: primary_percent, + bottom_percent, + has_error, + } + } else if settings.menu_bar_shows_percent { + TrayIconPlan::Percent { + percent: primary_percent, + has_error, + } + } else { + TrayIconPlan::Bars { + primary_percent, + secondary_percent, + has_error, + } + } +} + +fn fallback_percents( + healthy: &[&ProviderUsageSnapshot], + settings: &Settings, +) -> (f64, Option) { + ( + healthy + .iter() + .map(|snapshot| selected_tray_percents(snapshot, settings).0) + .fold(0.0_f64, f64::max), + None, + ) +} + +fn ordered_snapshot_refs<'a>( + settings: &Settings, + snapshots: &'a [ProviderUsageSnapshot], +) -> Vec<&'a ProviderUsageSnapshot> { + let order = settings + .provider_display_order_names() + .into_iter() + .enumerate() + .map(|(index, provider_id)| (provider_id, index)) + .collect::>(); + let mut ordered = snapshots.iter().collect::>(); + ordered.sort_by(|a, b| { + let a_order = order.get(&a.provider_id); + let b_order = order.get(&b.provider_id); + match (a_order, b_order) { + (Some(a_order), Some(b_order)) if a_order != b_order => a_order.cmp(b_order), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + _ => a.display_name.cmp(&b.display_name), + } + }); + ordered +} + +fn provider_status_label( + snapshot: &ProviderUsageSnapshot, + settings: &Settings, + language: Language, +) -> (String, String) { + let provider = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id); + let preference = provider + .map(|id| settings.get_provider_metric(id)) + .unwrap_or_default(); + if preference == MetricPreference::MonthlyPlan + && let Some(cost) = snapshot.cost.as_ref() + { + let amount = if !cost.formatted_used.is_empty() { + cost.formatted_used.clone() + } else { + crate::commands::format_cost_amount(cost) + }; + return ( + snapshot.provider_id.clone(), + format!("{} {}", snapshot.display_name, amount), + ); + } + + let label = crate::commands::compact_tray_status_label(headline_window(snapshot), language); + ( + snapshot.provider_id.clone(), + format!("{} {}", snapshot.display_name, label), + ) +} + +/// Window that headline tray surfaces should label for a provider. +pub(crate) fn headline_window(snapshot: &ProviderUsageSnapshot) -> &RateWindowSnapshot { + if snapshot.provider_id == "codex" { + codex_lane_headline_window(snapshot) + } else { + &snapshot.primary + } +} + +/// Pick the first non-informational Codex lane in session, weekly, monthly order. +pub(crate) fn codex_lane_headline_window(snapshot: &ProviderUsageSnapshot) -> &RateWindowSnapshot { + if !snapshot.primary.is_informational { + return &snapshot.primary; + } + if let Some(ref secondary) = snapshot.secondary + && !secondary.is_informational + { + return secondary; + } + if let Some(ref tertiary) = snapshot.tertiary + && !tertiary.is_informational + { + return tertiary; + } + &snapshot.primary +} + +/// Resolve a stable top/bottom pair while retaining stale saved preferences. +fn pick_stacked_tray_providers<'a>( + healthy: &'a [&'a ProviderUsageSnapshot], + settings: &Settings, +) -> Option<(&'a ProviderUsageSnapshot, &'a ProviderUsageSnapshot)> { + if healthy.len() < 2 { + return None; + } + + let preferred = |provider_id: Option<&str>| { + provider_id.and_then(|id| { + healthy + .iter() + .copied() + .find(|snapshot| snapshot.provider_id == id) + }) + }; + let preferred_bottom = preferred(settings.stacked_tray_bottom_provider.as_deref()); + let top = preferred(settings.stacked_tray_top_provider.as_deref()).or_else(|| { + healthy.iter().copied().find(|snapshot| { + preferred_bottom.map(|bottom| bottom.provider_id.as_str()) + != Some(snapshot.provider_id.as_str()) + }) + })?; + let bottom = preferred_bottom + .filter(|snapshot| snapshot.provider_id != top.provider_id) + .or_else(|| { + healthy + .iter() + .copied() + .find(|snapshot| snapshot.provider_id != top.provider_id) + })?; + + Some((top, bottom)) +} + +fn pick_tray_provider<'a>( + healthy: &'a [&'a ProviderUsageSnapshot], + prefer_highest: bool, +) -> Option<&'a ProviderUsageSnapshot> { + if prefer_highest { + healthy.iter().copied().max_by(|a, b| { + a.primary + .used_percent + .partial_cmp(&b.primary.used_percent) + .unwrap_or(std::cmp::Ordering::Equal) + }) + } else { + healthy.first().copied() + } +} + +fn selected_tray_percents( + snapshot: &ProviderUsageSnapshot, + settings: &Settings, +) -> (f64, Option) { + let (selected, companion) = + crate::usage_metric::selected_usage_icon_windows(snapshot, settings); + ( + display_metric_percent(&selected, settings.show_as_used), + companion + .as_ref() + .map(|window| display_metric_percent(window, settings.show_as_used)), + ) +} + +fn display_metric_percent(window: &RateWindowSnapshot, show_as_used: bool) -> f64 { + if window.is_informational { + return 0.0; + } + if window.is_exhausted || window.used_percent >= 100.0 { + return if show_as_used { 100.0 } else { 0.0 }; + } + + let used = window.used_percent.clamp(0.0, 100.0); + if show_as_used { used } else { 100.0 - used } +} + +#[cfg(test)] +mod tests { + use super::*; + use codexbar::core::{ProviderId, ProviderStateKind}; + + fn fake_snapshot(id: &str, display_name: &str, used_percent: f64) -> ProviderUsageSnapshot { + fake_snapshot_with(id, display_name, used_percent, None, None, None) + } + + fn fake_snapshot_with( + id: &str, + display_name: &str, + used_percent: f64, + secondary_percent: Option, + tertiary_percent: Option, + cost: Option<(f64, f64)>, + ) -> ProviderUsageSnapshot { + let window = |percent: f64| RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }; + + ProviderUsageSnapshot { + provider_id: id.into(), + display_name: display_name.into(), + primary: window(used_percent), + primary_label: None, + secondary: secondary_percent.map(window), + secondary_label: None, + model_specific: None, + tertiary: tertiary_percent.map(window), + tertiary_label: None, + extra_rate_windows: Vec::new(), + inventory: Vec::new(), + display_details: Vec::new(), + cost: cost.map(|(used, limit)| crate::commands::CostSnapshotBridge { + used, + limit: Some(limit), + remaining: Some((limit - used).max(0.0)), + currency_code: "USD".to_string(), + currency_symbol: None, + period: "monthly".to_string(), + resets_at: None, + formatted_used: format!("${used:.2}"), + formatted_limit: Some(format!("${limit:.2}")), + balance: None, + formatted_balance: None, + balance_updated_at: None, + account_id: None, + daily: Vec::new(), + always_visible: false, + }), + plan_name: None, + account_email: None, + subscription: None, + source_label: String::new(), + has_successful_claude_cli_quota: false, + updated_at: "2025-01-01T00:00:00Z".into(), + error: None, + error_state: ProviderStateKind::Ready, + pace: None, + account_organization: None, + tray_status_label: None, + fetch_duration_ms: None, + wayfinder_usage: None, + session_equivalent_forecast: None, + } + } + + #[test] + fn single_plan_uses_highest_provider_for_icon_and_summary() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 72.0, + secondary_percent: None, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![("status_summary".to_string(), "Claude 72%".to_string())] + ); + } + + #[test] + fn per_provider_plan_preserves_configured_order_for_status_rows() { + let settings = Settings { + tray_icon_mode: TrayIconMode::PerProvider, + provider_order: codexbar::settings::normalize_provider_order(&[ + "claude".to_string(), + "codex".to_string(), + ]), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let labels = + TrayPresentationPlan::resolve(&settings, &snapshots).status_labels(Language::English); + + assert_eq!( + labels, + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); + } + + #[test] + fn stacked_plan_resolves_distinct_preferences_once() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + fake_snapshot("gemini", "Gemini", 44.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 72.0, + bottom_percent: 30.0, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); + } + + #[test] + fn stacked_plan_falls_back_around_stale_and_duplicate_preferences() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("missing".to_string()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 30.0, + bottom_percent: 72.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); + assert_eq!(plan.status_rows[1].snapshot.provider_id, "claude"); + } + + #[test] + fn one_provider_stacked_plan_preserves_secondary_window_fallback() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + ..Settings::default() + }; + let snapshots = vec![fake_snapshot_with( + "codex", + "Codex", + 30.0, + Some(65.0), + None, + None, + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 65.0, + bottom_percent: 30.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows.len(), 1); + } + + #[test] + fn all_errors_produce_error_styled_zero_percent_plan() { + let settings = Settings { + menu_bar_shows_percent: true, + ..Settings::default() + }; + let mut snapshot = fake_snapshot("codex", "Codex", 30.0); + snapshot.error = Some("offline".to_string()); + let snapshots = vec![snapshot]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Percent { + percent: 0.0, + has_error: true, + } + ); + assert!(plan.status_rows.is_empty()); + } + + #[test] + fn plan_uses_selected_metric_and_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshots = vec![fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 85.0, + secondary_percent: Some(80.0), + has_error: false, + } + ); + } + + #[test] + fn render_icon_delegates_to_resolved_stacked_renderer() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 40.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.render_icon(), + render_stacked_bar_icon_rgba(72.0, 40.0, false) + ); + } + + #[test] + fn codex_headline_skips_informational_primary() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + assert_eq!(codex_lane_headline_window(&snapshot).used_percent, 25.0); + } + fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { + crate::commands::NamedRateWindowSnapshot { + id: "additional_budget".to_string(), + title: "Additional Budget".to_string(), + fallback_lane: false, + window: crate::commands::RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }, + } + } + + #[test] + fn selected_tray_percent_uses_cursor_extra_usage_cost() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(20.0)); + } + + #[test] + fn selected_tray_percent_tracks_extra_rate_window() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); + } + + #[test] + fn copilot_automatic_tracks_highest_extra_rate_window() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + } + + #[test] + fn selected_tray_percent_respects_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 85.0); + assert_eq!(secondary, Some(80.0)); + } + + #[test] + fn exhausted_automatic_window_never_renders_as_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(40.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); + } + + #[test] + fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(100.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = false; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); + } + + #[test] + fn missing_automatic_window_does_not_look_like_available_remaining_progress() { + let settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); + snapshot.primary.is_informational = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(remaining, 0.0); + } + + #[test] + fn selected_tray_percent_falls_back_when_extra_usage_missing() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 72.0); + } + + #[test] + fn single_meaningful_secondary_quota_uses_full_single_meter() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); + } + + #[test] + fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { + let settings = Settings::default(); + let mut snapshot = + fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, Some(30.0)); + } + + #[test] + fn two_meaningful_quotas_keep_two_meter_layout() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); + let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(40.0)); + } + + #[test] + fn informational_primary_skips_session_and_automatic_phantom_zero() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + // Session preference must not paint the synthetic 0% primary; + // it falls through to Automatic which prefers weekly (42%). + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); + assert_ne!(primary, 0.0); + + // Automatic also prefers weekly over informational primary. + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); + } + + #[test] + fn claude_automatic_prefers_weekly_when_model_exhausted() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 40.0, Some(22.0), None, None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 100.0, + remaining_percent: 0.0, + window_minutes: Some(10080), + resets_at: None, + reset_description: None, + is_exhausted: true, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 22.0); + + // Explicit model override is untouched. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Claude, MetricPreference::Model); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 100.0); + } + + #[test] + fn automatic_prefers_exhausted_weekly_over_low_session() { + let settings = Settings::default(); + let snapshot = fake_snapshot_with("codex", "Codex", 20.0, Some(100.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 100.0); + + // Explicit session override still wins. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Codex, MetricPreference::Session); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 20.0); + } + + #[test] + fn automatic_picks_highest_among_model_and_extra_windows() { + let settings = Settings::default(); + let mut snapshot = + fake_snapshot_with("gemini", "Gemini", 10.0, Some(30.0), Some(40.0), None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 55.0, + remaining_percent: 45.0, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + snapshot.extra_rate_windows.push(fake_extra_window(90.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 90.0); + } + + #[test] + fn f5_headline_prefers_non_informational_primary() { + let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); + } + + #[test] + fn f5_headline_falls_back_to_secondary_when_primary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); + } + + #[test] + fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); + snapshot.primary.is_informational = true; + snapshot.secondary.as_mut().unwrap().is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); + } + + #[test] + fn f5_headline_returns_primary_when_all_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); + snapshot.primary.is_informational = true; + if let Some(sec) = &mut snapshot.secondary { + sec.is_informational = true; + } + if let Some(ter) = &mut snapshot.tertiary { + ter.is_informational = true; + } + let headline = codex_lane_headline_window(&snapshot); + // Falls back to primary (the placeholder) when all are informational. + assert!(headline.is_informational); + } +} From 0837354b96fce915ca046e42a5b5bead9f5a1829 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:46:26 +0700 Subject: [PATCH 14/62] Fix tray presentation snapshot lifetimes --- .../src-tauri/src/tray_presentation.rs | 43 ++++++++++++++++++- 1 file changed, 41 insertions(+), 2 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs index e577f70312..8882a3e942 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs @@ -285,7 +285,7 @@ pub(crate) fn codex_lane_headline_window(snapshot: &ProviderUsageSnapshot) -> &R /// Resolve a stable top/bottom pair while retaining stale saved preferences. fn pick_stacked_tray_providers<'a>( - healthy: &'a [&'a ProviderUsageSnapshot], + healthy: &[&'a ProviderUsageSnapshot], settings: &Settings, ) -> Option<(&'a ProviderUsageSnapshot, &'a ProviderUsageSnapshot)> { if healthy.len() < 2 { @@ -320,7 +320,7 @@ fn pick_stacked_tray_providers<'a>( } fn pick_tray_provider<'a>( - healthy: &'a [&'a ProviderUsageSnapshot], + healthy: &[&'a ProviderUsageSnapshot], prefer_highest: bool, ) -> Option<&'a ProviderUsageSnapshot> { if prefer_highest { @@ -467,6 +467,24 @@ mod tests { ); } + #[test] + fn single_plan_borrows_selected_snapshot_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // `resolve` drops its temporary ordered/healthy vectors before returning. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); + } + #[test] fn per_provider_plan_preserves_configured_order_for_status_rows() { let settings = Settings { @@ -527,6 +545,27 @@ mod tests { ); } + #[test] + fn stacked_plan_borrows_both_snapshots_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // The plan retains references to the caller-owned snapshots, not the + // temporary vector of references used during selection. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); + assert!(std::ptr::eq(plan.status_rows[1].snapshot, &snapshots[0])); + } + #[test] fn stacked_plan_falls_back_around_stale_and_duplicate_preferences() { let settings = Settings { From 299aadba2978369bf2a8eda7ce6301734f63d059 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:53:03 +0700 Subject: [PATCH 15/62] Retire Crof provider safely --- .../providers/icons/ProviderIcon-crof.svg | 3 - .../providers/providerIcons.test.ts | 4 + .../src/components/providers/providerIcons.ts | 3 - apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 2 +- .../surfaces/settings/tabs/ProvidersTab.tsx | 1 - .../desktop-tauri/src/test/providerCatalog.ts | 1 - rust/assets/icons/ProviderIcon-crof.svg | 3 - rust/src/cli/serve/dashboard/icons.rs | 4 - .../dashboard/icons/ProviderIcon-crof.svg | 3 - rust/src/core/provider.rs | 11 +- rust/src/core/provider_factory.rs | 17 +- rust/src/core/token_accounts.rs | 1 - rust/src/providers/crof/mod.rs | 249 ------------------ rust/src/providers/mod.rs | 2 - rust/src/settings/api_keys.rs | 9 - rust/src/settings/raw.rs | 58 +++- rust/src/settings/tests.rs | 45 ++++ 17 files changed, 115 insertions(+), 301 deletions(-) delete mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg delete mode 100644 rust/assets/icons/ProviderIcon-crof.svg delete mode 100644 rust/src/cli/serve/dashboard/icons/ProviderIcon-crof.svg delete mode 100644 rust/src/providers/crof/mod.rs diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg deleted file mode 100644 index fdde018b8f..0000000000 --- a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts index 27c677024b..a204e1155a 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts @@ -8,4 +8,8 @@ describe("provider icon registry", () => { expect(PROVIDER_ICON_REGISTRY[id], id).toBeDefined(); } }); + + it("does not expose the retired Crof provider", () => { + expect(PROVIDER_ICON_REGISTRY).not.toHaveProperty("crof"); + }); }); diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index ccfa304078..540527b7d5 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -14,7 +14,6 @@ import coderabbit from "./icons/ProviderIcon-coderabbit.svg?raw"; import codex from "./icons/ProviderIcon-codex.svg?raw"; import commandcode from "./icons/ProviderIcon-commandcode.svg?raw"; import copilot from "./icons/ProviderIcon-copilot.svg?raw"; -import crof from "./icons/ProviderIcon-crof.svg?raw"; import crossmodel from "./icons/ProviderIcon-crossmodel.svg?raw"; import cursor from "./icons/ProviderIcon-cursor.svg?raw"; import deepgram from "./icons/ProviderIcon-deepgram.svg?raw"; @@ -102,7 +101,6 @@ const RAW: Record = { codex: tint(codex), commandcode: tint(commandcode), copilot: tint(copilot), - crof: tint(crof), crossmodel: tint(crossmodel), cursor: tint(cursor), deepgram: tint(deepgram), @@ -218,7 +216,6 @@ export const PROVIDER_ICON_REGISTRY: Record = { mimo: { id: "mimo", brandColor: "#ff6900", fallbackLetter: "M", svgPath: RAW.mimo }, doubao: { id: "doubao", brandColor: "#2563eb", fallbackLetter: "D", svgPath: RAW.doubao }, commandcode: { id: "commandcode", brandColor: "#44ff00", fallbackLetter: "C", svgPath: RAW.commandcode }, - crof: { id: "crof", brandColor: "#7c3aed", fallbackLetter: "C", svgPath: RAW.crof }, crossmodel: { id: "crossmodel", brandColor: "#c084fc", fallbackLetter: "X", svgPath: RAW.crossmodel }, qoder: { id: "qoder", brandColor: "#2563eb", fallbackLetter: "Q", svgPath: RAW.qoder }, replicate: { id: "replicate", brandColor: "#000000", fallbackLetter: "R", svgPath: RAW.replicate }, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..09ead3c1d3 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -30,7 +30,7 @@ import { const HAS_DASHBOARD = new Set([ "abacus", "alibaba", "alibabatokenplan", "amp", "augment", "azureopenai", "bedrock", "claude", "codex", "codebuff", - "aiand", "commandcode", "copilot", "crof", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", + "aiand", "commandcode", "copilot", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", "doubao", "elevenlabs", "factory", "gemini", "grok", "groq", "infini", "jetbrains", "kilo", "kimi", "kimik2", "kiro", "manus", "replicate", "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx index 48c200779c..ccf5295187 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx @@ -242,7 +242,6 @@ function providerSourceHintShort( case "clinepass": case "neuralwatt": case "doubao": - case "crof": case "stepfun": case "venice": case "openaiapi": diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index c8ab2daa5d..dd27d5a469 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -48,7 +48,6 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["mimo", "Xiaomi MiMo"], ["doubao", "Doubao"], ["commandcode", "Command Code"], - ["crof", "Crof"], ["stepfun", "StepFun"], ["venice", "Venice"], ["openaiapi", "OpenAI API"], diff --git a/rust/assets/icons/ProviderIcon-crof.svg b/rust/assets/icons/ProviderIcon-crof.svg deleted file mode 100644 index fdde018b8f..0000000000 --- a/rust/assets/icons/ProviderIcon-crof.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 6ac230cc63..cfdd22c6c6 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -121,10 +121,6 @@ static ICONS: &[(&str, &[u8])] = &[ "ProviderIcon-copilot", include_bytes!("icons/ProviderIcon-copilot.svg"), ), - ( - "ProviderIcon-crof", - include_bytes!("icons/ProviderIcon-crof.svg"), - ), ( "ProviderIcon-cursor", include_bytes!("icons/ProviderIcon-cursor.svg"), diff --git a/rust/src/cli/serve/dashboard/icons/ProviderIcon-crof.svg b/rust/src/cli/serve/dashboard/icons/ProviderIcon-crof.svg deleted file mode 100644 index fdde018b8f..0000000000 --- a/rust/src/cli/serve/dashboard/icons/ProviderIcon-crof.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index d6a38df551..ee639a7ea3 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -55,7 +55,6 @@ pub enum ProviderId { MiMo, Doubao, CommandCode, - Crof, StepFun, Venice, OpenAIApi, @@ -139,7 +138,6 @@ impl ProviderId { ProviderId::MiMo, ProviderId::Doubao, ProviderId::CommandCode, - ProviderId::Crof, ProviderId::StepFun, ProviderId::Venice, ProviderId::OpenAIApi, @@ -226,7 +224,6 @@ impl ProviderId { ProviderId::MiMo => "mimo", ProviderId::Doubao => "doubao", ProviderId::CommandCode => "commandcode", - ProviderId::Crof => "crof", ProviderId::StepFun => "stepfun", ProviderId::Venice => "venice", ProviderId::OpenAIApi => "openaiapi", @@ -310,7 +307,6 @@ impl ProviderId { ProviderId::MiMo => "Xiaomi MiMo", ProviderId::Doubao => "Doubao", ProviderId::CommandCode => "Command Code", - ProviderId::Crof => "Crof", ProviderId::StepFun => "StepFun", ProviderId::Venice => "Venice", ProviderId::OpenAIApi => "OpenAI API", @@ -406,7 +402,6 @@ impl ProviderId { ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, - ProviderId::Crof => None, ProviderId::StepFun => None, ProviderId::OpenAIApi => None, ProviderId::ElevenLabs => None, @@ -490,7 +485,6 @@ impl ProviderId { } "doubao" | "ark" | "volcengine" => Some(ProviderId::Doubao), "commandcode" | "command-code" | "command code" => Some(ProviderId::CommandCode), - "crof" => Some(ProviderId::Crof), "stepfun" | "step-fun" | "step fun" => Some(ProviderId::StepFun), "venice" => Some(ProviderId::Venice), "openaiapi" | "openai-api" | "openai api" | "openai-balance" => { @@ -1048,7 +1042,6 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::MiMo => "#FF6900", ProviderId::Doubao => "#2563EB", ProviderId::CommandCode => "#44FF00", - ProviderId::Crof => "#7C3AED", ProviderId::StepFun => "#999999", ProviderId::Venice => "#111827", ProviderId::OpenAIApi => "#10A37F", @@ -1096,7 +1089,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 77); + assert_eq!(all.len(), 76); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1121,7 +1114,6 @@ mod tests { assert!(all.contains(&ProviderId::MiMo)); assert!(all.contains(&ProviderId::Doubao)); assert!(all.contains(&ProviderId::CommandCode)); - assert!(all.contains(&ProviderId::Crof)); assert!(all.contains(&ProviderId::StepFun)); assert!(all.contains(&ProviderId::Venice)); assert!(all.contains(&ProviderId::OpenAIApi)); @@ -1231,6 +1223,7 @@ mod tests { Some(ProviderId::Antigravity) ); assert_eq!(ProviderId::from_cli_name("zed"), Some(ProviderId::Zed)); + assert_eq!(ProviderId::from_cli_name("crof"), None); assert_eq!(ProviderId::from_cli_name("unknown"), None); assert_eq!( ProviderId::from_cli_name("code-rabbit"), diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index aff4b6f31d..04b139d707 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -10,14 +10,14 @@ use crate::providers::{ AbacusProvider, AiAndProvider, AlibabaProvider, AlibabaTokenPlanProvider, AmpProvider, AntigravityProvider, AugmentProvider, AzureOpenAIProvider, BedrockProvider, ChutesProvider, ClaudeProvider, ClinePassProvider, CodeBuddyProvider, CodeRabbitProvider, CodebuffProvider, - CodexProvider, CommandCodeProvider, CopilotProvider, CrofProvider, CrossModelProvider, - CursorProvider, DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, - DoubaoProvider, ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, - GrokProvider, GroqProvider, HuggingFaceProvider, InfiniProvider, JetBrainsProvider, - KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, - LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, - MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, NousProvider, - OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, + CodexProvider, CommandCodeProvider, CopilotProvider, CrossModelProvider, CursorProvider, + DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, DoubaoProvider, + ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, GrokProvider, + GroqProvider, HuggingFaceProvider, InfiniProvider, JetBrainsProvider, KiloProvider, + KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, LongCatProvider, + ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, MuseProvider, + NanoGPTProvider, NeuralwattProvider, NotionProvider, NousProvider, OllamaProvider, + OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PiProvider, PoeProvider, QoderProvider, QwenCloudProvider, ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, @@ -73,7 +73,6 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::MiMo => Box::new(MiMoProvider::new()), ProviderId::Doubao => Box::new(DoubaoProvider::new()), ProviderId::CommandCode => Box::new(CommandCodeProvider::new()), - ProviderId::Crof => Box::new(CrofProvider::new()), ProviderId::StepFun => Box::new(StepFunProvider::new()), ProviderId::Venice => Box::new(VeniceProvider::new()), ProviderId::OpenAIApi => Box::new(OpenAIApiProvider::new()), diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index a8f3cc71ae..8780e978c5 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -354,7 +354,6 @@ impl TokenAccountSupport { | ProviderId::DeepSeek | ProviderId::Windsurf | ProviderId::Doubao - | ProviderId::Crof | ProviderId::StepFun | ProviderId::Venice | ProviderId::OpenAIApi diff --git a/rust/src/providers/crof/mod.rs b/rust/src/providers/crof/mod.rs deleted file mode 100644 index 4abef6b4ca..0000000000 --- a/rust/src/providers/crof/mod.rs +++ /dev/null @@ -1,249 +0,0 @@ -//! Crof provider implementation. -//! -//! Fetches API key based credit/request quota data from Crof. - -use async_trait::async_trait; -use reqwest::Client; -use serde::Deserialize; - -use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, -}; - -const CROF_USAGE_URL: &str = "https://crof.ai/usage_api/"; -const CROF_CREDENTIAL_TARGET: &str = "codexbar-crof"; -const BROWSER_USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"; - -#[derive(Debug, Deserialize)] -struct CrofUsageResponse { - credits: f64, - #[serde(default, rename = "requests_plan")] - requests_plan: Option, - #[serde(default, rename = "usable_requests")] - usable_requests: Option, -} - -pub struct CrofProvider { - metadata: ProviderMetadata, - client: Client, -} - -impl CrofProvider { - pub fn new() -> Self { - Self { - metadata: ProviderMetadata { - id: ProviderId::Crof, - display_name: "Crof", - session_label: "Balance", - weekly_label: "Requests", - supports_opus: false, - supports_credits: true, - default_enabled: false, - is_primary: false, - dashboard_url: Some("https://crof.ai"), - status_page_url: None, - tertiary_label_key: None, - }, - client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(15)) - .build() - .unwrap_or_else(|_| Client::new()), - } - } - - fn api_key(api_key: Option<&str>) -> Result { - super_key( - api_key, - CROF_CREDENTIAL_TARGET, - &["CROF_API_KEY", "CROFAI_API_KEY"], - ) - } - - async fn fetch_api(&self, api_key: &str) -> Result { - let response = self - .client - .get(CROF_USAGE_URL) - .bearer_auth(api_key) - .header("Accept", "application/json") - .header("User-Agent", BROWSER_USER_AGENT) - .send() - .await?; - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - if status == reqwest::StatusCode::UNAUTHORIZED { - return Err(ProviderError::AuthRequired); - } - if status == reqwest::StatusCode::FORBIDDEN { - if body.contains("cloudflare") || body.contains("Error 1010") { - return Err(ProviderError::Other( - "Crof usage API blocked by Cloudflare (1010). Retry from the desktop app." - .into(), - )); - } - return Err(ProviderError::AuthRequired); - } - if !status.is_success() { - return Err(ProviderError::Other(format!( - "Crof API returned status {status}" - ))); - } - - let usage: CrofUsageResponse = serde_json::from_str(&body) - .map_err(|e| ProviderError::Parse(format!("Failed to parse Crof usage: {e}")))?; - Ok(snapshot_from_usage(&usage)) - } -} - -fn snapshot_from_usage(usage: &CrofUsageResponse) -> UsageSnapshot { - let credits = usage.credits.max(0.0); - let display = if credits <= 0.0 { - "$0.00".to_string() - } else if credits >= 0.01 { - format!("${:.2}", (credits * 100.0).floor() / 100.0) - } else { - format!("${credits:.4}") - }; - let mut primary = RateWindow::new(if credits > 0.0 { 0.0 } else { 100.0 }); - primary.reset_description = Some(display.clone()); - - let mut snapshot = UsageSnapshot::new(primary).with_login_method(format!("{display} balance")); - - if let (Some(plan), Some(usable)) = (usage.requests_plan, usage.usable_requests) { - let remaining = usable.max(0.0).min(plan.max(0.0)); - let remaining_percent = if plan > 0.0 { - ((remaining / plan) * 100.0).clamp(0.0, 100.0) - } else { - 0.0 - }; - let mut requests = RateWindow::new(100.0 - remaining_percent); - requests.reset_description = Some(format!("{remaining:.0} requests left")); - snapshot = snapshot.with_secondary(requests); - } - - snapshot -} - -impl Default for CrofProvider { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Provider for CrofProvider { - fn id(&self) -> ProviderId { - ProviderId::Crof - } - - fn metadata(&self) -> &ProviderMetadata { - &self.metadata - } - - async fn fetch_usage(&self, ctx: &FetchContext) -> Result { - match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => { - let api_key = Self::api_key(ctx.api_key.as_deref())?; - Ok(ProviderFetchResult::new( - self.fetch_api(&api_key).await?, - "api", - )) - } - SourceMode::Web | SourceMode::Cli => { - Err(ProviderError::UnsupportedSource(ctx.source_mode)) - } - } - } - - fn available_sources(&self) -> Vec { - vec![SourceMode::Auto, SourceMode::OAuth] - } -} - -fn super_key( - explicit: Option<&str>, - credential_target: &str, - env_names: &[&str], -) -> Result { - if let Some(key) = explicit - && !key.trim().is_empty() - { - return Ok(key.trim().to_string()); - } - if let Ok(entry) = keyring::Entry::new(credential_target, "api_key") - && let Ok(key) = entry.get_password() - && !key.trim().is_empty() - { - return Ok(key); - } - for env in env_names { - if let Ok(key) = std::env::var(env) - && !key.trim().is_empty() - { - return Ok(key); - } - } - Err(ProviderError::NotInstalled(format!( - "API key not found. Set {} in Preferences or environment.", - env_names.join(" / ") - ))) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn crof_snapshot_formats_request_and_credit_windows() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 12.5, - requests_plan: Some(100.0), - usable_requests: Some(25.0), - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("$12.50") - ); - assert_eq!(snapshot.secondary.unwrap().used_percent, 75.0); - } - - #[test] - fn crof_payg_balance_only_does_not_require_request_quota() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 3.019, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!(snapshot.primary.reset_description.as_deref(), Some("$3.01")); - assert!(snapshot.secondary.is_none()); - assert_eq!(snapshot.login_method.as_deref(), Some("$3.01 balance")); - } - - #[test] - fn crof_sub_cent_balance_is_not_exhausted() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 0.0073, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("$0.0073") - ); - } - - #[test] - fn crof_zero_balance_is_exhausted() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 0.0, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 100.0); - assert_eq!(snapshot.primary.reset_description.as_deref(), Some("$0.00")); - } -} diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index ad697ad321..6bd6432888 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -24,7 +24,6 @@ pub mod coderabbit; pub mod codex; pub mod commandcode; pub mod copilot; -pub mod crof; pub mod crossmodel; pub mod cursor; pub mod deepgram; @@ -104,7 +103,6 @@ pub use coderabbit::CodeRabbitProvider; pub use codex::CodexProvider; pub use commandcode::CommandCodeProvider; pub use copilot::CopilotProvider; -pub use crof::CrofProvider; pub use crossmodel::CrossModelProvider; pub use cursor::CursorProvider; pub use deepgram::DeepgramProvider; diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index 3a7b166386..515a20ad09 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -428,15 +428,6 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: Some("https://console.volcengine.com/ark/region:ark+cn-beijing/usage"), }, - ProviderConfigInfo { - id: ProviderId::Crof, - name: "Crof", - requires_api_key: true, - api_key_env_var: Some("CROF_API_KEY"), - api_key_help: Some("Get your API key from Crof."), - config_file_path: None, - dashboard_url: Some("https://crof.ai"), - }, ProviderConfigInfo { id: ProviderId::StepFun, name: "StepFun", diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..ecd5d6c220 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -1,4 +1,41 @@ use super::*; +use serde::Deserializer; +use serde::de::{IgnoredAny, MapAccess, Visitor}; +use std::fmt; + +fn deserialize_provider_configs<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + struct ProviderConfigsVisitor; + + impl<'de> Visitor<'de> for ProviderConfigsVisitor { + type Value = HashMap; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a map of provider IDs to provider settings") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut configs = HashMap::with_capacity(map.size_hint().unwrap_or(0)); + while let Some(key) = map.next_key::()? { + if let Some(provider_id) = ProviderId::from_cli_name(&key) { + configs.insert(provider_id, map.next_value()?); + } else { + map.next_value::()?; + } + } + Ok(configs) + } + } + + deserializer.deserialize_map(ProviderConfigsVisitor) +} /// Raw on-disk shape of [`Settings`] used purely for deserialization. /// @@ -52,6 +89,7 @@ pub(super) struct RawSettings { show_all_token_accounts_in_menu: bool, // ── New unified per-provider map ───────────────────────────────── + #[serde(default, deserialize_with = "deserialize_provider_configs")] provider_configs: HashMap, // ── Legacy flat per-provider fields (migrated on load) ─────────── @@ -297,6 +335,8 @@ impl Default for RawSettings { impl From for Settings { fn from(raw: RawSettings) -> Self { let mut provider_configs = raw.provider_configs; + let is_known_provider = + |provider_id: &String| ProviderId::from_cli_name(provider_id).is_some(); // Helper closures to lazily insert per-provider configs from legacy // flat fields. Existing `provider_configs` entries take precedence. @@ -515,7 +555,11 @@ impl From for Settings { }; Settings { - enabled_providers: raw.enabled_providers, + enabled_providers: raw + .enabled_providers + .into_iter() + .filter(&is_known_provider) + .collect(), refresh_interval_secs: raw.refresh_interval_secs, adaptive_refresh: raw.adaptive_refresh, refresh_all_providers_on_menu_open: raw.refresh_all_providers_on_menu_open, @@ -549,7 +593,11 @@ impl From for Settings { disable_keychain_access: raw.disable_keychain_access, hide_personal_info: raw.hide_personal_info, update_channel: raw.update_channel, - provider_metrics: raw.provider_metrics, + provider_metrics: raw + .provider_metrics + .into_iter() + .filter(|(provider_id, _)| is_known_provider(provider_id)) + .collect(), provider_order: if raw.provider_order.is_empty() { Vec::new() } else { @@ -578,7 +626,11 @@ impl From for Settings { float_bar_orientation: normalize_float_bar_orientation(&raw.float_bar_orientation), float_bar_style: normalize_float_bar_style(&raw.float_bar_style), float_bar_click_through: raw.float_bar_click_through, - float_bar_provider_ids: raw.float_bar_provider_ids, + float_bar_provider_ids: raw + .float_bar_provider_ids + .into_iter() + .filter(&is_known_provider) + .collect(), float_bar_dark_text: raw.float_bar_dark_text, float_bar_show_reset_inline: raw.float_bar_show_reset_inline, float_bar_show_cost: raw.float_bar_show_cost, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..3e30260d65 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -1045,6 +1045,51 @@ fn test_new_format_provider_configs_only() { assert_eq!(settings.api_region(ProviderId::Zai), "global"); } +#[test] +fn retired_provider_config_is_ignored_until_explicit_save() { + let original = r#"{ + "enabled_providers": ["codex", "crof"], + "refresh_interval_secs": 300, + "provider_metrics": { "codex": "weekly", "crof": "session" }, + "float_bar_provider_ids": ["codex", "crof"], + "provider_configs": { + "crof": { "api_token": "retired-fixture-key" }, + "codex": { "cookie_source": "manual", "openai_web_extras": false }, + "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } + } + }"#; + let original_bytes = original.as_bytes().to_vec(); + + let settings: Settings = + serde_json::from_str(original).expect("load settings with retired key"); + + assert_eq!(original.as_bytes(), original_bytes); + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert!(!settings.openai_web_extras(ProviderId::Codex)); + assert_eq!( + settings.enabled_providers, + HashSet::from(["codex".to_string()]) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex"]); + assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); + assert_eq!( + settings.manual_cookie_header(ProviderId::Alibaba), + "ali=PLACEHOLDER" + ); + + let saved = serde_json::to_string(&settings).expect("serialize sanitized settings"); + let saved_value: serde_json::Value = serde_json::from_str(&saved).unwrap(); + let saved_configs = saved_value["provider_configs"].as_object().unwrap(); + assert!(!saved_configs.contains_key("crof")); + assert!(saved_configs.contains_key("codex")); + assert!(saved_configs.contains_key("alibaba")); + assert!( + !saved.contains("\"crof\""), + "saved settings retained Crof: {saved}" + ); +} + /// Default `Settings` should serialize WITHOUT a `provider_configs` /// field (empty map skipped). #[test] From 1ccb4de0ca04fb367367b15ca28edf53e7943749 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:33:46 +0700 Subject: [PATCH 16/62] Canonicalize loaded provider identifiers --- rust/src/settings/raw.rs | 61 ++++++++++++++++++++++++++++++-------- rust/src/settings/tests.rs | 27 +++++++++++++++++ 2 files changed, 75 insertions(+), 13 deletions(-) diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index ecd5d6c220..abe68d2cfe 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -3,6 +3,51 @@ use serde::Deserializer; use serde::de::{IgnoredAny, MapAccess, Visitor}; use std::fmt; +fn canonical_provider_id(raw: &str) -> Option { + ProviderId::from_cli_name(raw).map(|provider| provider.cli_name().to_string()) +} + +fn canonicalize_provider_id_list(ids: impl IntoIterator) -> Vec { + let mut seen = HashSet::new(); + ids.into_iter() + .filter_map(|raw| canonical_provider_id(&raw)) + .filter(|canonical| seen.insert(canonical.clone())) + .collect() +} + +fn canonicalize_provider_metrics( + metrics: HashMap, +) -> HashMap { + let mut entries = metrics + .into_iter() + .filter_map(|(raw, preference)| { + let canonical = canonical_provider_id(&raw)?; + let canonical_spelling = raw.eq_ignore_ascii_case(&canonical); + Some((canonical, canonical_spelling, raw, preference)) + }) + .collect::>(); + + // HashMap iteration order is unstable. Sort before resolving aliases so a + // canonical spelling always wins and alias-only collisions are repeatable. + entries.sort_by(|left, right| { + left.0 + .cmp(&right.0) + .then_with(|| left.1.cmp(&right.1)) + .then_with(|| { + left.2 + .to_ascii_lowercase() + .cmp(&right.2.to_ascii_lowercase()) + }) + .then_with(|| left.2.cmp(&right.2)) + }); + + let mut canonical = HashMap::with_capacity(entries.len()); + for (provider_id, _, _, preference) in entries { + canonical.insert(provider_id, preference); + } + canonical +} + fn deserialize_provider_configs<'de, D>( deserializer: D, ) -> Result, D::Error> @@ -335,8 +380,6 @@ impl Default for RawSettings { impl From for Settings { fn from(raw: RawSettings) -> Self { let mut provider_configs = raw.provider_configs; - let is_known_provider = - |provider_id: &String| ProviderId::from_cli_name(provider_id).is_some(); // Helper closures to lazily insert per-provider configs from legacy // flat fields. Existing `provider_configs` entries take precedence. @@ -558,7 +601,7 @@ impl From for Settings { enabled_providers: raw .enabled_providers .into_iter() - .filter(&is_known_provider) + .filter_map(|provider_id| canonical_provider_id(&provider_id)) .collect(), refresh_interval_secs: raw.refresh_interval_secs, adaptive_refresh: raw.adaptive_refresh, @@ -593,11 +636,7 @@ impl From for Settings { disable_keychain_access: raw.disable_keychain_access, hide_personal_info: raw.hide_personal_info, update_channel: raw.update_channel, - provider_metrics: raw - .provider_metrics - .into_iter() - .filter(|(provider_id, _)| is_known_provider(provider_id)) - .collect(), + provider_metrics: canonicalize_provider_metrics(raw.provider_metrics), provider_order: if raw.provider_order.is_empty() { Vec::new() } else { @@ -626,11 +665,7 @@ impl From for Settings { float_bar_orientation: normalize_float_bar_orientation(&raw.float_bar_orientation), float_bar_style: normalize_float_bar_style(&raw.float_bar_style), float_bar_click_through: raw.float_bar_click_through, - float_bar_provider_ids: raw - .float_bar_provider_ids - .into_iter() - .filter(&is_known_provider) - .collect(), + float_bar_provider_ids: canonicalize_provider_id_list(raw.float_bar_provider_ids), float_bar_dark_text: raw.float_bar_dark_text, float_bar_show_reset_inline: raw.float_bar_show_reset_inline, float_bar_show_cost: raw.float_bar_show_cost, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index 3e30260d65..d214d6b008 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -1090,6 +1090,33 @@ fn retired_provider_config_is_ignored_until_explicit_save() { ); } +#[test] +fn provider_aliases_are_canonicalized_at_the_load_boundary() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["openai", "ClAuDe", "not-a-provider"], + "provider_metrics": { + "openai": "weekly", + "CoDeX": "session", + "not-a-provider": "weekly" + }, + "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"] + }"#, + ) + .expect("load settings containing provider aliases"); + + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string(), "codex".to_string()]) + ); + assert_eq!( + settings.provider_metrics.get("codex"), + Some(&MetricPreference::Session) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex", "claude"]); +} + /// Default `Settings` should serialize WITHOUT a `provider_configs` /// field (empty map skipped). #[test] From 6193e7c37dbc2f56e884f1957126b59350aeac96 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 17:19:47 +0700 Subject: [PATCH 17/62] Exclude inherited Codex fork baselines --- rust/src/core/jsonl_scanner.rs | 7 + rust/src/core/jsonl_scanner/codex.rs | 61 +++++++- rust/src/core/jsonl_scanner/codex/parser.rs | 158 ++++++++++++++++++++ rust/src/core/jsonl_scanner/tests.rs | 2 + rust/src/cost_scanner/codex.rs | 28 +++- rust/src/cost_scanner/tests/paginated.rs | 157 +++++++++++++++++++ 6 files changed, 408 insertions(+), 5 deletions(-) diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 883a31d89f..13e78bcfe9 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -373,6 +373,8 @@ pub(crate) struct CodexSessionMetadata { pub lineage: CodexSessionLineage, pub fork_timestamp: Option, pub history_base_thread_id: Option, + pub is_subagent: bool, + pub subagent_history_start_ordinal: Option, } /// Running totals for Codex token counting @@ -396,6 +398,10 @@ pub struct CodexForkAccountingState { pub inherited_totals: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub remaining_inherited_totals: Option, + /// True when the child log itself supplied enough copied-prefix history to + /// establish the inherited baseline without consulting a parent cache row. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub locally_resolved: bool, } /// Snapshot of the last validated cost report, persisted so spend surfaces keep @@ -456,6 +462,7 @@ pub struct CodexParseResult { pub fork_baseline: Option, /// Remaining inherited counters used when a fork emits last-only rows. pub remaining_inherited_totals: Option, + pub fork_baseline_locally_resolved: bool, } /// A billable Codex token-count delta. diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 30f159c0ca..1d88cc3f18 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -13,8 +13,9 @@ use parser::CodexParserState; /// Persisted Codex cache schema version. Version 0 predates 64-bit totals; /// version 1 can retain a terminal pause after treating a paginated v2 /// subagent's independent counters as an inherited fork. Version 3 adds -/// persisted paginated-fork accounting state. Rebuild older artifacts. -pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 3; +/// persisted paginated-fork accounting state. Version 4 reparses copied-prefix +/// subagents with locally inferred component baselines. Rebuild older artifacts. +pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 4; /// Whether a persisted Codex cache artifact matches the current schema. /// A mismatched artifact (e.g. a pre-64-bit cache from an older release) is @@ -186,6 +187,12 @@ impl JsonlScanner { .pointer("/source/subagent/thread_spawn") .is_some_and(Value::is_object)) }); + let is_subagent = payload.is_some_and(|value| { + value.get("thread_source").and_then(Value::as_str) == Some("subagent") + || value + .pointer("/source/subagent/thread_spawn") + .is_some_and(Value::is_object) + }); let history_base_thread_id = payload .and_then(|value| value.get("history_base")) .filter(|value| value.is_object()) @@ -220,6 +227,10 @@ impl JsonlScanner { payload.and_then(|value| nonempty_json_string(value.get("timestamp"))) }), history_base_thread_id, + is_subagent, + subagent_history_start_ordinal: payload + .and_then(|value| value.get("subagent_history_start_ordinal")) + .and_then(Value::as_i64), }); } @@ -364,6 +375,8 @@ impl JsonlScanner { None, None, max_bytes_to_read, + false, + None, ) } @@ -400,6 +413,8 @@ impl JsonlScanner { None, scan_target_size, max_bytes_to_read, + false, + None, ) } @@ -431,6 +446,8 @@ impl JsonlScanner { None, None, max_bytes_to_read, + false, + None, ) } @@ -456,6 +473,35 @@ impl JsonlScanner { cancel, scan_target_size, max_bytes_to_read, + false, + None, + ) + } + + pub(crate) fn parse_codex_file_with_inferred_fork_baseline( + file_path: &Path, + range: &CostUsageDayRange, + subagent_history_start_ordinal: Option, + cancel: Option<&AtomicBool>, + scan_target_size: Option, + max_bytes_to_read: Option, + ) -> std::io::Result { + Self::parse_codex_file_with_state_bounded_internal( + file_path, + range, + 0, + None, + None, + None, + None, + cancel, + true, + false, + None, + scan_target_size, + max_bytes_to_read, + true, + subagent_history_start_ordinal, ) } @@ -473,6 +519,8 @@ impl JsonlScanner { cancel: Option<&AtomicBool>, scan_target_size: Option, max_bytes_to_read: Option, + infer_fork_baseline: bool, + subagent_history_start_ordinal: Option, ) -> std::io::Result { Self::parse_codex_file_with_state_bounded_internal( file_path, @@ -488,6 +536,8 @@ impl JsonlScanner { remaining_inherited_totals, scan_target_size, max_bytes_to_read, + infer_fork_baseline, + subagent_history_start_ordinal, ) } @@ -509,6 +559,8 @@ impl JsonlScanner { remaining_inherited_totals: Option, scan_target_size: Option, max_bytes_to_read: Option, + infer_fork_baseline: bool, + subagent_history_start_ordinal: Option, ) -> std::io::Result { let file = File::open(file_path)?; // Session JSONL files are bounded by the cache budget; sizes fit i64. @@ -538,6 +590,9 @@ impl JsonlScanner { paginated_continuation, remaining_inherited_totals, ); + if infer_fork_baseline { + parser.enable_fork_baseline_inference(subagent_history_start_ordinal); + } let mut parsed_bytes = safe_start_offset; let mut committed_bytes = safe_start_offset; let mut cancelled = false; @@ -626,6 +681,7 @@ impl JsonlScanner { }; let is_complete = !cancelled && !budget_exhausted && parsed_bytes >= effective_target_size; let bytes_read = parsed_bytes.saturating_sub(safe_start_offset).max(0); + let fork_baseline_locally_resolved = parser.fork_baseline_locally_resolved(); Ok(CodexParseResult { records: parser.records, parsed_bytes, @@ -644,6 +700,7 @@ impl JsonlScanner { fork_baseline_ambiguous: parser.fork_baseline_ambiguous, fork_baseline: parser.fork_baseline, remaining_inherited_totals: parser.remaining_inherited_totals, + fork_baseline_locally_resolved, }) } diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index 2ec70f9cf1..d81cace664 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -22,6 +22,119 @@ pub(super) struct CodexParserState { paginated_continuation: bool, paginated_baseline_checked: bool, pub(super) fork_baseline_ambiguous: bool, + fork_baseline_inference: Option, +} + +#[derive(Debug)] +struct ForkBaselineInference { + explicit_start_ordinal: Option, + baseline: Option, + boundary_open: bool, + inherited_opening: bool, + locally_confirmed: bool, + resolved: bool, +} + +impl ForkBaselineInference { + fn new(explicit_start_ordinal: Option) -> Self { + Self { + explicit_start_ordinal, + baseline: explicit_start_ordinal.map(|_| CodexTotals { + input: 0, + cached: 0, + output: 0, + reasoning: None, + }), + boundary_open: false, + inherited_opening: false, + locally_confirmed: explicit_start_ordinal.is_some(), + resolved: false, + } + } + + fn observe_non_token(&mut self, obj: &Value) { + if obj.get("type").and_then(Value::as_str) == Some("turn_context") && self.inherited_opening + { + self.boundary_open = true; + } + } + + /// Return the baseline when this is the first owned token event. `None` + /// means the event is still part of the copied prefix. + fn observe_token(&mut self, obj: &Value) -> Option { + let payload = token_count_payload(obj)?; + let info = payload.get("info")?; + let total = read_token_totals(info.get("total_token_usage")?); + let last = read_token_totals(info.get("last_token_usage")?); + let ordinal = obj.get("ordinal").and_then(Value::as_i64); + + if let Some(start) = self.explicit_start_ordinal { + if ordinal.is_some_and(|ordinal| ordinal < start) { + self.baseline = Some(total); + return None; + } + self.boundary_open = true; + } else if self.baseline.is_none() { + if totals_contain_usage(&total) && !totals_contain_usage(&last) { + self.baseline = Some(total); + self.inherited_opening = true; + self.locally_confirmed = true; + } + return None; + } else if !self.boundary_open { + let changed = self + .baseline + .as_ref() + .is_some_and(|baseline| baseline != &total); + if self.inherited_opening && changed && totals_contain_usage(&last) { + self.boundary_open = true; + } else { + return None; + } + } + + let baseline = self.baseline.clone().unwrap_or(CodexTotals { + input: 0, + cached: 0, + output: 0, + reasoning: None, + }); + if total == baseline { + return None; + } + let copied_snapshot = + totals_contain_usage(&baseline) && total == last && totals_at_least(&total, &baseline); + if copied_snapshot { + self.baseline = Some(total); + self.locally_confirmed = true; + return None; + } + + let owned_baseline = totals_delta(&last, &total); + self.baseline = Some(owned_baseline.clone()); + self.locally_confirmed = true; + self.resolved = true; + Some(owned_baseline) + } +} + +fn totals_contain_usage(totals: &CodexTotals) -> bool { + totals.input > 0 || totals.cached > 0 || totals.output > 0 +} + +fn totals_at_least(total: &CodexTotals, baseline: &CodexTotals) -> bool { + total.input >= baseline.input + && total.cached >= baseline.cached + && total.output >= baseline.output +} + +fn totals_delta(last: &CodexTotals, total: &CodexTotals) -> CodexTotals { + CodexTotals { + input: total.input.saturating_sub(last.input).max(0), + cached: total.cached.saturating_sub(last.cached).max(0), + output: total.output.saturating_sub(last.output).max(0), + reasoning: subtract_optional(total.reasoning, last.reasoning), + } } impl CodexParserState { @@ -95,9 +208,24 @@ impl CodexParserState { paginated_continuation, paginated_baseline_checked: false, fork_baseline_ambiguous: false, + fork_baseline_inference: None, } } + pub(super) fn enable_fork_baseline_inference(&mut self, start_ordinal: Option) { + self.fork_baseline = None; + self.remaining_inherited_totals = None; + self.previous_totals = None; + self.totals_watermark = None; + self.fork_baseline_inference = Some(ForkBaselineInference::new(start_ordinal)); + } + + pub(super) fn fork_baseline_locally_resolved(&self) -> bool { + self.fork_baseline_inference + .as_ref() + .is_some_and(|inference| inference.locally_confirmed) + } + pub(super) fn process_line(&mut self, line: &str, range: &CostUsageDayRange) { self.process_line_with_source_offset(line, range, 0); } @@ -108,6 +236,36 @@ impl CodexParserState { range: &CostUsageDayRange, source_end_offset: i64, ) { + if self + .fork_baseline_inference + .as_ref() + .is_some_and(|inference| !inference.resolved) + { + let Ok(obj) = serde_json::from_str::(line) else { + return; + }; + if token_count_payload(&obj).is_some() { + let baseline = self + .fork_baseline_inference + .as_mut() + .and_then(|inference| inference.observe_token(&obj)); + let Some(baseline) = baseline else { return }; + self.fork_baseline = Some(baseline.clone()); + self.remaining_inherited_totals = Some(baseline.clone()); + self.previous_totals = Some(baseline.clone()); + self.totals_watermark = Some(baseline); + } else { + self.fork_baseline_inference + .as_mut() + .expect("inference exists") + .observe_non_token(&obj); + if obj.get("type").and_then(Value::as_str) == Some("turn_context") { + self.update_current_model(&obj); + } + return; + } + } + let event_candidate = is_candidate_codex_line(line); let bare_candidate = !event_candidate && line.contains("\"usage\""); if !event_candidate && !bare_candidate { diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index bcb1129cb1..cc4b2156ef 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -1043,6 +1043,8 @@ fn session_meta_pre_read_accepts_snake_and_camel_fork_identity() { lineage: CodexSessionLineage::Child, fork_timestamp: Some("2026-05-31T10:00:00Z".to_string()), history_base_thread_id: Some("history-snake".to_string()), + is_subagent: false, + subagent_history_start_ordinal: None, } ); diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index f054c6eb8f..cb152611a2 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -42,6 +42,13 @@ fn summary_from_cached_report( } fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { + if usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + { + return true; + } let uses_parent_baseline = usage.codex_lineage.uses_parent_baseline() || (matches!(usage.codex_lineage, CodexSessionLineage::Root) && usage.codex_forked_from_id.is_some()); @@ -456,6 +463,7 @@ impl CostScanner { })) }); let is_fork = codex_lineage.uses_parent_baseline(); + let locally_inferred_subagent = is_fork && session_metadata.is_subagent; let cached_fork_state_matches = cached_fork_accounting_state.as_ref().is_some_and(|state| { state.session_id == codex_session_id @@ -485,7 +493,7 @@ impl CostScanner { .as_deref() .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); - if is_fork && fork_baseline.is_none() { + if is_fork && fork_baseline.is_none() && !locally_inferred_subagent { cache.files.insert( path_key, CostUsageFileUsage { @@ -626,7 +634,16 @@ impl CostScanner { let parse_target_size = cached .as_ref() .and_then(|entry| codex_resumable_scan_target_size(size, entry)); - let parse_result = match if let Some(baseline) = fork_baseline.clone() { + let parse_result = match if locally_inferred_subagent { + JsonlScanner::parse_codex_file_with_inferred_fork_baseline( + path, + range, + session_metadata.subagent_history_start_ordinal, + cancel, + parse_target_size, + max_bytes_to_read, + ) + } else if let Some(baseline) = fork_baseline.clone() { JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( path, range, @@ -636,6 +653,8 @@ impl CostScanner { cancel, parse_target_size, max_bytes_to_read, + false, + None, ) } else { JsonlScanner::parse_codex_file_with_state_bounded( @@ -656,7 +675,9 @@ impl CostScanner { stats.token_timestamp_comparisons = stats .token_timestamp_comparisons .saturating_add(parse_result.token_timestamp_comparisons); - if parse_result.fork_baseline_ambiguous { + if parse_result.fork_baseline_ambiguous + || (locally_inferred_subagent && !parse_result.fork_baseline_locally_resolved) + { cache.files.insert( path_key, CostUsageFileUsage { @@ -707,6 +728,7 @@ impl CostScanner { fork_timestamp: codex_fork_timestamp.clone(), inherited_totals: Some(inherited_totals), remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), + locally_resolved: parse_result.fork_baseline_locally_resolved, }) } else { None diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index 592d484823..7e61f7df55 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -135,6 +135,163 @@ fn write_codex_paginated_continuation_fixture( path } +fn write_copied_prefix_subagent_fixture( + sessions_root: &Path, + name: &str, + base: DateTime, + owned: bool, +) -> PathBuf { + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let mut lines = vec![ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), + "payload": { + "id": "child-id", "forked_from_id": "missing-parent", + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": "missing-parent"}}} + } + }), + token_row(base, 2, [1_000, 900, 100], [0, 0, 0], "gpt-5.6-sol"), + serde_json::json!({ + "type": "turn_context", "ordinal": 10, "timestamp": base.to_rfc3339(), + "payload": {"model": "gpt-5.6-sol"} + }), + token_row( + base, + 12, + [1_000, 900, 100], + [1_000, 900, 100], + "gpt-5.6-sol", + ), + token_row( + base, + 13, + [5_000, 3_900, 500], + [5_000, 3_900, 500], + "gpt-5.6-sol", + ), + ]; + if owned { + lines.extend([ + token_row(base, 19, [5_050, 3_910, 505], [50, 10, 5], "gpt-5.6-sol"), + token_row( + base + Duration::seconds(1), + 20, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + token_row( + base + Duration::seconds(2), + 21, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + ]); + } + let body = lines + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn token_row( + timestamp: DateTime, + ordinal: i64, + total: [i64; 3], + last: [i64; 3], + model: &str, +) -> serde_json::Value { + serde_json::json!({ + "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": model, + "total_token_usage": { + "input_tokens": total[0], "cached_input_tokens": total[1], "output_tokens": total[2] + }, + "last_token_usage": { + "input_tokens": last[0], "cached_input_tokens": last[1], "output_tokens": last[2] + } + }} + }) +} + +#[test] +fn copied_prefix_subagent_infers_advancing_baseline_without_parent() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + Utc::now() - Duration::hours(1), + true, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 70); + assert_eq!(summary.cached_tokens, 15); + assert_eq!(summary.output_tokens, 10); + assert_eq!(summary.sessions_count, 1); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(!usage.codex_unresolved_fork_parent); + assert!( + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + ); + assert_eq!( + usage.days.values().next().unwrap()["gpt-5.6-sol"], + vec![70, 15, 10] + ); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 70); + assert!(stats.codex_history_read_paths.is_empty()); +} + +#[test] +fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + Utc::now() - Duration::hours(1), + false, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.output_tokens, 0); + assert_eq!(summary.sessions_count, 0); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(usage.days.is_empty()); + assert!(!usage.codex_unresolved_fork_parent); +} + #[test] fn paginated_continuation_raises_inherited_baseline_from_total_last() { let root = tempfile::tempdir().unwrap(); From f7d37e8913c7b280ff0147c5ab8a134336357d2d Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:19:50 +0700 Subject: [PATCH 18/62] Model Codex parser modes explicitly --- rust/src/core/jsonl_scanner/codex.rs | 113 ++++-------- rust/src/core/jsonl_scanner/codex/parser.rs | 183 +++++++++++++------- rust/src/core/jsonl_scanner/tests.rs | 24 ++- rust/src/cost_scanner/codex.rs | 2 - 4 files changed, 158 insertions(+), 164 deletions(-) diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 1d88cc3f18..108d4284d7 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -8,7 +8,7 @@ use helpers::{ BoundedJsonlLine, CODEX_JSONL_MAX_LINE_BYTES, nonempty_json_string, parse_rfc3339_timestamp, read_bounded_jsonl_line, read_bounded_jsonl_line_until, session_meta_field, }; -use parser::CodexParserState; +use parser::{CodexParseMode, CodexParserState}; /// Persisted Codex cache schema version. Version 0 predates 64-bit totals; /// version 1 can retain a terminal pause after treating a paginated v2 @@ -364,19 +364,16 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - start_offset, - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, cancel, - false, - false, - None, None, max_bytes_to_read, - false, - None, + CodexParseMode::Standard { + start_offset, + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + }, ) } @@ -402,19 +399,16 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - start_offset, - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, cancel, - false, - false, - None, scan_target_size, max_bytes_to_read, - false, - None, + CodexParseMode::Standard { + start_offset, + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + }, ) } @@ -435,19 +429,14 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - 0, - None, - Some(initial_totals), - None, - None, cancel, - true, - false, - None, None, max_bytes_to_read, - false, - None, + CodexParseMode::ParentBaseline { + baseline: initial_totals, + paginated_continuation: false, + remaining_inherited_totals: None, + }, ) } @@ -473,8 +462,6 @@ impl JsonlScanner { cancel, scan_target_size, max_bytes_to_read, - false, - None, ) } @@ -489,19 +476,12 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - 0, - None, - None, - None, - None, cancel, - true, - false, - None, scan_target_size, max_bytes_to_read, - true, - subagent_history_start_ordinal, + CodexParseMode::InferSubagent { + start_ordinal: subagent_history_start_ordinal, + }, ) } @@ -519,48 +499,28 @@ impl JsonlScanner { cancel: Option<&AtomicBool>, scan_target_size: Option, max_bytes_to_read: Option, - infer_fork_baseline: bool, - subagent_history_start_ordinal: Option, ) -> std::io::Result { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - 0, - None, - Some(initial_totals), - None, - None, cancel, - true, - paginated_continuation, - remaining_inherited_totals, scan_target_size, max_bytes_to_read, - infer_fork_baseline, - subagent_history_start_ordinal, + CodexParseMode::ParentBaseline { + baseline: initial_totals, + paginated_continuation, + remaining_inherited_totals, + }, ) } - #[allow( - clippy::too_many_arguments, - reason = "resume state mirrors the persisted parser cache" - )] fn parse_codex_file_with_state_bounded_internal( file_path: &Path, range: &CostUsageDayRange, - start_offset: i64, - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, cancel: Option<&AtomicBool>, - fork_baseline_mode: bool, - paginated_continuation: bool, - remaining_inherited_totals: Option, scan_target_size: Option, max_bytes_to_read: Option, - infer_fork_baseline: bool, - subagent_history_start_ordinal: Option, + mode: CodexParseMode, ) -> std::io::Result { let file = File::open(file_path)?; // Session JSONL files are bounded by the cache budget; sizes fit i64. @@ -570,7 +530,7 @@ impl JsonlScanner { )] let file_size = file.metadata()?.len() as i64; - let safe_start_offset = start_offset.clamp(0, file_size); + let safe_start_offset = mode.start_offset().clamp(0, file_size); let requested_target_size = scan_target_size .unwrap_or(file_size) .max(safe_start_offset) @@ -581,18 +541,7 @@ impl JsonlScanner { reader.seek(SeekFrom::Start(safe_start_offset as u64))?; } - let mut parser = CodexParserState::with_timestamp_state_and_fork_options( - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, - fork_baseline_mode, - paginated_continuation, - remaining_inherited_totals, - ); - if infer_fork_baseline { - parser.enable_fork_baseline_inference(subagent_history_start_ordinal); - } + let mut parser = CodexParserState::from_mode(mode); let mut parsed_bytes = safe_start_offset; let mut committed_bytes = safe_start_offset; let mut cancelled = false; diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index d81cace664..c3aa7cb332 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -25,6 +25,33 @@ pub(super) struct CodexParserState { fork_baseline_inference: Option, } +pub(super) enum CodexParseMode { + Standard { + start_offset: i64, + initial_model: Option, + initial_totals: Option, + previous_token_timestamp: Option, + token_timestamps_monotonic: Option, + }, + ParentBaseline { + baseline: CodexTotals, + paginated_continuation: bool, + remaining_inherited_totals: Option, + }, + InferSubagent { + start_ordinal: Option, + }, +} + +impl CodexParseMode { + pub(super) fn start_offset(&self) -> i64 { + match self { + Self::Standard { start_offset, .. } => *start_offset, + Self::ParentBaseline { .. } | Self::InferSubagent { .. } => 0, + } + } +} + #[derive(Debug)] struct ForkBaselineInference { explicit_start_ordinal: Option, @@ -35,6 +62,11 @@ struct ForkBaselineInference { resolved: bool, } +enum ForkBaselineDecision { + SkipCopiedPrefix, + ProcessWithBaseline(CodexTotals), +} + impl ForkBaselineInference { fn new(explicit_start_ordinal: Option) -> Self { Self { @@ -59,19 +91,27 @@ impl ForkBaselineInference { } } - /// Return the baseline when this is the first owned token event. `None` - /// means the event is still part of the copied prefix. - fn observe_token(&mut self, obj: &Value) -> Option { - let payload = token_count_payload(obj)?; - let info = payload.get("info")?; - let total = read_token_totals(info.get("total_token_usage")?); - let last = read_token_totals(info.get("last_token_usage")?); + fn observe_token(&mut self, obj: &Value) -> ForkBaselineDecision { + let Some(payload) = token_count_payload(obj) else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(info) = payload.get("info") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(total_usage) = info.get("total_token_usage") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(last_usage) = info.get("last_token_usage") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let total = read_token_totals(total_usage); + let last = read_token_totals(last_usage); let ordinal = obj.get("ordinal").and_then(Value::as_i64); if let Some(start) = self.explicit_start_ordinal { if ordinal.is_some_and(|ordinal| ordinal < start) { self.baseline = Some(total); - return None; + return ForkBaselineDecision::SkipCopiedPrefix; } self.boundary_open = true; } else if self.baseline.is_none() { @@ -80,7 +120,7 @@ impl ForkBaselineInference { self.inherited_opening = true; self.locally_confirmed = true; } - return None; + return ForkBaselineDecision::SkipCopiedPrefix; } else if !self.boundary_open { let changed = self .baseline @@ -89,7 +129,7 @@ impl ForkBaselineInference { if self.inherited_opening && changed && totals_contain_usage(&last) { self.boundary_open = true; } else { - return None; + return ForkBaselineDecision::SkipCopiedPrefix; } } @@ -100,21 +140,21 @@ impl ForkBaselineInference { reasoning: None, }); if total == baseline { - return None; + return ForkBaselineDecision::SkipCopiedPrefix; } let copied_snapshot = totals_contain_usage(&baseline) && total == last && totals_at_least(&total, &baseline); if copied_snapshot { self.baseline = Some(total); self.locally_confirmed = true; - return None; + return ForkBaselineDecision::SkipCopiedPrefix; } let owned_baseline = totals_delta(&last, &total); self.baseline = Some(owned_baseline.clone()); self.locally_confirmed = true; self.resolved = true; - Some(owned_baseline) + ForkBaselineDecision::ProcessWithBaseline(owned_baseline) } } @@ -139,58 +179,74 @@ fn totals_delta(last: &CodexTotals, total: &CodexTotals) -> CodexTotals { impl CodexParserState { pub(super) fn new(initial_model: Option, initial_totals: Option) -> Self { - Self::with_timestamp_state(initial_model, initial_totals, None, None) - } - - fn with_timestamp_state( - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, - ) -> Self { - Self::with_timestamp_state_and_fork_mode( + Self::from_mode(CodexParseMode::Standard { + start_offset: 0, initial_model, initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, - false, - ) + previous_token_timestamp: None, + token_timestamps_monotonic: None, + }) } - pub(super) fn with_timestamp_state_and_fork_mode( - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, - fork_baseline_mode: bool, - ) -> Self { - Self::with_timestamp_state_and_fork_options( + pub(super) fn from_mode(mode: CodexParseMode) -> Self { + let ( initial_model, initial_totals, previous_token_timestamp, token_timestamps_monotonic, - fork_baseline_mode, - false, - None, - ) - } - - pub(super) fn with_timestamp_state_and_fork_options( - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, - fork_baseline_mode: bool, - paginated_continuation: bool, - remaining_inherited_totals: Option, - ) -> Self { + fork_baseline, + paginated_continuation, + remaining_inherited_totals, + fork_baseline_inference, + ) = match mode { + CodexParseMode::Standard { + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + .. + } => ( + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + None, + false, + None, + None, + ), + CodexParseMode::ParentBaseline { + baseline, + paginated_continuation, + remaining_inherited_totals, + } => { + let remaining_inherited_totals = + remaining_inherited_totals.or_else(|| Some(baseline.clone())); + ( + None, + Some(baseline.clone()), + None, + None, + Some(baseline), + paginated_continuation, + remaining_inherited_totals, + None, + ) + } + CodexParseMode::InferSubagent { start_ordinal } => ( + None, + None, + None, + None, + None, + false, + None, + Some(ForkBaselineInference::new(start_ordinal)), + ), + }; let previous_token_timestamp_parsed = previous_token_timestamp .as_deref() .and_then(parse_rfc3339_timestamp); - let fork_baseline = fork_baseline_mode.then(|| initial_totals.clone()).flatten(); - let remaining_inherited_totals = fork_baseline - .as_ref() - .and_then(|baseline| remaining_inherited_totals.or_else(|| Some(baseline.clone()))); Self { current_model: initial_model, previous_totals: initial_totals.clone(), @@ -208,18 +264,10 @@ impl CodexParserState { paginated_continuation, paginated_baseline_checked: false, fork_baseline_ambiguous: false, - fork_baseline_inference: None, + fork_baseline_inference, } } - pub(super) fn enable_fork_baseline_inference(&mut self, start_ordinal: Option) { - self.fork_baseline = None; - self.remaining_inherited_totals = None; - self.previous_totals = None; - self.totals_watermark = None; - self.fork_baseline_inference = Some(ForkBaselineInference::new(start_ordinal)); - } - pub(super) fn fork_baseline_locally_resolved(&self) -> bool { self.fork_baseline_inference .as_ref() @@ -245,11 +293,14 @@ impl CodexParserState { return; }; if token_count_payload(&obj).is_some() { - let baseline = self + let decision = self .fork_baseline_inference .as_mut() - .and_then(|inference| inference.observe_token(&obj)); - let Some(baseline) = baseline else { return }; + .expect("inference exists") + .observe_token(&obj); + let ForkBaselineDecision::ProcessWithBaseline(baseline) = decision else { + return; + }; self.fork_baseline = Some(baseline.clone()); self.remaining_inherited_totals = Some(baseline.clone()); self.previous_totals = Some(baseline.clone()); diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index cc4b2156ef..02d9257753 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -77,13 +77,11 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { output: 10, reasoning: Some(4), }; - let mut state = CodexParserState::with_timestamp_state_and_fork_mode( - None, - Some(baseline), - None, - None, - true, - ); + let mut state = CodexParserState::from_mode(CodexParseMode::ParentBaseline { + baseline, + paginated_continuation: false, + remaining_inherited_totals: None, + }); assert_eq!( state.apply_totals_delta(CodexTotals { input: 20, @@ -100,13 +98,11 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { output: 10, reasoning: None, }; - let mut state = CodexParserState::with_timestamp_state_and_fork_mode( - None, - Some(baseline_without_reasoning), - None, - None, - true, - ); + let mut state = CodexParserState::from_mode(CodexParseMode::ParentBaseline { + baseline: baseline_without_reasoning, + paginated_continuation: false, + remaining_inherited_totals: None, + }); assert_eq!( state.apply_totals_delta(CodexTotals { input: 20, diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index cb152611a2..5e0c4a8b72 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -653,8 +653,6 @@ impl CostScanner { cancel, parse_target_size, max_bytes_to_read, - false, - None, ) } else { JsonlScanner::parse_codex_file_with_state_bounded( From 340de7dabe0c804346d2fc914bb804a2be30f76c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:54:09 +0700 Subject: [PATCH 19/62] Persist inherited-only fork accounting --- rust/src/cost_scanner/codex.rs | 25 ++++++++++++------------ rust/src/cost_scanner/tests/paginated.rs | 9 +++++++++ 2 files changed, 21 insertions(+), 13 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 5e0c4a8b72..5b181311f9 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -715,19 +715,18 @@ impl CostScanner { bytes_read: parse_result.bytes_read, is_complete: parse_result.is_complete, }; - let codex_fork_accounting_state = if is_fork { - parse_result - .fork_baseline - .clone() - .map(|inherited_totals| CodexForkAccountingState { - session_id: codex_session_id.clone(), - forked_from_id: codex_forked_from_id.clone(), - history_base_thread_id: history_base_thread_id.clone(), - fork_timestamp: codex_fork_timestamp.clone(), - inherited_totals: Some(inherited_totals), - remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), - locally_resolved: parse_result.fork_baseline_locally_resolved, - }) + let codex_fork_accounting_state = if is_fork + && (parse_result.fork_baseline.is_some() || parse_result.fork_baseline_locally_resolved) + { + Some(CodexForkAccountingState { + session_id: codex_session_id.clone(), + forked_from_id: codex_forked_from_id.clone(), + history_base_thread_id: history_base_thread_id.clone(), + fork_timestamp: codex_fork_timestamp.clone(), + inherited_totals: parse_result.fork_baseline.clone(), + remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), + locally_resolved: parse_result.fork_baseline_locally_resolved, + }) } else { None }; diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index 7e61f7df55..c13c8255b9 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -290,6 +290,15 @@ fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { let usage = &cache.files[&child.to_string_lossy().to_string()]; assert!(usage.days.is_empty()); assert!(!usage.codex_unresolved_fork_parent); + let state = usage.codex_fork_accounting_state.as_ref().unwrap(); + assert!(state.locally_resolved); + assert!(state.inherited_totals.is_none()); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 0); + assert_eq!(cached.output_tokens, 0); + assert_eq!(cached.sessions_count, 0); + assert!(stats.codex_history_read_paths.is_empty()); } #[test] From 54b69cba351bb6c0bba75c44805a55db1a0de628 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:25:26 +0700 Subject: [PATCH 20/62] Estimate Antigravity local history costs --- .../src-tauri/src/commands/usage_spend.rs | 12 ++- rust/src/cli/cost.rs | 27 ++++++- rust/src/cli/serve/data.rs | 2 + .../providers/antigravity/local_sessions.rs | 80 ++++++++++++++++++- .../src/providers/antigravity/local_sqlite.rs | 41 ++++++++++ rust/src/providers/muse/local_usage/mod.rs | 1 + rust/src/spend_contract.rs | 16 +++- 7 files changed, 172 insertions(+), 7 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index 01f562e111..f2dab7fb08 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -590,12 +590,20 @@ fn build_usage_spend_summary( let seven = codexbar::providers::antigravity::local_sessions::summarize(7); let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); let mut spend = cached_spend(cached_snapshot); + spend.seven_day = seven.estimated_cost_usd; + spend.thirty_day = thirty.estimated_cost_usd; spend.seven_day_tokens = matches!(seven.coverage, LocalHistoryCoverage::Complete) .then_some(seven.total_tokens); spend.thirty_day_tokens = matches!(thirty.coverage, LocalHistoryCoverage::Complete) .then_some(thirty.total_tokens); - if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - spend.source = "local Antigravity history".to_string(); + if thirty.estimated_cost_usd.is_some() { + spend.source = if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { + "local Antigravity history · API list-price estimate".to_string() + } else { + "partial local Antigravity history · API list-price estimate".to_string() + }; + } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { + spend.source = "local Antigravity history · unpriced".to_string(); } spend } diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 8e9be16dc2..86dd742b05 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -386,7 +386,11 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - println!(" Local token history; dollar costs unavailable"); + if let Some(cost) = history.estimated_cost_usd { + println!(" API list-price estimate: ${cost:.2} (not billed spend)"); + } else { + println!(" Local token history; dollar costs unavailable"); + } } fn print_codex_session_output(result: &CostResult, days: u32) { @@ -625,6 +629,7 @@ mod tests { total_tokens: 12_345, session_count: 2, coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: None, }, 30, ); @@ -639,6 +644,7 @@ mod tests { total_tokens: 999, session_count: 1, coverage: LocalHistoryCoverage::Partial, + estimated_cost_usd: None, }, 30, ); @@ -646,6 +652,25 @@ mod tests { assert!(partial["tokens"]["total"].is_null()); assert_eq!(partial["historyCoverage"], "partial"); } + + #[test] + fn antigravity_json_labels_public_price_estimates() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + LocalTokenHistorySummary { + total_tokens: 1_000, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: Some(0.0125), + }, + 30, + ); + assert_eq!(payload["cost"]["total_usd"], 0.0125); + assert_eq!(payload["cost"]["currency"], "USD"); + assert!(payload["note"].as_str().unwrap().contains("not billed")); + } + #[test] fn provider_native_only_flag_default_false() { // Default CostArgs has provider_native_only = false (backward compat). diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index d1841db6bc..b615c69150 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -164,6 +164,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: None, }, 30, ); @@ -177,6 +178,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, + estimated_cost_usd: None, }, 30, ); diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 458bb97f69..ff1bc5daed 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -6,6 +6,8 @@ use std::path::{Path, PathBuf}; use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use serde_json::Value; +use crate::core::CostUsagePricing; + const MAX_SESSION_FILES: usize = 2048; const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; @@ -143,6 +145,7 @@ fn summarize_paths( let first_day = now.with_timezone(&Local).date_naive() - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); let mut total_tokens = 0_u64; + let mut estimated_cost_usd = None; let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); let mut complete = !truncated; @@ -163,6 +166,7 @@ fn summarize_paths( let mut reader = BufReader::new(file); let mut remaining = MAX_SESSION_FILE_BYTES; let mut path_had_usage = false; + let mut model = None::; loop { let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { Ok(Some(line)) => line, @@ -179,6 +183,16 @@ fn summarize_paths( continue; }; let kind = value.get("type").and_then(Value::as_str); + if kind == Some("session_meta") { + model = value + .get("modelId") + .or_else(|| value.get("model_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + continue; + } if kind != Some("usage") && value.get("input").is_none() { continue; } @@ -208,14 +222,28 @@ fn summarize_paths( let output = token_field(&value, &["output"]); let cache_read = token_field(&value, &["cacheRead", "cache_read"]); let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); + let reasoning = token_field( + &value, + &["reasoning", "reasoningTokens", "reasoning_tokens"], + ); let total = input .saturating_add(output) .saturating_add(cache_read) - .saturating_add(cache_write); + .saturating_add(cache_write) + .saturating_add(reasoning); if total == 0 { continue; } total_tokens = total_tokens.saturating_add(total); + if let Some(cost) = estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + ) { + estimated_cost_usd = checked_cost_sum(estimated_cost_usd, cost); + } path_had_usage = true; } if path_had_usage { @@ -233,9 +261,40 @@ fn summarize_paths( } else { LocalHistoryCoverage::Partial }, + estimated_cost_usd, } } +pub(super) fn estimate_cost_usd( + model: Option<&str>, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> Option { + let model = model.map(str::trim).filter(|value| !value.is_empty())?; + let input = i32::try_from(input).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let output = i32::try_from(output).ok()?; + let resolve = |candidate: &str| { + CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + }; + resolve(model).or_else(|| { + ["-tiered", "-low", "-thinking"] + .iter() + .find_map(|suffix| model.strip_suffix(suffix)) + .filter(|base| !base.is_empty()) + .and_then(resolve) + }) +} + +pub(super) fn checked_cost_sum(current: Option, cost: f64) -> Option { + let next = current.unwrap_or(0.0) + cost; + next.is_finite().then_some(next) +} + fn read_bounded_jsonl_line( reader: &mut R, remaining_file_bytes: &mut usize, @@ -290,6 +349,25 @@ fn token_field(value: &Value, keys: &[&str]) -> u64 { #[cfg(test)] mod tests { use super::*; + + #[test] + fn prices_known_models_and_provider_local_routing_variants() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) + .expect("routing suffix uses the base public price"); + assert!(direct > 0.0); + assert_eq!(direct, routed); + } + + #[test] + fn unknown_or_oversized_pricing_inputs_fail_closed() { + assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), + None + ); + } use rusqlite::Connection; #[test] diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index b050660411..ea8b3d3dd5 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -184,10 +184,26 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let mut total_tokens = 0_u64; + let mut estimated_cost_usd = None; let mut sessions = HashSet::new(); let mut rows: HashMap<(String, i64), Event> = HashMap::new(); let mut responses: HashMap<(String, String), Event> = HashMap::new(); + let mut label_models = HashMap::<(String, String), String>::new(); + let mut conflicting_labels = HashSet::<(String, String)>::new(); + for event in &events { + let (Some(label), Some(model)) = (event.turn.label.as_ref(), event.turn.model.as_ref()) + else { + continue; + }; + let key = (event.session.clone(), label.clone()); + if label_models.get(&key).is_some_and(|prior| prior != model) { + conflicting_labels.insert(key); + } else { + label_models.insert(key, model.clone()); + } + } + for event in events { let row_key = (event.session.clone(), event.row); if let Some(prior) = rows.get(&row_key) { @@ -234,6 +250,30 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } + if let Some(usage) = event.turn.usage.as_ref() { + let inherited_model = event.turn.label.as_ref().and_then(|label| { + let key = (event.session.clone(), label.clone()); + (!conflicting_labels.contains(&key)) + .then(|| label_models.get(&key)) + .flatten() + .map(String::as_str) + }); + let model = event.turn.model.as_deref().or(inherited_model); + let input = usage.system_prompt.checked_add(usage.new_input); + let output = usage.output.checked_add(usage.reasoning); + if let (Some(input), Some(output)) = (input, output) + && let Some(cost) = super::local_sessions::estimate_cost_usd( + model, + input, + usage.cache_read, + 0, + output, + ) + { + estimated_cost_usd = + super::local_sessions::checked_cost_sum(estimated_cost_usd, cost); + } + } sessions.insert(event.session); } @@ -245,6 +285,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } else { LocalHistoryCoverage::Partial }, + estimated_cost_usd, }) } diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 49a2f7ec34..7ee5b3f093 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -64,6 +64,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { total_tokens: report.total_tokens.unwrap_or(0), session_count: report.session_count, coverage: report.coverage, + estimated_cost_usd: None, } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index af9cdfef34..f43bcd0529 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -79,11 +79,14 @@ pub enum LocalHistoryCoverage { Unavailable, } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq)] pub struct LocalTokenHistorySummary { pub total_tokens: u64, pub session_count: usize, pub coverage: LocalHistoryCoverage, + /// Known subtotal priced at public API list rates. `None` means that no + /// local request in the selected window had a resolvable model price. + pub estimated_cost_usd: Option, } pub fn local_token_history_json( @@ -96,7 +99,10 @@ pub fn local_token_history_json( "provider": provider, "supported": true, "days_scanned": days, - "cost": {"total_usd": serde_json::Value::Null, "currency": serde_json::Value::Null}, + "cost": { + "total_usd": history.estimated_cost_usd, + "currency": history.estimated_cost_usd.map(|_| "USD") + }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, "sessions_count": complete.then_some(history.session_count), @@ -106,7 +112,11 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": "Local token history; dollar costs unavailable" + "note": if history.estimated_cost_usd.is_some() { + "Local token history estimated at public API list prices; not billed spend" + } else { + "Local token history; dollar costs unavailable" + } }) } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] From 7bb133aa070e60296554a2951fdc701530e3bc7a Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:07:37 +0700 Subject: [PATCH 21/62] Track local pricing coverage --- .../src-tauri/src/commands/usage_spend.rs | 27 +++++++--- .../src/lib/usageSpendSharing.test.ts | 5 ++ .../src/lib/usageSpendSharing.ts | 23 +++++++-- .../src/surfaces/TrayPanel.test.tsx | 16 +++++- .../surfaces/settings/tabs/UsageSpendTab.tsx | 20 +++++++- apps/desktop-tauri/src/types/bridge.ts | 7 +++ rust/src/cli/cost.rs | 49 ++++++++++++++++-- rust/src/cli/serve/data.rs | 4 +- .../providers/antigravity/local_sessions.rs | 46 +++++++++++++---- .../src/providers/antigravity/local_sqlite.rs | 24 ++++----- rust/src/providers/muse/local_usage/mod.rs | 2 +- rust/src/spend_contract.rs | 51 ++++++++++++++++--- 12 files changed, 220 insertions(+), 54 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index f2dab7fb08..49d84b7b34 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -26,6 +26,10 @@ pub struct UsageSpendRow { pub display_name: String, pub seven_day: Option, pub thirty_day: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub seven_day_estimate: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thirty_day_estimate: Option, pub seven_day_tokens: Option, pub thirty_day_tokens: Option, pub currency: String, @@ -496,6 +500,7 @@ fn build_usage_spend_summary( }) .unwrap_or_else(|| provider_id.clone()); + let mut local_cost_estimates = None; let spend = match provider_id.as_str() { "codex" => SpendValues { seven_day: codex_7_contract.known_cost_usd, @@ -590,21 +595,22 @@ fn build_usage_spend_summary( let seven = codexbar::providers::antigravity::local_sessions::summarize(7); let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); let mut spend = cached_spend(cached_snapshot); - spend.seven_day = seven.estimated_cost_usd; - spend.thirty_day = thirty.estimated_cost_usd; + spend.seven_day = seven.cost_estimate.total_usd(); + spend.thirty_day = thirty.cost_estimate.total_usd(); spend.seven_day_tokens = matches!(seven.coverage, LocalHistoryCoverage::Complete) .then_some(seven.total_tokens); spend.thirty_day_tokens = matches!(thirty.coverage, LocalHistoryCoverage::Complete) .then_some(thirty.total_tokens); - if thirty.estimated_cost_usd.is_some() { - spend.source = if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - "local Antigravity history · API list-price estimate".to_string() - } else { - "partial local Antigravity history · API list-price estimate".to_string() - }; + if thirty.cost_estimate.total_usd().is_some() { + spend.source = + "local Antigravity history · API list-price estimate".to_string(); + } else if thirty.cost_estimate.known_subtotal_usd.is_some() { + spend.source = + "local Antigravity history · known API list-price subtotal".to_string(); } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { spend.source = "local Antigravity history · unpriced".to_string(); } + local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } _ => cached_spend(cached_snapshot), @@ -626,11 +632,16 @@ fn build_usage_spend_summary( .collect() }) .unwrap_or_default(); + let (seven_day_estimate, thirty_day_estimate) = local_cost_estimates + .map(|(seven, thirty)| (Some(seven), Some(thirty))) + .unwrap_or((None, None)); rows.push(UsageSpendRow { provider_id: provider_id.clone(), display_name, seven_day: spend.seven_day, thirty_day: spend.thirty_day, + seven_day_estimate, + thirty_day_estimate, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, currency, diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index bf693ce7b8..d6748ff079 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { formatUsageSpendReportingDay, + formatSpendMetric, filterUsageSpendSummaryForOverview, renderUsageSpendSharePng, usageSpendShareFooter, @@ -10,6 +11,10 @@ import { import type { SpendContract, UsageSpendRow, UsageSpendSummary } from "../types/bridge"; describe("usage spend sharing", () => { + it("labels a mixed-pricing subtotal without presenting it as a total", () => { + expect(formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125)).toMatch(/^≥.* known/); + }); + it.each([ [0, "0 subscriptions"], [1, "1 subscription"], diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 2246e98505..4045eb9199 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -130,9 +130,14 @@ export function formatSpendMetric( tokens: number | null | undefined, currency: string, tokenLabel: string, + knownSubtotal?: number | null, ): string { const parts: string[] = []; - if (cost != null && Number.isFinite(cost)) parts.push(formatUsd(cost, currency)); + if (cost != null && Number.isFinite(cost)) { + parts.push(formatUsd(cost, currency)); + } else if (knownSubtotal != null && Number.isFinite(knownSubtotal)) { + parts.push(`≥${formatUsd(knownSubtotal, currency)} known`); + } if (tokens != null && Number.isFinite(tokens)) { parts.push(`${Math.max(0, tokens).toLocaleString()} ${tokenLabel}`); } @@ -202,8 +207,20 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri const y = y0 + (index + 1) * rowH; const cells = [ row.displayName, - formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, "tokens"), - formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, "tokens"), + formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + "tokens", + row.sevenDayEstimate?.knownSubtotalUsd, + ), + formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + "tokens", + row.thirtyDayEstimate?.knownSubtotalUsd, + ), row.currency || "USD", row.source, ]; diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 09a5b01e1c..ecefece9f2 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -339,13 +339,27 @@ describe("TrayPanel provider grid", () => { source: "hidden", includedInOverview: false, }, + { + providerId: "antigravity", + displayName: "Antigravity", + sevenDay: null, + thirtyDay: null, + thirtyDayEstimate: { + knownSubtotalUsd: 9, + coverage: { priced: 0, unpriced: 1, unmetered: 0, estimated: 1 }, + }, + currency: "USD", + source: "known subtotal", + includedInOverview: true, + }, ], }); renderTrayPanel([provider("codex", "Codex", 35)]); expect(await screen.findByRole("button", { name: "UsageSpendShare" })).toBeInTheDocument(); - expect(screen.getByText(/1 of 1 OverviewSpendProviderCoverage/)).toBeInTheDocument(); + expect(screen.getByText("$2.00")).toBeInTheDocument(); + expect(screen.getByText(/1 of 2 OverviewSpendProviderCoverage/)).toBeInTheDocument(); }); it("dismisses the tray panel on unmodified Escape", async () => { diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index 0f7cc7ee89..e7ee282ece 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -234,8 +234,24 @@ export default function UsageSpendTab(_props: TabProps) { {(summary?.rows ?? []).map((row) => ( {row.displayName} - {formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, t("UsageSpendTokens"))} - {formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, t("UsageSpendTokens"))} + + {formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + t("UsageSpendTokens"), + row.sevenDayEstimate?.knownSubtotalUsd, + )} + + + {formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + t("UsageSpendTokens"), + row.thirtyDayEstimate?.knownSubtotalUsd, + )} + {row.currency || "USD"} {row.source} diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 5aa0a40f27..d581262477 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -379,6 +379,8 @@ export interface UsageSpendRow { displayName: string; sevenDay: number | null; thirtyDay: number | null; + sevenDayEstimate?: LocalCostEstimate; + thirtyDayEstimate?: LocalCostEstimate; sevenDayTokens?: number | null; thirtyDayTokens?: number | null; currency: string; @@ -391,6 +393,11 @@ export interface UsageSpendRow { staleUpdatedAt?: string; } +export interface LocalCostEstimate { + knownSubtotalUsd: number | null; + coverage: CostCoverageCounts; +} + export interface UsageSpendSummary { rows: UsageSpendRow[]; contract: SpendContract; diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 86dd742b05..30eba40a42 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -386,8 +386,13 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - if let Some(cost) = history.estimated_cost_usd { + if let Some(cost) = history.cost_estimate.total_usd() { println!(" API list-price estimate: ${cost:.2} (not billed spend)"); + } else if let Some(cost) = history.cost_estimate.known_subtotal_usd { + println!( + " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", + history.cost_estimate.coverage.unpriced + ); } else { println!(" Local token history; dollar costs unavailable"); } @@ -629,7 +634,7 @@ mod tests { total_tokens: 12_345, session_count: 2, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -644,7 +649,7 @@ mod tests { total_tokens: 999, session_count: 1, coverage: LocalHistoryCoverage::Partial, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -662,15 +667,51 @@ mod tests { total_tokens: 1_000, session_count: 1, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: Some(0.0125), + cost_estimate: crate::spend_contract::LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: crate::spend_contract::CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + }, }, 30, ); assert_eq!(payload["cost"]["total_usd"], 0.0125); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); assert_eq!(payload["cost"]["currency"], "USD"); assert!(payload["note"].as_str().unwrap().contains("not billed")); } + #[test] + fn antigravity_json_keeps_mixed_pricing_as_a_known_subtotal() { + use crate::spend_contract::{ + CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, + }; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + LocalTokenHistorySummary { + total_tokens: 1_500, + session_count: 2, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + }, + }, + 30, + ); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); + assert_eq!(payload["cost"]["pricingCoverage"]["estimated"], 1); + assert_eq!(payload["cost"]["pricingCoverage"]["unpriced"], 1); + assert!(payload["note"].as_str().unwrap().contains("subtotal")); + } + #[test] fn provider_native_only_flag_default_false() { // Default CostArgs has provider_native_only = false (backward compat). diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index b615c69150..aa12414d6e 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -164,7 +164,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -178,7 +178,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index ff1bc5daed..3ab0e66e9c 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -145,7 +145,7 @@ fn summarize_paths( let first_day = now.with_timezone(&Local).date_naive() - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); let mut total_tokens = 0_u64; - let mut estimated_cost_usd = None; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); let mut complete = !truncated; @@ -235,15 +235,13 @@ fn summarize_paths( continue; } total_tokens = total_tokens.saturating_add(total); - if let Some(cost) = estimate_cost_usd( + cost_estimate.record_list_price(estimate_cost_usd( model.as_deref(), input, cache_read, cache_write, output.saturating_add(reasoning), - ) { - estimated_cost_usd = checked_cost_sum(estimated_cost_usd, cost); - } + )); path_had_usage = true; } if path_had_usage { @@ -261,7 +259,7 @@ fn summarize_paths( } else { LocalHistoryCoverage::Partial }, - estimated_cost_usd, + cost_estimate, } } @@ -290,11 +288,6 @@ pub(super) fn estimate_cost_usd( }) } -pub(super) fn checked_cost_sum(current: Option, cost: f64) -> Option { - let next = current.unwrap_or(0.0) + cost; - next.is_finite().then_some(next) -} - fn read_bounded_jsonl_line( reader: &mut R, remaining_file_bytes: &mut usize, @@ -368,6 +361,37 @@ mod tests { None ); } + + #[test] + fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { + let dir = tempfile::tempdir().unwrap(); + let known = dir.path().join("known.jsonl"); + let unknown = dir.path().join("unknown.jsonl"); + fs::write( + &known, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", + "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" + ), + ) + .unwrap(); + fs::write( + &unknown, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", + "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[known, unknown], now, 7, false); + + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert!(summary.cost_estimate.known_subtotal_usd.is_some()); + assert_eq!(summary.cost_estimate.total_usd(), None); + } use rusqlite::Connection; #[test] diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index ea8b3d3dd5..8e83eb8e33 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -184,7 +184,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let mut total_tokens = 0_u64; - let mut estimated_cost_usd = None; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions = HashSet::new(); let mut rows: HashMap<(String, i64), Event> = HashMap::new(); let mut responses: HashMap<(String, String), Event> = HashMap::new(); @@ -250,7 +250,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } - if let Some(usage) = event.turn.usage.as_ref() { + let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { let inherited_model = event.turn.label.as_ref().and_then(|label| { let key = (event.session.clone(), label.clone()); (!conflicting_labels.contains(&key)) @@ -261,19 +261,13 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL let model = event.turn.model.as_deref().or(inherited_model); let input = usage.system_prompt.checked_add(usage.new_input); let output = usage.output.checked_add(usage.reasoning); - if let (Some(input), Some(output)) = (input, output) - && let Some(cost) = super::local_sessions::estimate_cost_usd( - model, - input, - usage.cache_read, - 0, - output, - ) - { - estimated_cost_usd = - super::local_sessions::checked_cost_sum(estimated_cost_usd, cost); + if let (Some(input), Some(output)) = (input, output) { + super::local_sessions::estimate_cost_usd(model, input, usage.cache_read, 0, output) + } else { + None } - } + }); + cost_estimate.record_list_price(estimated_cost); sessions.insert(event.session); } @@ -285,7 +279,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } else { LocalHistoryCoverage::Partial }, - estimated_cost_usd, + cost_estimate, }) } diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 7ee5b3f093..24e84b2473 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -64,7 +64,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { total_tokens: report.total_tokens.unwrap_or(0), session_count: report.session_count, coverage: report.coverage, - estimated_cost_usd: None, + cost_estimate: Default::default(), } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index f43bcd0529..54c7d3490a 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -79,14 +79,45 @@ pub enum LocalHistoryCoverage { Unavailable, } -#[derive(Debug, Clone, Copy, Default, PartialEq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct LocalTokenHistorySummary { pub total_tokens: u64, pub session_count: usize, pub coverage: LocalHistoryCoverage, - /// Known subtotal priced at public API list rates. `None` means that no - /// local request in the selected window had a resolvable model price. - pub estimated_cost_usd: Option, + pub cost_estimate: LocalCostEstimate, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LocalCostEstimate { + /// Sum of requests whose models have public API list prices. This remains + /// a subtotal when one or more requests are unpriced. + pub known_subtotal_usd: Option, + pub coverage: CostCoverageCounts, +} + +impl LocalCostEstimate { + pub fn total_usd(&self) -> Option { + if self.coverage.unpriced == 0 && self.coverage.unmetered == 0 { + self.known_subtotal_usd + } else { + None + } + } + + pub(crate) fn record_list_price(&mut self, cost: Option) { + let Some(cost) = cost.filter(|value| value.is_finite() && *value >= 0.0) else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + return; + }; + let next = self.known_subtotal_usd.unwrap_or(0.0) + cost; + if next.is_finite() { + self.known_subtotal_usd = Some(next); + self.coverage.estimated = self.coverage.estimated.saturating_add(1); + } else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + } + } } pub fn local_token_history_json( @@ -95,13 +126,17 @@ pub fn local_token_history_json( days: u32, ) -> serde_json::Value { let complete = history.coverage == LocalHistoryCoverage::Complete; + let total_usd = history.cost_estimate.total_usd(); + let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; serde_json::json!({ "provider": provider, "supported": true, "days_scanned": days, "cost": { - "total_usd": history.estimated_cost_usd, - "currency": history.estimated_cost_usd.map(|_| "USD") + "total_usd": total_usd, + "known_subtotal_usd": known_subtotal_usd, + "currency": known_subtotal_usd.map(|_| "USD"), + "pricingCoverage": history.cost_estimate.coverage, }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, @@ -112,8 +147,10 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": if history.estimated_cost_usd.is_some() { + "note": if total_usd.is_some() { "Local token history estimated at public API list prices; not billed spend" + } else if known_subtotal_usd.is_some() { + "Known public API list-price subtotal; some local requests are unpriced" } else { "Local token history; dollar costs unavailable" } From 33697f174c69a7148aecd6cbb1d7cd65aefa95df Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:47:09 +0700 Subject: [PATCH 22/62] Fix local pricing coverage totals --- .../src-tauri/src/commands/usage_spend.rs | 93 +++++++++++++++---- .../src/lib/usageSpendSharing.test.ts | 8 ++ rust/src/cli/cost.rs | 77 ++++++++++++--- rust/src/cli/serve/data.rs | 11 +-- .../providers/antigravity/local_sessions.rs | 2 +- rust/src/spend_contract.rs | 42 ++++++--- rust/src/spend_contract/tests.rs | 53 +++++++++++ 7 files changed, 238 insertions(+), 48 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index 49d84b7b34..ee7ae528c8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -591,25 +591,10 @@ fn build_usage_spend_summary( spend } "antigravity" => { - use codexbar::providers::antigravity::local_sessions::LocalHistoryCoverage; let seven = codexbar::providers::antigravity::local_sessions::summarize(7); let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); - let mut spend = cached_spend(cached_snapshot); - spend.seven_day = seven.cost_estimate.total_usd(); - spend.thirty_day = thirty.cost_estimate.total_usd(); - spend.seven_day_tokens = matches!(seven.coverage, LocalHistoryCoverage::Complete) - .then_some(seven.total_tokens); - spend.thirty_day_tokens = matches!(thirty.coverage, LocalHistoryCoverage::Complete) - .then_some(thirty.total_tokens); - if thirty.cost_estimate.total_usd().is_some() { - spend.source = - "local Antigravity history · API list-price estimate".to_string(); - } else if thirty.cost_estimate.known_subtotal_usd.is_some() { - spend.source = - "local Antigravity history · known API list-price subtotal".to_string(); - } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - spend.source = "local Antigravity history · unpriced".to_string(); - } + let spend = + antigravity_spend_values(cached_spend(cached_snapshot), &seven, &thirty); local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } @@ -720,6 +705,29 @@ fn total_token_mix(mix: &codexbar::spend_contract::SpendTokenMix) -> Option saw.then_some(total) } +fn antigravity_spend_values( + mut spend: SpendValues, + seven: &codexbar::spend_contract::LocalTokenHistorySummary, + thirty: &codexbar::spend_contract::LocalTokenHistorySummary, +) -> SpendValues { + use codexbar::spend_contract::LocalHistoryCoverage; + + spend.seven_day = seven.total_usd(); + spend.thirty_day = thirty.total_usd(); + spend.seven_day_tokens = + (seven.coverage == LocalHistoryCoverage::Complete).then_some(seven.total_tokens); + spend.thirty_day_tokens = + (thirty.coverage == LocalHistoryCoverage::Complete).then_some(thirty.total_tokens); + if spend.thirty_day.is_some() { + spend.source = "local Antigravity history · API list-price estimate".to_string(); + } else if thirty.cost_estimate.known_subtotal_usd.is_some() { + spend.source = "local Antigravity history · known API list-price subtotal".to_string(); + } else if thirty.coverage == LocalHistoryCoverage::Complete { + spend.source = "local Antigravity history · unpriced".to_string(); + } + spend +} + fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { let Some(snapshot) = snapshot else { return SpendValues { @@ -797,6 +805,27 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { mod cache_key_tests { use super::*; + fn local_history( + total_tokens: u64, + coverage: codexbar::spend_contract::LocalHistoryCoverage, + known_subtotal_usd: Option, + unpriced: u32, + ) -> codexbar::spend_contract::LocalTokenHistorySummary { + codexbar::spend_contract::LocalTokenHistorySummary { + total_tokens, + session_count: if total_tokens > 0 { 1 } else { 0 }, + coverage, + cost_estimate: codexbar::spend_contract::LocalCostEstimate { + known_subtotal_usd, + coverage: codexbar::spend_contract::CostCoverageCounts { + estimated: if known_subtotal_usd.is_some() { 1 } else { 0 }, + unpriced, + ..Default::default() + }, + }, + } + } + #[test] fn invalidated_owner_clears_orphaned_indexing_activity() { let mut coordinator = UsageSpendCoordinator::default(); @@ -853,4 +882,34 @@ mod cache_key_tests { assert!(include_in_shared_overview("claude", false, true)); assert!(!include_in_shared_overview("codex", false, false)); } + + #[test] + fn antigravity_partial_history_exposes_only_the_known_subtotal() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(100, LocalHistoryCoverage::Partial, Some(1.25), 0); + let thirty = local_history(200, LocalHistoryCoverage::Partial, Some(2.50), 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, None); + assert_eq!(spend.thirty_day, None); + assert_eq!(spend.seven_day_tokens, None); + assert_eq!(spend.thirty_day_tokens, None); + assert!(spend.source.contains("known API list-price subtotal")); + } + + #[test] + fn antigravity_complete_empty_history_is_a_known_zero() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let thirty = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, Some(0.0)); + assert_eq!(spend.thirty_day, Some(0.0)); + assert_eq!(spend.seven_day_tokens, Some(0)); + assert_eq!(spend.thirty_day_tokens, Some(0)); + assert!(spend.source.contains("API list-price estimate")); + } } diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index d6748ff079..39c95c9a3f 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -15,6 +15,14 @@ describe("usage spend sharing", () => { expect(formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125)).toMatch(/^≥.* known/); }); + it("renders a complete known-zero total instead of a subtotal", () => { + const metric = formatSpendMetric(0, 0, "USD", "tokens", 9); + expect(metric).not.toBe("—"); + expect(metric).not.toContain("≥"); + expect(metric).not.toContain("9.00"); + expect(metric).toContain("0 tokens"); + }); + it.each([ [0, "0 subscriptions"], [1, "1 subscription"], diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 30eba40a42..819b7abb59 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -286,7 +286,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b println!("{title}"); } - if let Some(history) = result.token_history { + if let Some(history) = result.token_history.as_ref() { print_local_token_history(history, days); } else if group_by == CostGroupBy::Session && result.provider == "codex" { print_codex_session_output(result, days); @@ -372,7 +372,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b } } -fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySummary, days: u32) { +fn print_local_token_history(history: &crate::spend_contract::LocalTokenHistorySummary, days: u32) { use crate::spend_contract::LocalHistoryCoverage; match history.coverage { LocalHistoryCoverage::Complete if history.total_tokens == 0 => { @@ -386,13 +386,17 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - if let Some(cost) = history.cost_estimate.total_usd() { + if let Some(cost) = history.total_usd() { println!(" API list-price estimate: ${cost:.2} (not billed spend)"); } else if let Some(cost) = history.cost_estimate.known_subtotal_usd { - println!( - " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", - history.cost_estimate.coverage.unpriced - ); + if history.coverage == LocalHistoryCoverage::Complete { + println!( + " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", + history.cost_estimate.coverage.unpriced + ); + } else { + println!(" Known API list-price subtotal: ${cost:.2} (history incomplete)"); + } } else { println!(" Local token history; dollar costs unavailable"); } @@ -468,7 +472,7 @@ fn build_json_payloads(results: &[CostResult], days: u32) -> Vec) -> String { let history = crate::providers::antigravity::local_sessions::summarize(30); results.push(crate::spend_contract::local_token_history_json( "antigravity", - history, + &history, 30, )); continue; } if provider_id == ProviderId::Muse { let report = crate::providers::muse::local_usage::scan(30, None); + let history = report.into(); results.push(crate::spend_contract::local_token_history_json( - "muse", - report.into(), - 30, + "muse", &history, 30, )); continue; } @@ -160,7 +159,7 @@ mod tests { use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; let complete = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, @@ -174,7 +173,7 @@ mod tests { let partial = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 3ab0e66e9c..3191e88810 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -390,7 +390,7 @@ mod tests { assert_eq!(summary.cost_estimate.coverage.estimated, 1); assert_eq!(summary.cost_estimate.coverage.unpriced, 1); assert!(summary.cost_estimate.known_subtotal_usd.is_some()); - assert_eq!(summary.cost_estimate.total_usd(), None); + assert_eq!(summary.total_usd(), None); } use rusqlite::Connection; diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index 54c7d3490a..722acec630 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -87,6 +87,21 @@ pub struct LocalTokenHistorySummary { pub cost_estimate: LocalCostEstimate, } +impl LocalTokenHistorySummary { + /// Return a complete list-price total only when both the history scan and + /// pricing coverage are complete. A complete scan with no token usage is + /// a known zero even though there were no requests to price. + pub fn total_usd(&self) -> Option { + if self.coverage != LocalHistoryCoverage::Complete { + return None; + } + if self.total_tokens == 0 { + return Some(0.0); + } + self.cost_estimate.complete_total_usd() + } +} + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct LocalCostEstimate { @@ -97,7 +112,7 @@ pub struct LocalCostEstimate { } impl LocalCostEstimate { - pub fn total_usd(&self) -> Option { + fn complete_total_usd(&self) -> Option { if self.coverage.unpriced == 0 && self.coverage.unmetered == 0 { self.known_subtotal_usd } else { @@ -122,12 +137,21 @@ impl LocalCostEstimate { pub fn local_token_history_json( provider: &str, - history: LocalTokenHistorySummary, + history: &LocalTokenHistorySummary, days: u32, ) -> serde_json::Value { let complete = history.coverage == LocalHistoryCoverage::Complete; - let total_usd = history.cost_estimate.total_usd(); + let total_usd = history.total_usd(); let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; + let note = if total_usd.is_some() { + "Local token history estimated at public API list prices; not billed spend" + } else if known_subtotal_usd.is_some() && !complete { + "Known public API list-price subtotal; local history is incomplete" + } else if known_subtotal_usd.is_some() { + "Known public API list-price subtotal; some local requests are unpriced" + } else { + "Local token history; dollar costs unavailable" + }; serde_json::json!({ "provider": provider, "supported": true, @@ -135,8 +159,8 @@ pub fn local_token_history_json( "cost": { "total_usd": total_usd, "known_subtotal_usd": known_subtotal_usd, - "currency": known_subtotal_usd.map(|_| "USD"), - "pricingCoverage": history.cost_estimate.coverage, + "currency": total_usd.or(known_subtotal_usd).map(|_| "USD"), + "pricingCoverage": &history.cost_estimate.coverage, }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, @@ -147,13 +171,7 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": if total_usd.is_some() { - "Local token history estimated at public API list prices; not billed spend" - } else if known_subtotal_usd.is_some() { - "Known public API list-price subtotal; some local requests are unpriced" - } else { - "Local token history; dollar costs unavailable" - } + "note": note, }) } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] diff --git a/rust/src/spend_contract/tests.rs b/rust/src/spend_contract/tests.rs index fb74ff09de..ac32e2e5eb 100644 --- a/rust/src/spend_contract/tests.rs +++ b/rust/src/spend_contract/tests.rs @@ -1,5 +1,58 @@ use super::*; +#[test] +fn local_history_total_requires_complete_scan_and_pricing() { + let priced = LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + }; + let partial_history = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: priced.clone(), + }; + assert_eq!(partial_history.total_usd(), None); + assert_eq!(partial_history.cost_estimate.known_subtotal_usd, Some(1.25)); + + let mixed_pricing = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + }, + }; + assert_eq!(mixed_pricing.total_usd(), None); + assert_eq!(mixed_pricing.cost_estimate.known_subtotal_usd, Some(1.25)); + + let complete = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: priced, + }; + assert_eq!(complete.total_usd(), Some(1.25)); +} + +#[test] +fn complete_empty_local_history_has_a_known_zero_total() { + let history = LocalTokenHistorySummary { + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }; + + assert_eq!(history.total_usd(), Some(0.0)); +} + #[test] fn coverage_ratio_counts_estimated_as_covered() { let coverage = CostCoverageCounts { From d5128cc96df740cc6fc5236a3915d76574966590 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:12:47 +0700 Subject: [PATCH 23/62] Fix stacked tray presentation structure --- .../src-tauri/src/tray_presentation.rs | 707 +----------------- .../src-tauri/src/tray_presentation_tests.rs | 686 +++++++++++++++++ rust/src/tray/render.rs | 212 ++---- 3 files changed, 790 insertions(+), 815 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs index 8882a3e942..6292a2c1fa 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs @@ -58,12 +58,15 @@ impl<'a> TrayPresentationPlan<'a> { settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal"; let selected = pick_tray_provider(&healthy, prefer_highest); - let (primary_percent, secondary_percent, status_rows) = match settings.tray_icon_mode { + let (icon, status_rows) = match settings.tray_icon_mode { TrayIconMode::Stacked => { if let Some((top, bottom)) = pick_stacked_tray_providers(&healthy, settings) { ( - selected_tray_percents(top, settings).0, - Some(selected_tray_percents(bottom, settings).0), + TrayIconPlan::Stacked { + top_percent: selected_tray_percents(top, settings).0, + bottom_percent: selected_tray_percents(bottom, settings).0, + has_error, + }, vec![ TrayStatusRow { key: TrayStatusKey::Provider, @@ -87,7 +90,12 @@ impl<'a> TrayPresentationPlan<'a> { }) .into_iter() .collect(); - (percents.0, percents.1, rows) + ( + resolve_single_provider_icon_plan( + settings, percents.0, percents.1, has_error, + ), + rows, + ) } } TrayIconMode::PerProvider => { @@ -102,7 +110,10 @@ impl<'a> TrayPresentationPlan<'a> { snapshot, }) .collect(); - (percents.0, percents.1, rows) + ( + resolve_single_provider_icon_plan(settings, percents.0, percents.1, has_error), + rows, + ) } TrayIconMode::Single => { let percents = selected @@ -115,12 +126,13 @@ impl<'a> TrayPresentationPlan<'a> { }) .into_iter() .collect(); - (percents.0, percents.1, rows) + ( + resolve_single_provider_icon_plan(settings, percents.0, percents.1, has_error), + rows, + ) } }; - let icon = resolve_icon_plan(settings, primary_percent, secondary_percent, has_error); - Self { settings, icon, @@ -161,21 +173,13 @@ impl<'a> TrayPresentationPlan<'a> { } } -fn resolve_icon_plan( +fn resolve_single_provider_icon_plan( settings: &Settings, primary_percent: f64, secondary_percent: Option, has_error: bool, ) -> TrayIconPlan { - if settings.tray_icon_mode == TrayIconMode::Stacked - && let Some(bottom_percent) = secondary_percent - { - TrayIconPlan::Stacked { - top_percent: primary_percent, - bottom_percent, - has_error, - } - } else if settings.menu_bar_shows_percent { + if settings.menu_bar_shows_percent { TrayIconPlan::Percent { percent: primary_percent, has_error, @@ -362,668 +366,5 @@ fn display_metric_percent(window: &RateWindowSnapshot, show_as_used: bool) -> f6 } #[cfg(test)] -mod tests { - use super::*; - use codexbar::core::{ProviderId, ProviderStateKind}; - - fn fake_snapshot(id: &str, display_name: &str, used_percent: f64) -> ProviderUsageSnapshot { - fake_snapshot_with(id, display_name, used_percent, None, None, None) - } - - fn fake_snapshot_with( - id: &str, - display_name: &str, - used_percent: f64, - secondary_percent: Option, - tertiary_percent: Option, - cost: Option<(f64, f64)>, - ) -> ProviderUsageSnapshot { - let window = |percent: f64| RateWindowSnapshot { - used_percent: percent, - remaining_percent: 100.0 - percent, - window_minutes: None, - resets_at: None, - reset_description: None, - is_exhausted: false, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }; - - ProviderUsageSnapshot { - provider_id: id.into(), - display_name: display_name.into(), - primary: window(used_percent), - primary_label: None, - secondary: secondary_percent.map(window), - secondary_label: None, - model_specific: None, - tertiary: tertiary_percent.map(window), - tertiary_label: None, - extra_rate_windows: Vec::new(), - inventory: Vec::new(), - display_details: Vec::new(), - cost: cost.map(|(used, limit)| crate::commands::CostSnapshotBridge { - used, - limit: Some(limit), - remaining: Some((limit - used).max(0.0)), - currency_code: "USD".to_string(), - currency_symbol: None, - period: "monthly".to_string(), - resets_at: None, - formatted_used: format!("${used:.2}"), - formatted_limit: Some(format!("${limit:.2}")), - balance: None, - formatted_balance: None, - balance_updated_at: None, - account_id: None, - daily: Vec::new(), - always_visible: false, - }), - plan_name: None, - account_email: None, - subscription: None, - source_label: String::new(), - has_successful_claude_cli_quota: false, - updated_at: "2025-01-01T00:00:00Z".into(), - error: None, - error_state: ProviderStateKind::Ready, - pace: None, - account_organization: None, - tray_status_label: None, - fetch_duration_ms: None, - wayfinder_usage: None, - session_equivalent_forecast: None, - } - } - - #[test] - fn single_plan_uses_highest_provider_for_icon_and_summary() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Single, - menu_bar_shows_highest_usage: true, - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Bars { - primary_percent: 72.0, - secondary_percent: None, - has_error: false, - } - ); - assert_eq!( - plan.status_labels(Language::English), - vec![("status_summary".to_string(), "Claude 72%".to_string())] - ); - } - - #[test] - fn single_plan_borrows_selected_snapshot_from_stable_input() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Single, - menu_bar_shows_highest_usage: true, - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - // `resolve` drops its temporary ordered/healthy vectors before returning. - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); - } - - #[test] - fn per_provider_plan_preserves_configured_order_for_status_rows() { - let settings = Settings { - tray_icon_mode: TrayIconMode::PerProvider, - provider_order: codexbar::settings::normalize_provider_order(&[ - "claude".to_string(), - "codex".to_string(), - ]), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = - TrayPresentationPlan::resolve(&settings, &snapshots).status_labels(Language::English); - - assert_eq!( - labels, - vec![ - ("claude".to_string(), "Claude 72%".to_string()), - ("codex".to_string(), "Codex 30%".to_string()), - ] - ); - } - - #[test] - fn stacked_plan_resolves_distinct_preferences_once() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("claude".to_string()), - stacked_tray_bottom_provider: Some("codex".to_string()), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - fake_snapshot("gemini", "Gemini", 44.0), - ]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Stacked { - top_percent: 72.0, - bottom_percent: 30.0, - has_error: false, - } - ); - assert_eq!( - plan.status_labels(Language::English), - vec![ - ("claude".to_string(), "Claude 72%".to_string()), - ("codex".to_string(), "Codex 30%".to_string()), - ] - ); - } - - #[test] - fn stacked_plan_borrows_both_snapshots_from_stable_input() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("claude".to_string()), - stacked_tray_bottom_provider: Some("codex".to_string()), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - // The plan retains references to the caller-owned snapshots, not the - // temporary vector of references used during selection. - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); - assert!(std::ptr::eq(plan.status_rows[1].snapshot, &snapshots[0])); - } - - #[test] - fn stacked_plan_falls_back_around_stale_and_duplicate_preferences() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("missing".to_string()), - stacked_tray_bottom_provider: Some("claude".to_string()), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Stacked { - top_percent: 30.0, - bottom_percent: 72.0, - has_error: false, - } - ); - assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); - assert_eq!(plan.status_rows[1].snapshot.provider_id, "claude"); - } - - #[test] - fn one_provider_stacked_plan_preserves_secondary_window_fallback() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - ..Settings::default() - }; - let snapshots = vec![fake_snapshot_with( - "codex", - "Codex", - 30.0, - Some(65.0), - None, - None, - )]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Stacked { - top_percent: 65.0, - bottom_percent: 30.0, - has_error: false, - } - ); - assert_eq!(plan.status_rows.len(), 1); - } - - #[test] - fn all_errors_produce_error_styled_zero_percent_plan() { - let settings = Settings { - menu_bar_shows_percent: true, - ..Settings::default() - }; - let mut snapshot = fake_snapshot("codex", "Codex", 30.0); - snapshot.error = Some("offline".to_string()); - let snapshots = vec![snapshot]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Percent { - percent: 0.0, - has_error: true, - } - ); - assert!(plan.status_rows.is_empty()); - } - - #[test] - fn plan_uses_selected_metric_and_remaining_display_mode() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshots = vec![fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - )]; - - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.icon, - TrayIconPlan::Bars { - primary_percent: 85.0, - secondary_percent: Some(80.0), - has_error: false, - } - ); - } - - #[test] - fn render_icon_delegates_to_resolved_stacked_renderer() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Stacked, - stacked_tray_top_provider: Some("claude".to_string()), - stacked_tray_bottom_provider: Some("codex".to_string()), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 40.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - let plan = TrayPresentationPlan::resolve(&settings, &snapshots); - - assert_eq!( - plan.render_icon(), - render_stacked_bar_icon_rgba(72.0, 40.0, false) - ); - } - - #[test] - fn codex_headline_skips_informational_primary() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); - snapshot.primary.is_informational = true; - - assert_eq!(codex_lane_headline_window(&snapshot).used_percent, 25.0); - } - fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { - crate::commands::NamedRateWindowSnapshot { - id: "additional_budget".to_string(), - title: "Additional Budget".to_string(), - fallback_lane: false, - window: crate::commands::RateWindowSnapshot { - used_percent: percent, - remaining_percent: 100.0 - percent, - window_minutes: None, - resets_at: None, - reset_description: None, - is_exhausted: false, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }, - } - } - - #[test] - fn selected_tray_percent_uses_cursor_extra_usage_cost() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(20.0)); - } - - #[test] - fn selected_tray_percent_tracks_extra_rate_window() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn copilot_automatic_tracks_highest_extra_rate_window() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - } - - #[test] - fn selected_tray_percent_respects_remaining_display_mode() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 85.0); - assert_eq!(secondary, Some(80.0)); - } - - #[test] - fn exhausted_automatic_window_never_renders_as_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(40.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(100.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = false; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn missing_automatic_window_does_not_look_like_available_remaining_progress() { - let settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); - snapshot.primary.is_informational = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(remaining, 0.0); - } - - #[test] - fn selected_tray_percent_falls_back_when_extra_usage_missing() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 72.0); - } - - #[test] - fn single_meaningful_secondary_quota_uses_full_single_meter() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { - let settings = Settings::default(); - let mut snapshot = - fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, Some(30.0)); - } - - #[test] - fn two_meaningful_quotas_keep_two_meter_layout() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); - let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(40.0)); - } - - #[test] - fn informational_primary_skips_session_and_automatic_phantom_zero() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - // Session preference must not paint the synthetic 0% primary; - // it falls through to Automatic which prefers weekly (42%). - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - assert_ne!(primary, 0.0); - - // Automatic also prefers weekly over informational primary. - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - } - - #[test] - fn claude_automatic_prefers_weekly_when_model_exhausted() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot_with("claude", "Claude", 40.0, Some(22.0), None, None); - snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { - used_percent: 100.0, - remaining_percent: 0.0, - window_minutes: Some(10080), - resets_at: None, - reset_description: None, - is_exhausted: true, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 22.0); - - // Explicit model override is untouched. - let mut overridden = settings.clone(); - overridden.set_provider_metric(ProviderId::Claude, MetricPreference::Model); - let (primary, _) = selected_tray_percents(&snapshot, &overridden); - assert_eq!(primary, 100.0); - } - - #[test] - fn automatic_prefers_exhausted_weekly_over_low_session() { - let settings = Settings::default(); - let snapshot = fake_snapshot_with("codex", "Codex", 20.0, Some(100.0), None, None); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 100.0); - - // Explicit session override still wins. - let mut overridden = settings.clone(); - overridden.set_provider_metric(ProviderId::Codex, MetricPreference::Session); - let (primary, _) = selected_tray_percents(&snapshot, &overridden); - assert_eq!(primary, 20.0); - } - - #[test] - fn automatic_picks_highest_among_model_and_extra_windows() { - let settings = Settings::default(); - let mut snapshot = - fake_snapshot_with("gemini", "Gemini", 10.0, Some(30.0), Some(40.0), None); - snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { - used_percent: 55.0, - remaining_percent: 45.0, - window_minutes: None, - resets_at: None, - reset_description: None, - is_exhausted: false, - is_informational: false, - reserve_percent: None, - reserve_description: None, - reserve_will_last_to_reset: false, - reserve_eta_seconds: None, - }); - snapshot.extra_rate_windows.push(fake_extra_window(90.0)); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 90.0); - } - - #[test] - fn f5_headline_prefers_non_informational_primary() { - let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_secondary_when_primary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); - snapshot.primary.is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); - snapshot.primary.is_informational = true; - snapshot.secondary.as_mut().unwrap().is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_returns_primary_when_all_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); - snapshot.primary.is_informational = true; - if let Some(sec) = &mut snapshot.secondary { - sec.is_informational = true; - } - if let Some(ter) = &mut snapshot.tertiary { - ter.is_informational = true; - } - let headline = codex_lane_headline_window(&snapshot); - // Falls back to primary (the placeholder) when all are informational. - assert!(headline.is_informational); - } -} +#[path = "tray_presentation_tests.rs"] +mod tests; diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs new file mode 100644 index 0000000000..3fce3f03d7 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs @@ -0,0 +1,686 @@ +use super::*; +use codexbar::core::{ProviderId, ProviderStateKind}; + +fn fake_snapshot(id: &str, display_name: &str, used_percent: f64) -> ProviderUsageSnapshot { + fake_snapshot_with(id, display_name, used_percent, None, None, None) +} + +fn fake_snapshot_with( + id: &str, + display_name: &str, + used_percent: f64, + secondary_percent: Option, + tertiary_percent: Option, + cost: Option<(f64, f64)>, +) -> ProviderUsageSnapshot { + let window = |percent: f64| RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }; + + ProviderUsageSnapshot { + provider_id: id.into(), + display_name: display_name.into(), + primary: window(used_percent), + primary_label: None, + secondary: secondary_percent.map(window), + secondary_label: None, + model_specific: None, + tertiary: tertiary_percent.map(window), + tertiary_label: None, + extra_rate_windows: Vec::new(), + inventory: Vec::new(), + display_details: Vec::new(), + cost: cost.map(|(used, limit)| crate::commands::CostSnapshotBridge { + used, + limit: Some(limit), + remaining: Some((limit - used).max(0.0)), + currency_code: "USD".to_string(), + currency_symbol: None, + period: "monthly".to_string(), + resets_at: None, + formatted_used: format!("${used:.2}"), + formatted_limit: Some(format!("${limit:.2}")), + balance: None, + formatted_balance: None, + balance_updated_at: None, + account_id: None, + daily: Vec::new(), + always_visible: false, + }), + plan_name: None, + account_email: None, + subscription: None, + source_label: String::new(), + has_successful_claude_cli_quota: false, + updated_at: "2025-01-01T00:00:00Z".into(), + error: None, + error_state: ProviderStateKind::Ready, + pace: None, + account_organization: None, + tray_status_label: None, + fetch_duration_ms: None, + wayfinder_usage: None, + session_equivalent_forecast: None, + } +} + +#[test] +fn single_plan_uses_highest_provider_for_icon_and_summary() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 72.0, + secondary_percent: None, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![("status_summary".to_string(), "Claude 72%".to_string())] + ); +} + +#[test] +fn single_plan_borrows_selected_snapshot_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // `resolve` drops its temporary ordered/healthy vectors before returning. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); +} + +#[test] +fn per_provider_plan_preserves_configured_order_for_status_rows() { + let settings = Settings { + tray_icon_mode: TrayIconMode::PerProvider, + provider_order: codexbar::settings::normalize_provider_order(&[ + "claude".to_string(), + "codex".to_string(), + ]), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let labels = + TrayPresentationPlan::resolve(&settings, &snapshots).status_labels(Language::English); + + assert_eq!( + labels, + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); +} + +#[test] +fn stacked_plan_resolves_distinct_preferences_once() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + fake_snapshot("gemini", "Gemini", 44.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 72.0, + bottom_percent: 30.0, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); +} + +#[test] +fn stacked_plan_borrows_both_snapshots_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // The plan retains references to the caller-owned snapshots, not the + // temporary vector of references used during selection. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); + assert!(std::ptr::eq(plan.status_rows[1].snapshot, &snapshots[0])); +} + +#[test] +fn stacked_plan_falls_back_around_stale_and_duplicate_preferences() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("missing".to_string()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 30.0, + bottom_percent: 72.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); + assert_eq!(plan.status_rows[1].snapshot.provider_id, "claude"); +} + +#[test] +fn one_provider_stacked_mode_falls_back_to_single_provider_bars() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + ..Settings::default() + }; + let snapshots = vec![fake_snapshot_with( + "codex", + "Codex", + 30.0, + Some(65.0), + None, + None, + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 65.0, + secondary_percent: Some(30.0), + has_error: false, + } + ); + assert_eq!(plan.status_rows.len(), 1); +} + +#[test] +fn one_healthy_provider_never_uses_stacked_renderer() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + menu_bar_shows_percent: true, + ..Settings::default() + }; + let healthy = fake_snapshot("codex", "Codex", 30.0); + let mut failed = fake_snapshot("claude", "Claude", 72.0); + failed.error = Some("offline".to_string()); + let snapshots = vec![healthy, failed]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Percent { + percent: 30.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows.len(), 1); + assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); +} + +#[test] +fn all_errors_produce_error_styled_zero_percent_plan() { + let settings = Settings { + menu_bar_shows_percent: true, + ..Settings::default() + }; + let mut snapshot = fake_snapshot("codex", "Codex", 30.0); + snapshot.error = Some("offline".to_string()); + let snapshots = vec![snapshot]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Percent { + percent: 0.0, + has_error: true, + } + ); + assert!(plan.status_rows.is_empty()); +} + +#[test] +fn plan_uses_selected_metric_and_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshots = vec![fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 85.0, + secondary_percent: Some(80.0), + has_error: false, + } + ); +} + +#[test] +fn render_icon_delegates_to_resolved_stacked_renderer() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 40.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.render_icon(), + render_stacked_bar_icon_rgba(72.0, 40.0, false) + ); +} + +#[test] +fn codex_headline_skips_informational_primary() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + assert_eq!(codex_lane_headline_window(&snapshot).used_percent, 25.0); +} +fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { + crate::commands::NamedRateWindowSnapshot { + id: "additional_budget".to_string(), + title: "Additional Budget".to_string(), + fallback_lane: false, + window: crate::commands::RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }, + } +} + +#[test] +fn selected_tray_percent_uses_cursor_extra_usage_cost() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(20.0)); +} + +#[test] +fn selected_tray_percent_tracks_extra_rate_window() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); +} + +#[test] +fn copilot_automatic_tracks_highest_extra_rate_window() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); +} + +#[test] +fn selected_tray_percent_respects_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 85.0); + assert_eq!(secondary, Some(80.0)); +} + +#[test] +fn exhausted_automatic_window_never_renders_as_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(40.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); +} + +#[test] +fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(100.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = false; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); +} + +#[test] +fn missing_automatic_window_does_not_look_like_available_remaining_progress() { + let settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); + snapshot.primary.is_informational = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(remaining, 0.0); +} + +#[test] +fn selected_tray_percent_falls_back_when_extra_usage_missing() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 72.0); +} + +#[test] +fn single_meaningful_secondary_quota_uses_full_single_meter() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); +} + +#[test] +fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, Some(30.0)); +} + +#[test] +fn two_meaningful_quotas_keep_two_meter_layout() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); + let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(40.0)); +} + +#[test] +fn informational_primary_skips_session_and_automatic_phantom_zero() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + // Session preference must not paint the synthetic 0% primary; + // it falls through to Automatic which prefers weekly (42%). + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); + assert_ne!(primary, 0.0); + + // Automatic also prefers weekly over informational primary. + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); +} + +#[test] +fn claude_automatic_prefers_weekly_when_model_exhausted() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 40.0, Some(22.0), None, None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 100.0, + remaining_percent: 0.0, + window_minutes: Some(10080), + resets_at: None, + reset_description: None, + is_exhausted: true, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 22.0); + + // Explicit model override is untouched. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Claude, MetricPreference::Model); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 100.0); +} + +#[test] +fn automatic_prefers_exhausted_weekly_over_low_session() { + let settings = Settings::default(); + let snapshot = fake_snapshot_with("codex", "Codex", 20.0, Some(100.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 100.0); + + // Explicit session override still wins. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Codex, MetricPreference::Session); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 20.0); +} + +#[test] +fn automatic_picks_highest_among_model_and_extra_windows() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("gemini", "Gemini", 10.0, Some(30.0), Some(40.0), None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 55.0, + remaining_percent: 45.0, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + snapshot.extra_rate_windows.push(fake_extra_window(90.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 90.0); +} + +#[test] +fn f5_headline_prefers_non_informational_primary() { + let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_falls_back_to_secondary_when_primary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); + snapshot.primary.is_informational = true; + snapshot.secondary.as_mut().unwrap().is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_returns_primary_when_all_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); + snapshot.primary.is_informational = true; + if let Some(sec) = &mut snapshot.secondary { + sec.is_informational = true; + } + if let Some(ter) = &mut snapshot.tertiary { + ter.is_informational = true; + } + let headline = codex_lane_headline_window(&snapshot); + // Falls back to primary (the placeholder) when all are informational. + assert!(headline.is_informational); +} diff --git a/rust/src/tray/render.rs b/rust/src/tray/render.rs index 1787f60f3c..02c4ba668e 100644 --- a/rust/src/tray/render.rs +++ b/rust/src/tray/render.rs @@ -10,6 +10,62 @@ use super::icon::UsageLevel; /// Side length of the generated tray icon in pixels. pub const TRAY_ICON_SIZE: u32 = 32; +const ICON_INSET: u32 = 2; +const BAR_LEFT: u32 = 4; +const BAR_RIGHT: u32 = TRAY_ICON_SIZE - 4; +const ICON_BACKGROUND_RGB: [u8; 3] = [60, 60, 70]; +const BAR_BACKGROUND: Rgba = Rgba([80, 80, 90, 255]); + +fn new_icon_canvas(has_error: bool) -> RgbaImage { + let mut image: RgbaImage = ImageBuffer::new(TRAY_ICON_SIZE, TRAY_ICON_SIZE); + let background = Rgba([ + ICON_BACKGROUND_RGB[0], + ICON_BACKGROUND_RGB[1], + ICON_BACKGROUND_RGB[2], + if has_error { 180 } else { 255 }, + ]); + for y in ICON_INSET..TRAY_ICON_SIZE - ICON_INSET { + for x in ICON_INSET..TRAY_ICON_SIZE - ICON_INSET { + image.put_pixel(x, y, background); + } + } + image +} + +fn usage_color(percent: f64, has_error: bool) -> Rgba { + let (r, g, b) = UsageLevel::from_percent(percent).color(); + if has_error { + #[allow( + clippy::cast_possible_truncation, + reason = "mean of three u8 channels is bounded to 0..=255" + )] + let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; + Rgba([gray, gray, gray, 255]) + } else { + Rgba([r, g, b, 255]) + } +} + +fn draw_bar_row(image: &mut RgbaImage, y_start: u32, y_end: u32, percent: f64, has_error: bool) { + let bar_width = BAR_RIGHT - BAR_LEFT; + #[allow( + clippy::cast_possible_truncation, + reason = "percent is clamped to 0..=100 and scaled to a 24-pixel meter" + )] + let fill = ((percent.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; + let fill_end = (BAR_LEFT + fill).min(BAR_RIGHT); + let color = usage_color(percent, has_error); + + for y in y_start..y_end { + for x in BAR_LEFT..BAR_RIGHT { + image.put_pixel(x, y, BAR_BACKGROUND); + } + for x in BAR_LEFT..fill_end { + image.put_pixel(x, y, color); + } + } +} + /// Render a usage-bar tray icon as raw RGBA bytes. /// /// - `session_percent`: primary bar fill (0–100), colour-coded by [`UsageLevel`] @@ -24,73 +80,19 @@ pub fn render_bar_icon_rgba( weekly_percent: Option, has_error: bool, ) -> (Vec, u32, u32) { - const SZ: u32 = TRAY_ICON_SIZE; - let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); - - for pixel in img.pixels_mut() { - *pixel = Rgba([0, 0, 0, 0]); - } - - let bg_alpha: u8 = if has_error { 180 } else { 255 }; - let bg_color = Rgba([60, 60, 70, bg_alpha]); - for y in 2..SZ - 2 { - for x in 2..SZ - 2 { - img.put_pixel(x, y, bg_color); - } - } - - let color_for = |percent: f64| -> (u8, u8, u8) { - let (r, g, b) = UsageLevel::from_percent(percent).color(); - if has_error { - // Average of three u8 colour channels: sum ≤ 765, so /3 ≤ 255 fits u8. - #[allow( - clippy::cast_possible_truncation, - reason = "mean of three u8 channels; r+g+b ≤ 765, divided by 3 is ≤ 255 and fits u8" - )] - let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; - (gray, gray, gray) - } else { - (r, g, b) - } - }; - - let bar_left = 4u32; - let bar_right = SZ - 4; - let bar_width = bar_right - bar_left; - - // pct is clamped to 0–100, scaled by bar_width (≤ SZ = 32), so the result fits u32. - #[allow( - clippy::cast_possible_truncation, - reason = "pct clamped to 0–100 and scaled by bar_width ≤ 32; result is a small pixel count that fits u32" - )] - let fill_px = |pct: f64| ((pct.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; - - let mut draw_bar = |y_start: u32, y_end: u32, pct: f64| { - let (r, g, b) = color_for(pct); - let fill_end = (bar_left + fill_px(pct)).min(bar_right); - for y in y_start..y_end { - for x in bar_left..bar_right { - img.put_pixel(x, y, Rgba([80, 80, 90, 255])); - } - } - for y in y_start..y_end { - for x in bar_left..fill_end { - img.put_pixel(x, y, Rgba([r, g, b, 255])); - } - } - }; + let mut image = new_icon_canvas(has_error); match weekly_percent { Some(weekly) => { - draw_bar(8, 15, session_percent); // session bar (top, thicker) - draw_bar(18, 23, weekly); // weekly bar (bottom, thinner) + draw_bar_row(&mut image, 8, 15, session_percent, has_error); + draw_bar_row(&mut image, 18, 23, weekly, has_error); } None => { - draw_bar(10, 22, session_percent); // single thick bar (centred) + draw_bar_row(&mut image, 10, 22, session_percent, has_error); } } - (img.into_raw(), SZ, SZ) + (image.into_raw(), TRAY_ICON_SIZE, TRAY_ICON_SIZE) } /// Render two providers as equally prominent stacked usage meters. @@ -103,72 +105,16 @@ pub fn render_stacked_bar_icon_rgba( bottom_percent: f64, has_error: bool, ) -> (Vec, u32, u32) { - const SZ: u32 = TRAY_ICON_SIZE; - let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); - - for pixel in img.pixels_mut() { - *pixel = Rgba([0, 0, 0, 0]); - } - - let bg_alpha = if has_error { 180 } else { 255 }; - for y in 2..SZ - 2 { - for x in 2..SZ - 2 { - img.put_pixel(x, y, Rgba([60, 60, 70, bg_alpha])); - } - } - - let bar_left = 4u32; - let bar_right = SZ - 4; - let bar_width = bar_right - bar_left; - let mut draw_provider = |y_start: u32, y_end: u32, percent: f64| { - let (r, g, b) = UsageLevel::from_percent(percent).color(); - let color = if has_error { - #[allow( - clippy::cast_possible_truncation, - reason = "mean of three u8 channels is bounded to 0..=255" - )] - let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; - Rgba([gray, gray, gray, 255]) - } else { - Rgba([r, g, b, 255]) - }; - #[allow( - clippy::cast_possible_truncation, - reason = "percent is clamped to 0..=100 and scaled to a 24-pixel meter" - )] - let fill = ((percent.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; - let fill_end = (bar_left + fill).min(bar_right); - - for y in y_start..y_end { - for x in bar_left..bar_right { - img.put_pixel(x, y, Rgba([80, 80, 90, 255])); - } - for x in bar_left..fill_end { - img.put_pixel(x, y, color); - } - } - }; - - draw_provider(6, 14, top_percent); - draw_provider(18, 26, bottom_percent); - (img.into_raw(), SZ, SZ) + let mut image = new_icon_canvas(has_error); + draw_bar_row(&mut image, 6, 14, top_percent, has_error); + draw_bar_row(&mut image, 18, 26, bottom_percent, has_error); + (image.into_raw(), TRAY_ICON_SIZE, TRAY_ICON_SIZE) } /// Render a compact numeric percent tray icon as raw RGBA bytes. pub fn render_percent_icon_rgba(percent: f64, has_error: bool) -> (Vec, u32, u32) { const SZ: u32 = TRAY_ICON_SIZE; - let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); - - for pixel in img.pixels_mut() { - *pixel = Rgba([0, 0, 0, 0]); - } - - let bg_alpha: u8 = if has_error { 180 } else { 255 }; - for y in 2..SZ - 2 { - for x in 2..SZ - 2 { - img.put_pixel(x, y, Rgba([60, 60, 70, bg_alpha])); - } - } + let mut img = new_icon_canvas(has_error); // percent clamped to 0–100 before rounding, so the cast to u32 cannot truncate. #[allow( @@ -195,18 +141,7 @@ pub fn render_percent_icon_rgba(percent: f64, has_error: bool) -> (Vec, u32, let start_x = (SZ.saturating_sub(text_width)) / 2; let start_y = (SZ.saturating_sub(text_height)) / 2; - let (r, g, b) = UsageLevel::from_percent(percent).color(); - let color = if has_error { - // Average of three u8 colour channels: sum ≤ 765, so /3 ≤ 255 fits u8. - #[allow( - clippy::cast_possible_truncation, - reason = "mean of three u8 channels; r+g+b ≤ 765, divided by 3 is ≤ 255 and fits u8" - )] - let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; - Rgba([gray, gray, gray, 255]) - } else { - Rgba([r, g, b, 255]) - }; + let color = usage_color(percent, has_error); let mut x = start_x; for ch in text.chars() { @@ -385,4 +320,17 @@ mod tests { assert_eq!(pixel(8, 20), [80, 80, 90, 255]); assert_eq!(pixel(8, 15), [60, 60, 70, 255]); } + + #[test] + fn normal_and_stacked_bars_share_error_color_policy() { + let (normal, width, _) = render_bar_icon_rgba(100.0, None, true); + let (stacked, _, _) = render_stacked_bar_icon_rgba(100.0, 0.0, true); + let pixel = |rgba: &[u8], x: u32, y: u32| { + let index = ((y * width + x) * 4) as usize; + &rgba[index..index + 4] + }; + + assert_eq!(pixel(&normal, 8, 12), pixel(&stacked, 8, 8)); + assert_eq!(pixel(&normal, 8, 12)[0], pixel(&normal, 8, 12)[1]); + } } From 58e984b2067baf7d8043e8823641f7199f6d487f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:17:54 +0700 Subject: [PATCH 24/62] Fix partial overview spend expectation --- apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index ecefece9f2..e36f1e8f1e 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -358,7 +358,7 @@ describe("TrayPanel provider grid", () => { renderTrayPanel([provider("codex", "Codex", 35)]); expect(await screen.findByRole("button", { name: "UsageSpendShare" })).toBeInTheDocument(); - expect(screen.getByText("$2.00")).toBeInTheDocument(); + expect(screen.getByText("~$2.00")).toBeInTheDocument(); expect(screen.getByText(/1 of 2 OverviewSpendProviderCoverage/)).toBeInTheDocument(); }); From 1badfe80a251a9b325275591ab44047a686d962c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:32:45 +0700 Subject: [PATCH 25/62] Fix Codex fork accounting precedence --- rust/src/core/jsonl_scanner/codex/parser.rs | 10 +- rust/src/core/jsonl_scanner/tests.rs | 29 ++++ rust/src/cost_scanner/codex.rs | 140 ++++++++++++++------ rust/src/cost_scanner/tests/paginated.rs | 46 ++++++- 4 files changed, 181 insertions(+), 44 deletions(-) diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index c3aa7cb332..f33316e378 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -108,8 +108,14 @@ impl ForkBaselineInference { let last = read_token_totals(last_usage); let ordinal = obj.get("ordinal").and_then(Value::as_i64); - if let Some(start) = self.explicit_start_ordinal { - if ordinal.is_some_and(|ordinal| ordinal < start) { + if let Some(start) = self.explicit_start_ordinal + && !self.boundary_open + { + let Some(ordinal) = ordinal else { + self.baseline = Some(total); + return ForkBaselineDecision::SkipCopiedPrefix; + }; + if ordinal < start { self.baseline = Some(total); return ForkBaselineDecision::SkipCopiedPrefix; } diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index 02d9257753..9296583fb8 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -115,6 +115,35 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { assert!(!state.fork_baseline_ambiguous); } +#[test] +fn inferred_fork_waits_for_present_explicit_start_ordinal() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + + state.process_line( + r#"{"timestamp":"2026-09-22T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.6-sol","total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":10},"last_token_usage":{"input_tokens":0,"cached_input_tokens":0,"output_tokens":0}}}}"#, + &range, + ); + + assert!(state.records.is_empty()); + assert!(state.fork_baseline.is_none()); + + state.process_line( + r#"{"ordinal":10,"timestamp":"2026-09-22T10:00:01Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.6-sol","total_token_usage":{"input_tokens":110,"cached_input_tokens":22,"output_tokens":11},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":1}}}}"#, + &range, + ); + + assert_eq!(state.records.len(), 1); + assert_eq!(state.records[0].input, 10); + assert_eq!(state.records[0].cached, 2); + assert_eq!(state.records[0].output, 1); +} + #[test] fn codex_token_pipeline_preserves_counts_above_i32_max() { let parsed = read_token_totals(&serde_json::json!({ diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 5b181311f9..56a0f9724a 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -13,6 +13,41 @@ use pending_range::{ }; use reconciliation::*; +#[derive(Debug)] +enum CodexAccountingMode { + Standard, + ValidatedBaseline { + baseline: crate::core::CodexTotals, + paginated_continuation: bool, + remaining_inherited_totals: Option, + locally_resolved: bool, + }, + InferSubagent { + start_ordinal: Option, + }, + Unresolved, +} + +impl CodexAccountingMode { + fn is_unresolved(&self) -> bool { + matches!(self, Self::Unresolved) + } + + fn infers_subagent_baseline(&self) -> bool { + matches!(self, Self::InferSubagent { .. }) + } + + fn locally_resolved(&self) -> bool { + matches!( + self, + Self::ValidatedBaseline { + locally_resolved: true, + .. + } + ) + } +} + fn summary_from_cached_report( report: &CachedCostReport, period_start: NaiveDate, @@ -463,7 +498,6 @@ impl CostScanner { })) }); let is_fork = codex_lineage.uses_parent_baseline(); - let locally_inferred_subagent = is_fork && session_metadata.is_subagent; let cached_fork_state_matches = cached_fork_accounting_state.as_ref().is_some_and(|state| { state.session_id == codex_session_id @@ -471,29 +505,44 @@ impl CostScanner { && state.history_base_thread_id == history_base_thread_id && state.fork_timestamp == codex_fork_timestamp }); - let fork_baseline = cached_fork_accounting_state + let matching_cached_fork_state = cached_fork_accounting_state .as_ref() - .filter(|_| cached_fork_state_matches) - .and_then(|state| state.inherited_totals.clone()) - .or_else(|| { - is_fork - .then_some(codex_forked_from_id.as_deref()) - .flatten() - .and_then(|parent_id| { - codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) - }) + .filter(|_| cached_fork_state_matches); + let cached_fork_baseline = + matching_cached_fork_state.and_then(|state| state.inherited_totals.clone()); + let parent_fork_baseline = is_fork + .then_some(codex_forked_from_id.as_deref()) + .flatten() + .and_then(|parent_id| { + codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) }); - let remaining_inherited_totals = cached_fork_accounting_state - .as_ref() - .filter(|_| cached_fork_state_matches) - .and_then(|state| state.remaining_inherited_totals.clone()); + let fork_baseline = cached_fork_baseline.or(parent_fork_baseline); + let remaining_inherited_totals = + matching_cached_fork_state.and_then(|state| state.remaining_inherited_totals.clone()); let paginated_continuation = is_fork && codex_forked_from_id.is_some() && history_base_thread_id .as_deref() .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); + let accounting_mode = if !is_fork { + CodexAccountingMode::Standard + } else if let Some(baseline) = fork_baseline { + CodexAccountingMode::ValidatedBaseline { + baseline, + paginated_continuation, + remaining_inherited_totals, + locally_resolved: matching_cached_fork_state + .is_some_and(|state| state.locally_resolved), + } + } else if session_metadata.is_subagent { + CodexAccountingMode::InferSubagent { + start_ordinal: session_metadata.subagent_history_start_ordinal, + } + } else { + CodexAccountingMode::Unresolved + }; - if is_fork && fork_baseline.is_none() && !locally_inferred_subagent { + if accounting_mode.is_unresolved() { cache.files.insert( path_key, CostUsageFileUsage { @@ -634,38 +683,46 @@ impl CostScanner { let parse_target_size = cached .as_ref() .and_then(|entry| codex_resumable_scan_target_size(size, entry)); - let parse_result = match if locally_inferred_subagent { - JsonlScanner::parse_codex_file_with_inferred_fork_baseline( + let parse_result = match match &accounting_mode { + CodexAccountingMode::Standard => JsonlScanner::parse_codex_file_with_state_bounded( path, range, - session_metadata.subagent_history_start_ordinal, + 0, + None, + None, + None, + None, cancel, - parse_target_size, max_bytes_to_read, - ) - } else if let Some(baseline) = fork_baseline.clone() { - JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( - path, - range, + ), + CodexAccountingMode::ValidatedBaseline { baseline, paginated_continuation, - remaining_inherited_totals.clone(), - cancel, - parse_target_size, - max_bytes_to_read, - ) - } else { - JsonlScanner::parse_codex_file_with_state_bounded( + remaining_inherited_totals, + .. + } => JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( path, range, - 0, - None, - None, - None, - None, + baseline.clone(), + *paginated_continuation, + remaining_inherited_totals.clone(), cancel, + parse_target_size, max_bytes_to_read, - ) + ), + CodexAccountingMode::InferSubagent { start_ordinal } => { + JsonlScanner::parse_codex_file_with_inferred_fork_baseline( + path, + range, + *start_ordinal, + cancel, + parse_target_size, + max_bytes_to_read, + ) + } + CodexAccountingMode::Unresolved => { + unreachable!("unresolved forks return before parsing") + } } { Ok(result) => result, Err(_) => return CodexFileScanOutcome::default(), @@ -674,7 +731,8 @@ impl CostScanner { .token_timestamp_comparisons .saturating_add(parse_result.token_timestamp_comparisons); if parse_result.fork_baseline_ambiguous - || (locally_inferred_subagent && !parse_result.fork_baseline_locally_resolved) + || (accounting_mode.infers_subagent_baseline() + && !parse_result.fork_baseline_locally_resolved) { cache.files.insert( path_key, @@ -715,6 +773,8 @@ impl CostScanner { bytes_read: parse_result.bytes_read, is_complete: parse_result.is_complete, }; + let locally_resolved = + accounting_mode.locally_resolved() || parse_result.fork_baseline_locally_resolved; let codex_fork_accounting_state = if is_fork && (parse_result.fork_baseline.is_some() || parse_result.fork_baseline_locally_resolved) { @@ -725,7 +785,7 @@ impl CostScanner { fork_timestamp: codex_fork_timestamp.clone(), inherited_totals: parse_result.fork_baseline.clone(), remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), - locally_resolved: parse_result.fork_baseline_locally_resolved, + locally_resolved, }) } else { None diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index c13c8255b9..a225233a7d 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -138,6 +138,7 @@ fn write_codex_paginated_continuation_fixture( fn write_copied_prefix_subagent_fixture( sessions_root: &Path, name: &str, + parent_id: &str, base: DateTime, owned: bool, ) -> PathBuf { @@ -152,10 +153,10 @@ fn write_copied_prefix_subagent_fixture( serde_json::json!({ "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), "payload": { - "id": "child-id", "forked_from_id": "missing-parent", + "id": "child-id", "forked_from_id": parent_id, "subagent_history_start_ordinal": 10, "thread_source": "subagent", - "source": {"subagent": {"thread_spawn": {"parent_thread_id": "missing-parent"}}} + "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} } }), token_row(base, 2, [1_000, 900, 100], [0, 0, 0], "gpt-5.6-sol"), @@ -236,6 +237,7 @@ fn copied_prefix_subagent_infers_advancing_baseline_without_parent() { let child = write_copied_prefix_subagent_fixture( &sessions, "child.jsonl", + "missing-parent", Utc::now() - Duration::hours(1), true, ); @@ -275,6 +277,7 @@ fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { let child = write_copied_prefix_subagent_fixture( &sessions, "child.jsonl", + "missing-parent", Utc::now() - Duration::hours(1), false, ); @@ -301,6 +304,45 @@ fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { assert!(stats.codex_history_read_paths.is_empty()); } +#[test] +fn copied_prefix_subagent_prefers_validated_parent_baseline() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "parent-id", + base + Duration::seconds(10), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + + assert_eq!(state.inherited_totals.as_ref().unwrap().input, 1_000); + assert!(!state.locally_resolved); +} + #[test] fn paginated_continuation_raises_inherited_baseline_from_total_last() { let root = tempfile::tempdir().unwrap(); From 5dcb845c766f11b3dad6e4a19705242526db1587 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:41:24 +0700 Subject: [PATCH 26/62] Fix tray render test lifetime --- rust/src/tray/render.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/rust/src/tray/render.rs b/rust/src/tray/render.rs index 02c4ba668e..d0d464e7b6 100644 --- a/rust/src/tray/render.rs +++ b/rust/src/tray/render.rs @@ -327,7 +327,12 @@ mod tests { let (stacked, _, _) = render_stacked_bar_icon_rgba(100.0, 0.0, true); let pixel = |rgba: &[u8], x: u32, y: u32| { let index = ((y * width + x) * 4) as usize; - &rgba[index..index + 4] + [ + rgba[index], + rgba[index + 1], + rgba[index + 2], + rgba[index + 3], + ] }; assert_eq!(pixel(&normal, 8, 12), pixel(&stacked, 8, 8)); From d3be0188aa0c4672f9ad3115b6c516f24fa908f2 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:47:15 +0700 Subject: [PATCH 27/62] Fix Codex fork accounting test access --- rust/src/core/jsonl_scanner/tests.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index 9296583fb8..2614ebefd3 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -139,9 +139,9 @@ fn inferred_fork_waits_for_present_explicit_start_ordinal() { ); assert_eq!(state.records.len(), 1); - assert_eq!(state.records[0].input, 10); - assert_eq!(state.records[0].cached, 2); - assert_eq!(state.records[0].output, 1); + assert_eq!(state.records[0].0.input, 10); + assert_eq!(state.records[0].0.cached, 2); + assert_eq!(state.records[0].0.output, 1); } #[test] From f64861c0a4cf40919a215f6455673320e2c6f463 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 23:19:56 +0700 Subject: [PATCH 28/62] Fix Codex fork baseline provenance --- rust/src/cost_scanner/codex.rs | 76 ++++++++++++++++++------ rust/src/cost_scanner/tests/paginated.rs | 59 ++++++++++++++++++ 2 files changed, 117 insertions(+), 18 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 56a0f9724a..7a9d3068dc 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -16,11 +16,11 @@ use reconciliation::*; #[derive(Debug)] enum CodexAccountingMode { Standard, - ValidatedBaseline { + Baseline { baseline: crate::core::CodexTotals, paginated_continuation: bool, remaining_inherited_totals: Option, - locally_resolved: bool, + provenance: CodexBaselineProvenance, }, InferSubagent { start_ordinal: Option, @@ -28,6 +28,13 @@ enum CodexAccountingMode { Unresolved, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CodexBaselineProvenance { + ValidatedParent { replaces_cached_state: bool }, + CachedValidatedParent, + CachedLocalInference, +} + impl CodexAccountingMode { fn is_unresolved(&self) -> bool { matches!(self, Self::Unresolved) @@ -40,8 +47,20 @@ impl CodexAccountingMode { fn locally_resolved(&self) -> bool { matches!( self, - Self::ValidatedBaseline { - locally_resolved: true, + Self::Baseline { + provenance: CodexBaselineProvenance::CachedLocalInference, + .. + } + ) + } + + fn requires_cached_reparse(&self) -> bool { + matches!( + self, + Self::Baseline { + provenance: CodexBaselineProvenance::ValidatedParent { + replaces_cached_state: true + }, .. } ) @@ -508,17 +527,12 @@ impl CostScanner { let matching_cached_fork_state = cached_fork_accounting_state .as_ref() .filter(|_| cached_fork_state_matches); - let cached_fork_baseline = - matching_cached_fork_state.and_then(|state| state.inherited_totals.clone()); let parent_fork_baseline = is_fork .then_some(codex_forked_from_id.as_deref()) .flatten() .and_then(|parent_id| { codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) }); - let fork_baseline = cached_fork_baseline.or(parent_fork_baseline); - let remaining_inherited_totals = - matching_cached_fork_state.and_then(|state| state.remaining_inherited_totals.clone()); let paginated_continuation = is_fork && codex_forked_from_id.is_some() && history_base_thread_id @@ -526,13 +540,34 @@ impl CostScanner { .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); let accounting_mode = if !is_fork { CodexAccountingMode::Standard - } else if let Some(baseline) = fork_baseline { - CodexAccountingMode::ValidatedBaseline { + } else if let Some(baseline) = parent_fork_baseline { + let reparse_cached_file = matching_cached_fork_state.is_some_and(|state| { + state.locally_resolved || state.inherited_totals.as_ref() != Some(&baseline) + }); + let cached_parent_state = matching_cached_fork_state.filter(|state| { + !state.locally_resolved && state.inherited_totals.as_ref() == Some(&baseline) + }); + CodexAccountingMode::Baseline { baseline, paginated_continuation, - remaining_inherited_totals, - locally_resolved: matching_cached_fork_state - .is_some_and(|state| state.locally_resolved), + remaining_inherited_totals: cached_parent_state + .and_then(|state| state.remaining_inherited_totals.clone()), + provenance: CodexBaselineProvenance::ValidatedParent { + replaces_cached_state: reparse_cached_file, + }, + } + } else if let Some(state) = matching_cached_fork_state + && let Some(baseline) = state.inherited_totals.clone() + { + CodexAccountingMode::Baseline { + baseline, + paginated_continuation, + remaining_inherited_totals: state.remaining_inherited_totals.clone(), + provenance: if state.locally_resolved { + CodexBaselineProvenance::CachedLocalInference + } else { + CodexBaselineProvenance::CachedValidatedParent + }, } } else if session_metadata.is_subagent { CodexAccountingMode::InferSubagent { @@ -575,6 +610,7 @@ impl CostScanner { && cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) && (entry.codex_file_identity.is_none() || identity_matches_cached(entry)) && !cached_identity_changed + && !accounting_mode.requires_cached_reparse() { let (session_cost, has_tokens) = add_codex_days_map_to_summary(summary, &entry.days, range); @@ -680,9 +716,13 @@ impl CostScanner { } } - let parse_target_size = cached - .as_ref() - .and_then(|entry| codex_resumable_scan_target_size(size, entry)); + let parse_target_size = (!accounting_mode.requires_cached_reparse()) + .then(|| { + cached + .as_ref() + .and_then(|entry| codex_resumable_scan_target_size(size, entry)) + }) + .flatten(); let parse_result = match match &accounting_mode { CodexAccountingMode::Standard => JsonlScanner::parse_codex_file_with_state_bounded( path, @@ -695,7 +735,7 @@ impl CostScanner { cancel, max_bytes_to_read, ), - CodexAccountingMode::ValidatedBaseline { + CodexAccountingMode::Baseline { baseline, paginated_continuation, remaining_inherited_totals, diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index a225233a7d..be8d3972e5 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -343,6 +343,65 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { assert!(!state.locally_resolved); } +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "parent-id", + base + Duration::seconds(10), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (_, _, inferred_cache) = scanner.scan_codex_detailed_with_cache(None); + let inferred_state = inferred_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(inferred_state.locally_resolved); + assert_eq!( + inferred_state.inherited_totals.as_ref().unwrap().input, + 5_000 + ); + + write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + + let (_, stats, validated_cache) = scanner.scan_codex_detailed_with_cache(None); + let validated_state = validated_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(!validated_state.locally_resolved); + assert_eq!( + validated_state.inherited_totals.as_ref().unwrap().input, + 1_000 + ); + assert!( + stats + .codex_history_read_paths + .contains(&child.to_string_lossy().to_string()), + "the unchanged child must be reparsed when baseline provenance changes" + ); +} + #[test] fn paginated_continuation_raises_inherited_baseline_from_total_last() { let root = tempfile::tempdir().unwrap(); From 3521a2389252a21d9b6a6639fd9a604350907f31 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 23:33:53 +0700 Subject: [PATCH 29/62] Fix Codex parent baseline transition order --- rust/src/cost_scanner/codex.rs | 40 +++++++++++++------ rust/src/cost_scanner/codex/logical_target.rs | 32 +++++++++++++++ rust/src/cost_scanner/codex/scan.rs | 1 + rust/src/cost_scanner/tests/paginated.rs | 32 +++++++++++++-- 4 files changed, 88 insertions(+), 17 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 7a9d3068dc..1f6d41eb9b 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -96,24 +96,38 @@ fn summary_from_cached_report( } fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { - if usage + let locally_resolved = usage .codex_fork_accounting_state .as_ref() - .is_some_and(|state| state.locally_resolved) - { - return true; - } + .is_some_and(|state| state.locally_resolved); let uses_parent_baseline = usage.codex_lineage.uses_parent_baseline() || (matches!(usage.codex_lineage, CodexSessionLineage::Root) && usage.codex_forked_from_id.is_some()); - !uses_parent_baseline - || usage - .codex_forked_from_id - .as_deref() - .is_some_and(|parent_id| { - codex_parent_baseline(cache, parent_id, usage.codex_fork_timestamp.as_deref()) - .is_some() - }) + if !uses_parent_baseline { + return true; + } + let parent_is_available = usage + .codex_forked_from_id + .as_deref() + .is_some_and(|parent_id| { + codex_parent_baseline(cache, parent_id, usage.codex_fork_timestamp.as_deref()).is_some() + }); + + // Local inference is safe only while no validated parent is available. + // Once the parent enters the cache, force the child through baseline + // replacement instead of accepting its unchanged-file fast path. + if locally_resolved { + !parent_is_available + } else { + parent_is_available + } +} + +fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) } /// Return a parent cumulative baseline only when exactly one cached session diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index a98fb1e63e..3738a85f87 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -42,10 +42,42 @@ pub(super) fn cached_codex_file_is_complete_for_range( && codex_scan_target_size(usage) == size && usage.parsed_bytes.unwrap_or(0) >= size && !usage.codex_unresolved_fork_parent + // Reconsider locally inferred children after this pass has + // had a chance to discover and cache their parent. + && !super::codex_fork_uses_local_inference(usage) && super::codex_fork_parent_is_safe(cache, usage) }) } +/// Process cached local-inference children after all other candidates. A +/// parent discovered in this pass must enter the cache before its unchanged +/// child can decide whether the inferred baseline is still authoritative. +pub(super) fn defer_codex_locally_inferred_candidates( + candidates: &mut Vec, + cache: &CostUsageCache, +) { + if candidates.len() < 2 { + return; + } + + let mut other = Vec::with_capacity(candidates.len()); + let mut locally_inferred = Vec::new(); + for candidate in candidates.drain(..) { + let path_key = candidate.path.to_string_lossy(); + if cache + .files + .get(path_key.as_ref()) + .is_some_and(super::codex_fork_uses_local_inference) + { + locally_inferred.push(candidate); + } else { + other.push(candidate); + } + } + other.extend(locally_inferred); + candidates.extend(other); +} + /// Give paths already in the durable queue their saved turn before newly /// discovered dirty paths. The scanner appends unfinished paths after this /// pass, making the queue a round-robin cursor instead of a newest-first loop. diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index e983adfd6a..ab91037a59 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -198,6 +198,7 @@ pub(super) fn scan_codex_detailed_with_cache( let mut pending_next = cache.codex_pending_paths.clone(); let pending_paths_before_pass = cache.codex_pending_paths.clone(); prioritize_codex_pending_candidates(&mut candidates, &pending_paths_before_pass); + defer_codex_locally_inferred_candidates(&mut candidates, &cache); if discovery_complete && !is_cancelled(cancel) { pending_next .retain(|path| !cached_codex_file_is_complete_for_range(&cache, path, scan_range)); diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index be8d3972e5..d377728c4f 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -326,6 +326,13 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { base + Duration::seconds(10), true, ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); let mut options = CostScanOptions::app_driven(); options.prefer_newest_codex_sessions_first = false; let scanner = CostScanner::new(7) @@ -343,8 +350,9 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { assert!(!state.locally_resolved); } -#[test] -fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { +fn assert_cached_inference_is_replaced_when_parent_appears( + prefer_newest_codex_sessions_first: bool, +) { let root = tempfile::tempdir().unwrap(); let sessions = root.path().join("sessions"); let cache_root = root.path().join("cache"); @@ -357,7 +365,7 @@ fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { true, ); let mut options = CostScanOptions::app_driven(); - options.prefer_newest_codex_sessions_first = false; + options.prefer_newest_codex_sessions_first = prefer_newest_codex_sessions_first; let scanner = CostScanner::new(7) .with_options(options) .with_cache_root(&cache_root) @@ -374,7 +382,7 @@ fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { 5_000 ); - write_codex_fork_session_fixture( + let parent = write_codex_fork_session_fixture( &sessions, "parent.jsonl", "parent-id", @@ -383,6 +391,12 @@ fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { base, &[1_000], ); + std::fs::OpenOptions::new() + .write(true) + .open(parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); let (_, stats, validated_cache) = scanner.scan_codex_detailed_with_cache(None); let validated_state = validated_cache.files[&child.to_string_lossy().to_string()] @@ -402,6 +416,16 @@ fn copied_prefix_subagent_replaces_cached_inference_when_parent_appears() { ); } +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_parent_is_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(true); +} + +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_child_would_be_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(false); +} + #[test] fn paginated_continuation_raises_inherited_baseline_from_total_last() { let root = tempfile::tempdir().unwrap(); From f7d1ad974c9e8d136610d1a59c55b5cde781266c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 23:43:42 +0700 Subject: [PATCH 30/62] Fix fork cache replacement test --- rust/src/cost_scanner/tests/paginated.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index d377728c4f..f4ce5fc138 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -391,6 +391,7 @@ fn assert_cached_inference_is_replaced_when_parent_appears( base, &[1_000], ); + let now = std::time::SystemTime::now(); std::fs::OpenOptions::new() .write(true) .open(parent) From 0735ff5a68a8ceb832b38916f89178eb7c8c410b Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 00:02:59 +0700 Subject: [PATCH 31/62] Reconcile inferred Codex forks in one scan --- rust/src/cost_scanner/codex/scan.rs | 72 ++++++++++++++++++++++++++++- 1 file changed, 70 insertions(+), 2 deletions(-) diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index ab91037a59..54eab62379 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -205,6 +205,7 @@ pub(super) fn scan_codex_detailed_with_cache( } let mut incomplete_processed = Vec::new(); + let mut locally_inferred_complete_paths = Vec::new(); for (index, candidate) in candidates.iter().enumerate() { if is_cancelled(cancel) || index >= candidate_limit @@ -272,8 +273,75 @@ pub(super) fn scan_codex_detailed_with_cache( if !outcome.is_complete || has_unconsumed_tail { incomplete_processed.push(key); stats.files_deferred = stats.files_deferred.saturating_add(1); - } else if let Some(plan) = codex_source_row_plan(&cache, &candidate.path, scan_range) { - apply_codex_source_row_plan(&mut cache, &key, plan); + } else { + if cache + .files + .get(&key) + .is_some_and(codex_fork_uses_local_inference) + { + locally_inferred_complete_paths.push(candidate.path.clone()); + } + if let Some(plan) = codex_source_row_plan(&cache, &candidate.path, scan_range) { + apply_codex_source_row_plan(&mut cache, &key, plan); + } + } + } + + // A child can be visited before its parent during a cold scan. Once the + // remaining candidates have populated the cache, replace that temporary + // local inference in the same refresh instead of publishing it for one + // cycle. Reconciliation still consumes the normal byte budget; work that + // no longer fits is queued for the next explicit refresh. + for path in locally_inferred_complete_paths { + let key = path.to_string_lossy().to_string(); + let parent_is_now_available = cache.files.get(&key).is_some_and(|usage| { + codex_fork_uses_local_inference(usage) && !codex_fork_parent_is_safe(&cache, usage) + }); + if !parent_is_now_available { + continue; + } + let allowance = + per_file_limit.min(refresh_byte_limit.saturating_sub(bytes_read_this_refresh)); + if is_cancelled(cancel) || allowance <= 0 { + if !pending_next.contains(&key) { + pending_next.push(key); + } + stats.files_deferred = stats.files_deferred.saturating_add(1); + continue; + } + + let outcome = scanner.parse_codex_file_bounded( + &path, + scan_range, + &mut summary, + &mut cache, + cancel, + &mut stats, + Some(allowance), + ); + bytes_read_this_refresh = bytes_read_this_refresh.saturating_add(outcome.bytes_read.max(0)); + stats.codex_bytes_read = stats + .codex_bytes_read + .saturating_add(u64::try_from(outcome.bytes_read.max(0)).unwrap_or(u64::MAX)); + pending_next.retain(|pending| pending != &key); + let observed_size = fs::metadata(&path) + .ok() + .map(|metadata| { + #[allow( + clippy::cast_possible_wrap, + reason = "file sizes are clamped to i64::MAX" + )] + let size = metadata.len().min(i64::MAX as u64) as i64; + size + }) + .unwrap_or(0); + let has_unconsumed_tail = cache + .files + .get(&key) + .is_some_and(|usage| codex_logical_target_has_unconsumed_tail(observed_size, usage)); + if !outcome.is_complete || has_unconsumed_tail { + incomplete_processed.push(key); + stats.files_deferred = stats.files_deferred.saturating_add(1); } } pending_next.extend(incomplete_processed); From 361655e80169d68270927107d0ec1dbce7d8b833 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 00:22:47 +0700 Subject: [PATCH 32/62] Order Codex scan work by lineage --- rust/src/cost_scanner/codex.rs | 27 ++- rust/src/cost_scanner/codex/logical_target.rs | 55 +++++++ rust/src/cost_scanner/codex/scan.rs | 143 ++++++---------- rust/src/cost_scanner/tests/paginated.rs | 155 +++++++++++++++++- 4 files changed, 277 insertions(+), 103 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 1f6d41eb9b..39f96eea2c 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -1,5 +1,5 @@ use super::*; -use crate::core::{CodexForkAccountingState, CodexSessionLineage}; +use crate::core::{CodexForkAccountingState, CodexSessionLineage, CodexSessionMetadata}; mod cache_days; mod logical_target; @@ -204,6 +204,11 @@ struct CodexScanCandidate { mtime_unix_ms: i64, } +struct CodexPreparedCandidate { + path: PathBuf, + session_metadata: CodexSessionMetadata, +} + #[derive(Debug, Clone, Copy, Default)] struct CodexFileScanOutcome { bytes_read: i64, @@ -381,7 +386,8 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, ) { - let _ = self.parse_codex_file_bounded(path, range, summary, cache, cancel, stats, None); + let _ = + self.parse_codex_file_bounded(path, range, summary, cache, cancel, stats, None, None); } #[allow( @@ -397,11 +403,14 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, max_bytes_to_read: Option, + prepared_session_metadata: Option<&CodexSessionMetadata>, ) -> CodexFileScanOutcome { if is_cancelled(cancel) { return CodexFileScanOutcome::default(); } - stats.files_seen = stats.files_seen.saturating_add(1); + if prepared_session_metadata.is_none() { + stats.files_seen = stats.files_seen.saturating_add(1); + } let metadata = match fs::metadata(path) { Ok(metadata) => metadata, @@ -461,10 +470,14 @@ impl CostScanner { }; } - stats.codex_metadata_read_paths.push(path_key.clone()); - stats.codex_read_receipt.metadata_reads = - stats.codex_read_receipt.metadata_reads.saturating_add(1); - let session_metadata = JsonlScanner::read_codex_session_metadata(path).unwrap_or_default(); + let session_metadata = if let Some(prepared) = prepared_session_metadata { + prepared.clone() + } else { + stats.codex_metadata_read_paths.push(path_key.clone()); + stats.codex_read_receipt.metadata_reads = + stats.codex_read_receipt.metadata_reads.saturating_add(1); + JsonlScanner::read_codex_session_metadata(path).unwrap_or_default() + }; let cached_identity_matches = cached .as_ref() .is_some_and(|entry| entry.mtime_unix_ms == mtime_ms && entry.size == size); diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 3738a85f87..d5b0350396 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -78,6 +78,61 @@ pub(super) fn defer_codex_locally_inferred_candidates( candidates.extend(other); } +/// Order one bounded work set so every uniquely identified parent is parsed +/// before its children. The sort is stable for unrelated candidates and falls +/// back to discovery order for duplicate identities or dependency cycles. +pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec) { + if candidates.len() < 2 { + return; + } + + let mut session_owners = HashMap::>::new(); + for (index, candidate) in candidates.iter().enumerate() { + let Some(session_id) = candidate.session_metadata.session_id.as_ref() else { + continue; + }; + session_owners + .entry(session_id.clone()) + .and_modify(|owner| *owner = None) + .or_insert(Some(index)); + } + let parent_indices = candidates + .iter() + .map(|candidate| { + candidate + .session_metadata + .forked_from_id + .as_ref() + .and_then(|parent_id| session_owners.get(parent_id)) + .copied() + .flatten() + }) + .collect::>(); + let mut remaining = candidates.drain(..).map(Some).collect::>(); + let mut ordered = Vec::with_capacity(remaining.len()); + + loop { + let mut progressed = false; + for index in 0..remaining.len() { + if remaining[index].is_none() { + continue; + } + let parent_is_ready = + parent_indices[index].is_none_or(|parent_index| remaining[parent_index].is_none()); + if parent_is_ready { + ordered.push(remaining[index].take().expect("candidate checked above")); + progressed = true; + } + } + if !progressed { + break; + } + } + + ordered.extend(remaining.into_iter().flatten()); + candidates.extend(ordered); +} + /// Give paths already in the durable queue their saved turn before newly /// discovered dirty paths. The scanner appends unfinished paths after this /// pass, making the queue a round-robin cursor instead of a newest-first loop. diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 54eab62379..b83fe08f6c 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -204,38 +204,49 @@ pub(super) fn scan_codex_detailed_with_cache( .retain(|path| !cached_codex_file_is_complete_for_range(&cache, path, scan_range)); } - let mut incomplete_processed = Vec::new(); - let mut locally_inferred_complete_paths = Vec::new(); - for (index, candidate) in candidates.iter().enumerate() { - if is_cancelled(cancel) - || index >= candidate_limit - || bytes_read_this_refresh >= refresh_byte_limit - { - for deferred in &candidates[index..] { - let key = deferred.path.to_string_lossy().to_string(); - if !pending_next.contains(&key) { - pending_next.push(key); - } - } - stats.files_deferred = stats.files_deferred.saturating_add( - u32::try_from((candidates.len() - index).min(u32::MAX as usize)) - .unwrap_or(u32::MAX), - ); - break; + // Admit one bounded set, inspect each admitted candidate once, and order + // that set by lineage before reading token history. This makes cold + // child-before-parent scans parent-first without a second parse pass. + let deferred_candidates = candidates.split_off(candidate_limit.min(candidates.len())); + let deferred_paths = deferred_candidates + .into_iter() + .map(|candidate| candidate.path) + .collect::>(); + let mut work_queue = Vec::with_capacity(candidates.len()); + let mut cancelled_during_preparation = Vec::new(); + for candidate in candidates { + if is_cancelled(cancel) { + cancelled_during_preparation.push(candidate.path); + continue; } + let key = candidate.path.to_string_lossy().to_string(); + stats.files_seen = stats.files_seen.saturating_add(1); + stats.codex_metadata_read_paths.push(key); + stats.codex_read_receipt.metadata_reads = + stats.codex_read_receipt.metadata_reads.saturating_add(1); + work_queue.push(CodexPreparedCandidate { + session_metadata: JsonlScanner::read_codex_session_metadata(&candidate.path) + .unwrap_or_default(), + path: candidate.path, + }); + } + let mut unprocessed = Vec::new(); + if !cancelled_during_preparation.is_empty() || is_cancelled(cancel) { + unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); + unprocessed.extend(cancelled_during_preparation); + } else { + order_codex_candidates_by_lineage(&mut work_queue); + } + let mut incomplete_processed = Vec::new(); + for (index, candidate) in work_queue.iter().enumerate() { let refresh_remaining = refresh_byte_limit.saturating_sub(bytes_read_this_refresh); let allowance = per_file_limit.min(refresh_remaining); - if allowance <= 0 { - for deferred in &candidates[index..] { - let key = deferred.path.to_string_lossy().to_string(); - if !pending_next.contains(&key) { - pending_next.push(key); - } - } - stats.files_deferred = stats.files_deferred.saturating_add( - u32::try_from((candidates.len() - index).min(u32::MAX as usize)) - .unwrap_or(u32::MAX), + if is_cancelled(cancel) || allowance <= 0 { + unprocessed.extend( + work_queue[index..] + .iter() + .map(|candidate| candidate.path.clone()), ); break; } @@ -248,6 +259,7 @@ pub(super) fn scan_codex_detailed_with_cache( cancel, &mut stats, Some(allowance), + Some(&candidate.session_metadata), ); bytes_read_this_refresh = bytes_read_this_refresh.saturating_add(outcome.bytes_read.max(0)); stats.codex_bytes_read = stats @@ -273,75 +285,18 @@ pub(super) fn scan_codex_detailed_with_cache( if !outcome.is_complete || has_unconsumed_tail { incomplete_processed.push(key); stats.files_deferred = stats.files_deferred.saturating_add(1); - } else { - if cache - .files - .get(&key) - .is_some_and(codex_fork_uses_local_inference) - { - locally_inferred_complete_paths.push(candidate.path.clone()); - } - if let Some(plan) = codex_source_row_plan(&cache, &candidate.path, scan_range) { - apply_codex_source_row_plan(&mut cache, &key, plan); - } + } else if let Some(plan) = codex_source_row_plan(&cache, &candidate.path, scan_range) { + apply_codex_source_row_plan(&mut cache, &key, plan); } } - - // A child can be visited before its parent during a cold scan. Once the - // remaining candidates have populated the cache, replace that temporary - // local inference in the same refresh instead of publishing it for one - // cycle. Reconciliation still consumes the normal byte budget; work that - // no longer fits is queued for the next explicit refresh. - for path in locally_inferred_complete_paths { + unprocessed.extend(deferred_paths); + stats.files_deferred = stats.files_deferred.saturating_add( + u32::try_from(unprocessed.len().min(u32::MAX as usize)).unwrap_or(u32::MAX), + ); + for path in unprocessed { let key = path.to_string_lossy().to_string(); - let parent_is_now_available = cache.files.get(&key).is_some_and(|usage| { - codex_fork_uses_local_inference(usage) && !codex_fork_parent_is_safe(&cache, usage) - }); - if !parent_is_now_available { - continue; - } - let allowance = - per_file_limit.min(refresh_byte_limit.saturating_sub(bytes_read_this_refresh)); - if is_cancelled(cancel) || allowance <= 0 { - if !pending_next.contains(&key) { - pending_next.push(key); - } - stats.files_deferred = stats.files_deferred.saturating_add(1); - continue; - } - - let outcome = scanner.parse_codex_file_bounded( - &path, - scan_range, - &mut summary, - &mut cache, - cancel, - &mut stats, - Some(allowance), - ); - bytes_read_this_refresh = bytes_read_this_refresh.saturating_add(outcome.bytes_read.max(0)); - stats.codex_bytes_read = stats - .codex_bytes_read - .saturating_add(u64::try_from(outcome.bytes_read.max(0)).unwrap_or(u64::MAX)); - pending_next.retain(|pending| pending != &key); - let observed_size = fs::metadata(&path) - .ok() - .map(|metadata| { - #[allow( - clippy::cast_possible_wrap, - reason = "file sizes are clamped to i64::MAX" - )] - let size = metadata.len().min(i64::MAX as u64) as i64; - size - }) - .unwrap_or(0); - let has_unconsumed_tail = cache - .files - .get(&key) - .is_some_and(|usage| codex_logical_target_has_unconsumed_tail(observed_size, usage)); - if !outcome.is_complete || has_unconsumed_tail { - incomplete_processed.push(key); - stats.files_deferred = stats.files_deferred.saturating_add(1); + if !pending_next.contains(&key) { + pending_next.push(key); } } pending_next.extend(incomplete_processed); diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index f4ce5fc138..a42ad77405 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -310,7 +310,7 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { let sessions = root.path().join("sessions"); let cache_root = root.path().join("cache"); let base = Utc::now() - Duration::hours(1); - write_codex_fork_session_fixture( + let parent = write_codex_fork_session_fixture( &sessions, "parent.jsonl", "parent-id", @@ -335,12 +335,16 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { .unwrap(); let mut options = CostScanOptions::app_driven(); options.prefer_newest_codex_sessions_first = false; + let parent_size = std::fs::metadata(&parent).unwrap().len(); + let child_size = std::fs::metadata(&child).unwrap().len(); + options.codex_max_session_file_bytes = + i64::try_from(parent_size.max(child_size)).expect("fixture size fits i64"); let scanner = CostScanner::new(7) .with_options(options) .with_cache_root(&cache_root) .with_sessions_dirs(vec![sessions]); - let (_, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); let state = cache.files[&child.to_string_lossy().to_string()] .codex_fork_accounting_state .as_ref() @@ -348,6 +352,153 @@ fn copied_prefix_subagent_prefers_validated_parent_baseline() { assert_eq!(state.inherited_totals.as_ref().unwrap().input, 1_000); assert!(!state.locally_resolved); + assert_eq!(stats.files_seen, 2); + assert_eq!(stats.codex_read_receipt.metadata_reads, 2); + assert_eq!(stats.codex_read_receipt.history_reads, 2); + assert_eq!( + stats.codex_bytes_read, + parent_size.saturating_add(child_size), + "one bounded parse per candidate must enforce the per-file allowance" + ); + assert_eq!( + stats.codex_history_read_paths, + vec![ + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); +} + +#[test] +fn candidate_limit_counts_each_child_parent_candidate_once() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "parent-id", + base + Duration::seconds(10), + true, + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); + std::fs::OpenOptions::new() + .write(true) + .open(&parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + options.codex_candidate_limit = 1; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 1); + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 1); + assert_eq!( + stats.codex_metadata_read_paths, + vec![child.to_string_lossy().to_string()] + ); + assert_eq!( + cache.codex_pending_paths, + vec![parent.to_string_lossy().to_string()] + ); +} + +#[test] +fn cold_scan_orders_multi_level_parent_chain_before_children() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let ancestor = write_codex_fork_session_fixture( + &sessions, + "ancestor.jsonl", + "ancestor-id", + None, + base, + base, + &[1_000], + ); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + Some("ancestor-id"), + base + Duration::seconds(10), + base + Duration::seconds(10), + &[1_500], + ); + let child = write_codex_fork_session_fixture( + &sessions, + "child.jsonl", + "child-id", + Some("parent-id"), + base + Duration::seconds(20), + base + Duration::seconds(20), + &[2_000], + ); + let now = std::time::SystemTime::now(); + for (path, age) in [(&child, 30), (&parent, 20), (&ancestor, 10)] { + std::fs::OpenOptions::new() + .write(true) + .open(path) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(age)) + .unwrap(); + } + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 3); + assert_eq!(stats.codex_read_receipt.metadata_reads, 3); + assert_eq!(stats.codex_read_receipt.history_reads, 3); + assert_eq!( + stats.codex_history_read_paths, + vec![ + ancestor.to_string_lossy().to_string(), + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); + let parent_state = cache.files[&parent.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + let child_state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert_eq!(parent_state.inherited_totals.as_ref().unwrap().input, 1_000); + assert_eq!(child_state.inherited_totals.as_ref().unwrap().input, 1_500); } fn assert_cached_inference_is_replaced_when_parent_appears( From 4d0b5993a5b772f491c950b68be2b79a6acdc2ee Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 00:40:49 +0700 Subject: [PATCH 33/62] Fail closed on ambiguous Codex lineage --- rust/src/cost_scanner/codex.rs | 25 +- rust/src/cost_scanner/codex/logical_target.rs | 61 +- rust/src/cost_scanner/codex/scan.rs | 3 +- rust/src/cost_scanner/tests.rs | 3 + rust/src/cost_scanner/tests/copied_prefix.rs | 578 ++++++++++++++++++ rust/src/cost_scanner/tests/paginated.rs | 443 -------------- 6 files changed, 641 insertions(+), 472 deletions(-) create mode 100644 rust/src/cost_scanner/tests/copied_prefix.rs diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 39f96eea2c..aab7def35d 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -207,6 +207,14 @@ struct CodexScanCandidate { struct CodexPreparedCandidate { path: PathBuf, session_metadata: CodexSessionMetadata, + lineage_disposition: CodexLineageDisposition, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +enum CodexLineageDisposition { + #[default] + Ready, + AmbiguousOrCyclic, } #[derive(Debug, Clone, Copy, Default)] @@ -403,12 +411,12 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, max_bytes_to_read: Option, - prepared_session_metadata: Option<&CodexSessionMetadata>, + prepared_candidate: Option<&CodexPreparedCandidate>, ) -> CodexFileScanOutcome { if is_cancelled(cancel) { return CodexFileScanOutcome::default(); } - if prepared_session_metadata.is_none() { + if prepared_candidate.is_none() { stats.files_seen = stats.files_seen.saturating_add(1); } @@ -454,6 +462,9 @@ impl CostScanner { // before reading even the bounded metadata prefix; raw token history // is only needed after freshness fails or a fork needs reconciliation. if let Some(entry) = cached.as_ref() + && prepared_candidate.is_none_or(|candidate| { + candidate.lineage_disposition == CodexLineageDisposition::Ready + }) && cache_entry_is_fresh(entry) && identity_matches_cached(entry) { @@ -470,8 +481,8 @@ impl CostScanner { }; } - let session_metadata = if let Some(prepared) = prepared_session_metadata { - prepared.clone() + let session_metadata = if let Some(prepared) = prepared_candidate { + prepared.session_metadata.clone() } else { stats.codex_metadata_read_paths.push(path_key.clone()); stats.codex_read_receipt.metadata_reads = @@ -565,7 +576,11 @@ impl CostScanner { && history_base_thread_id .as_deref() .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); - let accounting_mode = if !is_fork { + let accounting_mode = if prepared_candidate.is_some_and(|candidate| { + candidate.lineage_disposition == CodexLineageDisposition::AmbiguousOrCyclic + }) { + CodexAccountingMode::Unresolved + } else if !is_fork { CodexAccountingMode::Standard } else if let Some(baseline) = parent_fork_baseline { let reparse_cached_file = matching_cached_fork_state.is_some_and(|state| { diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index d5b0350396..b22b7ee6c3 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -79,48 +79,60 @@ pub(super) fn defer_codex_locally_inferred_candidates( } /// Order one bounded work set so every uniquely identified parent is parsed -/// before its children. The sort is stable for unrelated candidates and falls -/// back to discovery order for duplicate identities or dependency cycles. +/// before its children. Duplicate identities, cycles, and every dependent +/// candidate are marked unsafe so parsing cannot accept or infer a baseline +/// from ambiguous lineage. pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec) { - if candidates.len() < 2 { + if candidates.is_empty() { return; } - let mut session_owners = HashMap::>::new(); + let mut session_owners = HashMap::>::new(); for (index, candidate) in candidates.iter().enumerate() { let Some(session_id) = candidate.session_metadata.session_id.as_ref() else { continue; }; session_owners .entry(session_id.clone()) - .and_modify(|owner| *owner = None) - .or_insert(Some(index)); + .or_default() + .push(index); } - let parent_indices = candidates - .iter() - .map(|candidate| { - candidate - .session_metadata - .forked_from_id - .as_ref() - .and_then(|parent_id| session_owners.get(parent_id)) - .copied() - .flatten() - }) - .collect::>(); + let mut unsafe_lineage = vec![false; candidates.len()]; + for owners in session_owners.values().filter(|owners| owners.len() > 1) { + for &index in owners { + unsafe_lineage[index] = true; + } + } + let mut parent_indices = vec![None; candidates.len()]; + for (index, candidate) in candidates.iter().enumerate() { + let Some(parent_id) = candidate.session_metadata.forked_from_id.as_ref() else { + continue; + }; + match session_owners.get(parent_id).map(Vec::as_slice) { + Some([parent_index]) => parent_indices[index] = Some(*parent_index), + Some([]) | None => {} + Some(_) => unsafe_lineage[index] = true, + } + } + let mut remaining = candidates.drain(..).map(Some).collect::>(); let mut ordered = Vec::with_capacity(remaining.len()); + let mut completed = vec![false; remaining.len()]; loop { let mut progressed = false; for index in 0..remaining.len() { - if remaining[index].is_none() { + if remaining[index].is_none() || unsafe_lineage[index] { continue; } - let parent_is_ready = - parent_indices[index].is_none_or(|parent_index| remaining[parent_index].is_none()); + let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { + completed[parent_index] && !unsafe_lineage[parent_index] + }); if parent_is_ready { - ordered.push(remaining[index].take().expect("candidate checked above")); + let mut candidate = remaining[index].take().expect("candidate checked above"); + candidate.lineage_disposition = CodexLineageDisposition::Ready; + ordered.push(candidate); + completed[index] = true; progressed = true; } } @@ -129,7 +141,10 @@ pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec, + owned: bool, +) -> PathBuf { + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let mut lines = vec![ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), + "payload": { + "id": session_id, "forked_from_id": parent_id, + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} + } + }), + token_row(base, 2, [1_000, 900, 100], [0, 0, 0], "gpt-5.6-sol"), + serde_json::json!({ + "type": "turn_context", "ordinal": 10, "timestamp": base.to_rfc3339(), + "payload": {"model": "gpt-5.6-sol"} + }), + token_row( + base, + 12, + [1_000, 900, 100], + [1_000, 900, 100], + "gpt-5.6-sol", + ), + token_row( + base, + 13, + [5_000, 3_900, 500], + [5_000, 3_900, 500], + "gpt-5.6-sol", + ), + ]; + if owned { + lines.extend([ + token_row(base, 19, [5_050, 3_910, 505], [50, 10, 5], "gpt-5.6-sol"), + token_row( + base + Duration::seconds(1), + 20, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + token_row( + base + Duration::seconds(2), + 21, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + ]); + } + let body = lines + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn token_row( + timestamp: DateTime, + ordinal: i64, + total: [i64; 3], + last: [i64; 3], + model: &str, +) -> serde_json::Value { + serde_json::json!({ + "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": model, + "total_token_usage": { + "input_tokens": total[0], "cached_input_tokens": total[1], "output_tokens": total[2] + }, + "last_token_usage": { + "input_tokens": last[0], "cached_input_tokens": last[1], "output_tokens": last[2] + } + }} + }) +} + +#[test] +fn copied_prefix_subagent_infers_advancing_baseline_without_parent() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "missing-parent", + Utc::now() - Duration::hours(1), + true, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 70); + assert_eq!(summary.cached_tokens, 15); + assert_eq!(summary.output_tokens, 10); + assert_eq!(summary.sessions_count, 1); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(!usage.codex_unresolved_fork_parent); + assert!( + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + ); + assert_eq!( + usage.days.values().next().unwrap()["gpt-5.6-sol"], + vec![70, 15, 10] + ); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 70); + assert!(stats.codex_history_read_paths.is_empty()); +} + +#[test] +fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "missing-parent", + Utc::now() - Duration::hours(1), + false, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.output_tokens, 0); + assert_eq!(summary.sessions_count, 0); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(usage.days.is_empty()); + assert!(!usage.codex_unresolved_fork_parent); + let state = usage.codex_fork_accounting_state.as_ref().unwrap(); + assert!(state.locally_resolved); + assert!(state.inherited_totals.is_none()); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 0); + assert_eq!(cached.output_tokens, 0); + assert_eq!(cached.sessions_count, 0); + assert!(stats.codex_history_read_paths.is_empty()); +} + +#[test] +fn copied_prefix_subagent_prefers_validated_parent_baseline() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let parent_size = std::fs::metadata(&parent).unwrap().len(); + let child_size = std::fs::metadata(&child).unwrap().len(); + options.codex_max_session_file_bytes = + i64::try_from(parent_size.max(child_size)).expect("fixture size fits i64"); + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + let state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + + assert_eq!(state.inherited_totals.as_ref().unwrap().input, 1_000); + assert!(!state.locally_resolved); + assert_eq!(stats.files_seen, 2); + assert_eq!(stats.codex_read_receipt.metadata_reads, 2); + assert_eq!(stats.codex_read_receipt.history_reads, 2); + assert_eq!( + stats.codex_bytes_read, + parent_size.saturating_add(child_size), + "one bounded parse per candidate must enforce the per-file allowance" + ); + assert_eq!( + stats.codex_history_read_paths, + vec![ + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); +} + +#[test] +fn candidate_limit_counts_each_child_parent_candidate_once() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); + std::fs::OpenOptions::new() + .write(true) + .open(&parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + options.codex_candidate_limit = 1; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 1); + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 1); + assert_eq!( + stats.codex_metadata_read_paths, + vec![child.to_string_lossy().to_string()] + ); + assert_eq!( + cache.codex_pending_paths, + vec![parent.to_string_lossy().to_string()] + ); +} + +#[test] +fn cold_scan_orders_multi_level_parent_chain_before_children() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let ancestor = write_codex_fork_session_fixture( + &sessions, + "ancestor.jsonl", + "ancestor-id", + None, + base, + base, + &[1_000], + ); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + Some("ancestor-id"), + base + Duration::seconds(10), + base + Duration::seconds(10), + &[1_500], + ); + let child = write_codex_fork_session_fixture( + &sessions, + "child.jsonl", + "child-id", + Some("parent-id"), + base + Duration::seconds(20), + base + Duration::seconds(20), + &[2_000], + ); + let now = std::time::SystemTime::now(); + for (path, age) in [(&child, 30), (&parent, 20), (&ancestor, 10)] { + std::fs::OpenOptions::new() + .write(true) + .open(path) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(age)) + .unwrap(); + } + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 3); + assert_eq!(stats.codex_read_receipt.metadata_reads, 3); + assert_eq!(stats.codex_read_receipt.history_reads, 3); + assert_eq!( + stats.codex_history_read_paths, + vec![ + ancestor.to_string_lossy().to_string(), + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); + let parent_state = cache.files[&parent.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + let child_state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert_eq!(parent_state.inherited_totals.as_ref().unwrap().input, 1_000); + assert_eq!(child_state.inherited_totals.as_ref().unwrap().input, 1_500); +} + +fn assert_unsafe_lineage_is_unresolved( + summary: &CostSummary, + stats: &CostScanStats, + cache: &CostUsageCache, + paths: &[&Path], +) { + assert_eq!(summary.sessions_count, 0); + assert_eq!(summary.input_tokens, 0); + assert_eq!( + stats.codex_read_receipt.metadata_reads, + u32::try_from(paths.len()).expect("fixture count fits u32") + ); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert!(stats.codex_history_read_paths.is_empty()); + for path in paths { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.codex_fork_accounting_state.is_none()); + assert!(usage.days.is_empty()); + } +} + +#[test] +fn duplicate_parent_session_ids_fail_closed_with_their_child() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first_parent = write_codex_fork_session_fixture( + &sessions, + "first-parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let second_parent = write_codex_fork_session_fixture( + &sessions, + "second-parent.jsonl", + "parent-id", + None, + base + Duration::seconds(1), + base + Duration::seconds(1), + &[2_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(2), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved( + &summary, + &stats, + &cache, + &[&first_parent, &second_parent, &child], + ); +} + +#[test] +fn two_node_subagent_cycle_fails_closed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first = write_copied_prefix_subagent_fixture( + &sessions, + "first.jsonl", + "first-id", + "second-id", + base, + true, + ); + let second = write_copied_prefix_subagent_fixture( + &sessions, + "second.jsonl", + "second-id", + "first-id", + base + Duration::seconds(1), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved(&summary, &stats, &cache, &[&first, &second]); +} + +#[test] +fn self_referential_subagent_fails_closed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let session = write_copied_prefix_subagent_fixture( + &sessions, + "self-cycle.jsonl", + "self-id", + "self-id", + Utc::now() - Duration::hours(1), + true, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved(&summary, &stats, &cache, &[&session]); +} + +fn assert_cached_inference_is_replaced_when_parent_appears( + prefer_newest_codex_sessions_first: bool, +) { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = prefer_newest_codex_sessions_first; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (_, _, inferred_cache) = scanner.scan_codex_detailed_with_cache(None); + let inferred_state = inferred_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(inferred_state.locally_resolved); + assert_eq!( + inferred_state.inherited_totals.as_ref().unwrap().input, + 5_000 + ); + + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); + + let (_, stats, validated_cache) = scanner.scan_codex_detailed_with_cache(None); + let validated_state = validated_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(!validated_state.locally_resolved); + assert_eq!( + validated_state.inherited_totals.as_ref().unwrap().input, + 1_000 + ); + assert!( + stats + .codex_history_read_paths + .contains(&child.to_string_lossy().to_string()), + "the unchanged child must be reparsed when baseline provenance changes" + ); +} + +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_parent_is_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(true); +} + +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_child_would_be_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(false); +} diff --git a/rust/src/cost_scanner/tests/paginated.rs b/rust/src/cost_scanner/tests/paginated.rs index a42ad77405..592d484823 100644 --- a/rust/src/cost_scanner/tests/paginated.rs +++ b/rust/src/cost_scanner/tests/paginated.rs @@ -135,449 +135,6 @@ fn write_codex_paginated_continuation_fixture( path } -fn write_copied_prefix_subagent_fixture( - sessions_root: &Path, - name: &str, - parent_id: &str, - base: DateTime, - owned: bool, -) -> PathBuf { - let day = base.with_timezone(&Local).date_naive(); - let day_dir = sessions_root - .join(day.format("%Y").to_string()) - .join(day.format("%m").to_string()) - .join(day.format("%d").to_string()); - std::fs::create_dir_all(&day_dir).unwrap(); - let path = day_dir.join(name); - let mut lines = vec![ - serde_json::json!({ - "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), - "payload": { - "id": "child-id", "forked_from_id": parent_id, - "subagent_history_start_ordinal": 10, - "thread_source": "subagent", - "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} - } - }), - token_row(base, 2, [1_000, 900, 100], [0, 0, 0], "gpt-5.6-sol"), - serde_json::json!({ - "type": "turn_context", "ordinal": 10, "timestamp": base.to_rfc3339(), - "payload": {"model": "gpt-5.6-sol"} - }), - token_row( - base, - 12, - [1_000, 900, 100], - [1_000, 900, 100], - "gpt-5.6-sol", - ), - token_row( - base, - 13, - [5_000, 3_900, 500], - [5_000, 3_900, 500], - "gpt-5.6-sol", - ), - ]; - if owned { - lines.extend([ - token_row(base, 19, [5_050, 3_910, 505], [50, 10, 5], "gpt-5.6-sol"), - token_row( - base + Duration::seconds(1), - 20, - [5_070, 3_915, 510], - [20, 5, 5], - "gpt-5.6-sol", - ), - token_row( - base + Duration::seconds(2), - 21, - [5_070, 3_915, 510], - [20, 5, 5], - "gpt-5.6-sol", - ), - ]); - } - let body = lines - .into_iter() - .map(|line| line.to_string()) - .collect::>() - .join("\n") - + "\n"; - std::fs::write(&path, body).unwrap(); - path -} - -fn token_row( - timestamp: DateTime, - ordinal: i64, - total: [i64; 3], - last: [i64; 3], - model: &str, -) -> serde_json::Value { - serde_json::json!({ - "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), - "payload": {"type": "token_count", "info": { - "model": model, - "total_token_usage": { - "input_tokens": total[0], "cached_input_tokens": total[1], "output_tokens": total[2] - }, - "last_token_usage": { - "input_tokens": last[0], "cached_input_tokens": last[1], "output_tokens": last[2] - } - }} - }) -} - -#[test] -fn copied_prefix_subagent_infers_advancing_baseline_without_parent() { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let child = write_copied_prefix_subagent_fixture( - &sessions, - "child.jsonl", - "missing-parent", - Utc::now() - Duration::hours(1), - true, - ); - let scanner = CostScanner::new(7) - .with_options(CostScanOptions::app_driven()) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions]); - - let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); - assert_eq!(summary.input_tokens, 70); - assert_eq!(summary.cached_tokens, 15); - assert_eq!(summary.output_tokens, 10); - assert_eq!(summary.sessions_count, 1); - let usage = &cache.files[&child.to_string_lossy().to_string()]; - assert!(!usage.codex_unresolved_fork_parent); - assert!( - usage - .codex_fork_accounting_state - .as_ref() - .is_some_and(|state| state.locally_resolved) - ); - assert_eq!( - usage.days.values().next().unwrap()["gpt-5.6-sol"], - vec![70, 15, 10] - ); - - let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); - assert_eq!(cached.input_tokens, 70); - assert!(stats.codex_history_read_paths.is_empty()); -} - -#[test] -fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let child = write_copied_prefix_subagent_fixture( - &sessions, - "child.jsonl", - "missing-parent", - Utc::now() - Duration::hours(1), - false, - ); - let scanner = CostScanner::new(7) - .with_options(CostScanOptions::app_driven()) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions]); - - let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); - assert_eq!(summary.input_tokens, 0); - assert_eq!(summary.output_tokens, 0); - assert_eq!(summary.sessions_count, 0); - let usage = &cache.files[&child.to_string_lossy().to_string()]; - assert!(usage.days.is_empty()); - assert!(!usage.codex_unresolved_fork_parent); - let state = usage.codex_fork_accounting_state.as_ref().unwrap(); - assert!(state.locally_resolved); - assert!(state.inherited_totals.is_none()); - - let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); - assert_eq!(cached.input_tokens, 0); - assert_eq!(cached.output_tokens, 0); - assert_eq!(cached.sessions_count, 0); - assert!(stats.codex_history_read_paths.is_empty()); -} - -#[test] -fn copied_prefix_subagent_prefers_validated_parent_baseline() { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let base = Utc::now() - Duration::hours(1); - let parent = write_codex_fork_session_fixture( - &sessions, - "parent.jsonl", - "parent-id", - None, - base, - base, - &[1_000], - ); - let child = write_copied_prefix_subagent_fixture( - &sessions, - "child.jsonl", - "parent-id", - base + Duration::seconds(10), - true, - ); - let now = std::time::SystemTime::now(); - std::fs::OpenOptions::new() - .write(true) - .open(&child) - .unwrap() - .set_modified(now - std::time::Duration::from_secs(20)) - .unwrap(); - let mut options = CostScanOptions::app_driven(); - options.prefer_newest_codex_sessions_first = false; - let parent_size = std::fs::metadata(&parent).unwrap().len(); - let child_size = std::fs::metadata(&child).unwrap().len(); - options.codex_max_session_file_bytes = - i64::try_from(parent_size.max(child_size)).expect("fixture size fits i64"); - let scanner = CostScanner::new(7) - .with_options(options) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions]); - - let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); - let state = cache.files[&child.to_string_lossy().to_string()] - .codex_fork_accounting_state - .as_ref() - .unwrap(); - - assert_eq!(state.inherited_totals.as_ref().unwrap().input, 1_000); - assert!(!state.locally_resolved); - assert_eq!(stats.files_seen, 2); - assert_eq!(stats.codex_read_receipt.metadata_reads, 2); - assert_eq!(stats.codex_read_receipt.history_reads, 2); - assert_eq!( - stats.codex_bytes_read, - parent_size.saturating_add(child_size), - "one bounded parse per candidate must enforce the per-file allowance" - ); - assert_eq!( - stats.codex_history_read_paths, - vec![ - parent.to_string_lossy().to_string(), - child.to_string_lossy().to_string(), - ] - ); -} - -#[test] -fn candidate_limit_counts_each_child_parent_candidate_once() { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let base = Utc::now() - Duration::hours(1); - let parent = write_codex_fork_session_fixture( - &sessions, - "parent.jsonl", - "parent-id", - None, - base, - base, - &[1_000], - ); - let child = write_copied_prefix_subagent_fixture( - &sessions, - "child.jsonl", - "parent-id", - base + Duration::seconds(10), - true, - ); - let now = std::time::SystemTime::now(); - std::fs::OpenOptions::new() - .write(true) - .open(&child) - .unwrap() - .set_modified(now - std::time::Duration::from_secs(20)) - .unwrap(); - std::fs::OpenOptions::new() - .write(true) - .open(&parent) - .unwrap() - .set_modified(now - std::time::Duration::from_secs(10)) - .unwrap(); - let mut options = CostScanOptions::app_driven(); - options.prefer_newest_codex_sessions_first = false; - options.codex_candidate_limit = 1; - let scanner = CostScanner::new(7) - .with_options(options) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions]); - - let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); - - assert_eq!(stats.files_seen, 1); - assert_eq!(stats.codex_read_receipt.metadata_reads, 1); - assert_eq!(stats.codex_read_receipt.history_reads, 1); - assert_eq!( - stats.codex_metadata_read_paths, - vec![child.to_string_lossy().to_string()] - ); - assert_eq!( - cache.codex_pending_paths, - vec![parent.to_string_lossy().to_string()] - ); -} - -#[test] -fn cold_scan_orders_multi_level_parent_chain_before_children() { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let base = Utc::now() - Duration::hours(1); - let ancestor = write_codex_fork_session_fixture( - &sessions, - "ancestor.jsonl", - "ancestor-id", - None, - base, - base, - &[1_000], - ); - let parent = write_codex_fork_session_fixture( - &sessions, - "parent.jsonl", - "parent-id", - Some("ancestor-id"), - base + Duration::seconds(10), - base + Duration::seconds(10), - &[1_500], - ); - let child = write_codex_fork_session_fixture( - &sessions, - "child.jsonl", - "child-id", - Some("parent-id"), - base + Duration::seconds(20), - base + Duration::seconds(20), - &[2_000], - ); - let now = std::time::SystemTime::now(); - for (path, age) in [(&child, 30), (&parent, 20), (&ancestor, 10)] { - std::fs::OpenOptions::new() - .write(true) - .open(path) - .unwrap() - .set_modified(now - std::time::Duration::from_secs(age)) - .unwrap(); - } - let mut options = CostScanOptions::app_driven(); - options.prefer_newest_codex_sessions_first = false; - let scanner = CostScanner::new(7) - .with_options(options) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions]); - - let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); - - assert_eq!(stats.files_seen, 3); - assert_eq!(stats.codex_read_receipt.metadata_reads, 3); - assert_eq!(stats.codex_read_receipt.history_reads, 3); - assert_eq!( - stats.codex_history_read_paths, - vec![ - ancestor.to_string_lossy().to_string(), - parent.to_string_lossy().to_string(), - child.to_string_lossy().to_string(), - ] - ); - let parent_state = cache.files[&parent.to_string_lossy().to_string()] - .codex_fork_accounting_state - .as_ref() - .unwrap(); - let child_state = cache.files[&child.to_string_lossy().to_string()] - .codex_fork_accounting_state - .as_ref() - .unwrap(); - assert_eq!(parent_state.inherited_totals.as_ref().unwrap().input, 1_000); - assert_eq!(child_state.inherited_totals.as_ref().unwrap().input, 1_500); -} - -fn assert_cached_inference_is_replaced_when_parent_appears( - prefer_newest_codex_sessions_first: bool, -) { - let root = tempfile::tempdir().unwrap(); - let sessions = root.path().join("sessions"); - let cache_root = root.path().join("cache"); - let base = Utc::now() - Duration::hours(1); - let child = write_copied_prefix_subagent_fixture( - &sessions, - "child.jsonl", - "parent-id", - base + Duration::seconds(10), - true, - ); - let mut options = CostScanOptions::app_driven(); - options.prefer_newest_codex_sessions_first = prefer_newest_codex_sessions_first; - let scanner = CostScanner::new(7) - .with_options(options) - .with_cache_root(&cache_root) - .with_sessions_dirs(vec![sessions.clone()]); - - let (_, _, inferred_cache) = scanner.scan_codex_detailed_with_cache(None); - let inferred_state = inferred_cache.files[&child.to_string_lossy().to_string()] - .codex_fork_accounting_state - .as_ref() - .unwrap(); - assert!(inferred_state.locally_resolved); - assert_eq!( - inferred_state.inherited_totals.as_ref().unwrap().input, - 5_000 - ); - - let parent = write_codex_fork_session_fixture( - &sessions, - "parent.jsonl", - "parent-id", - None, - base, - base, - &[1_000], - ); - let now = std::time::SystemTime::now(); - std::fs::OpenOptions::new() - .write(true) - .open(parent) - .unwrap() - .set_modified(now - std::time::Duration::from_secs(10)) - .unwrap(); - - let (_, stats, validated_cache) = scanner.scan_codex_detailed_with_cache(None); - let validated_state = validated_cache.files[&child.to_string_lossy().to_string()] - .codex_fork_accounting_state - .as_ref() - .unwrap(); - assert!(!validated_state.locally_resolved); - assert_eq!( - validated_state.inherited_totals.as_ref().unwrap().input, - 1_000 - ); - assert!( - stats - .codex_history_read_paths - .contains(&child.to_string_lossy().to_string()), - "the unchanged child must be reparsed when baseline provenance changes" - ); -} - -#[test] -fn copied_prefix_subagent_replaces_cached_inference_when_parent_is_visited_first() { - assert_cached_inference_is_replaced_when_parent_appears(true); -} - -#[test] -fn copied_prefix_subagent_replaces_cached_inference_when_child_would_be_visited_first() { - assert_cached_inference_is_replaced_when_parent_appears(false); -} - #[test] fn paginated_continuation_raises_inherited_baseline_from_total_last() { let root = tempfile::tempdir().unwrap(); From 9b2361adc3c9702926158d74b1fc5cc43dabf1f0 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 01:25:36 +0700 Subject: [PATCH 34/62] Validate Codex lineage across refreshes --- rust/src/cost_scanner/codex.rs | 190 ++++++++++------ rust/src/cost_scanner/codex/logical_target.rs | 163 +++++++++++--- rust/src/cost_scanner/codex/scan.rs | 15 +- rust/src/cost_scanner/tests.rs | 3 + rust/src/cost_scanner/tests/lineage_cache.rs | 212 ++++++++++++++++++ 5 files changed, 492 insertions(+), 91 deletions(-) create mode 100644 rust/src/cost_scanner/tests/lineage_cache.rs diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index aab7def35d..23022e2226 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -95,31 +95,42 @@ fn summary_from_cached_report( } } +#[derive(Debug, Clone, PartialEq, Eq)] +enum CodexParentResolution { + Absent, + Safe(crate::core::CodexTotals), + Unsafe, +} + +fn codex_usage_uses_parent(usage: &CostUsageFileUsage) -> bool { + usage.codex_lineage.uses_parent_baseline() + || (matches!(usage.codex_lineage, CodexSessionLineage::Root) + && usage.codex_forked_from_id.is_some()) +} + fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { let locally_resolved = usage .codex_fork_accounting_state .as_ref() .is_some_and(|state| state.locally_resolved); - let uses_parent_baseline = usage.codex_lineage.uses_parent_baseline() - || (matches!(usage.codex_lineage, CodexSessionLineage::Root) - && usage.codex_forked_from_id.is_some()); - if !uses_parent_baseline { + if !codex_usage_uses_parent(usage) { return true; } - let parent_is_available = usage - .codex_forked_from_id - .as_deref() - .is_some_and(|parent_id| { - codex_parent_baseline(cache, parent_id, usage.codex_fork_timestamp.as_deref()).is_some() - }); + let parent_resolution = + usage + .codex_forked_from_id + .as_deref() + .map_or(CodexParentResolution::Unsafe, |parent_id| { + codex_parent_resolution(cache, parent_id, usage.codex_fork_timestamp.as_deref()) + }); - // Local inference is safe only while no validated parent is available. - // Once the parent enters the cache, force the child through baseline - // replacement instead of accepting its unchanged-file fast path. + // Local inference is safe only while the parent is genuinely absent. + // An owner that is ambiguous, stale, locally inferred, cyclic, or + // transitively unsafe must fail closed instead of looking absent. if locally_resolved { - !parent_is_available + matches!(parent_resolution, CodexParentResolution::Absent) } else { - parent_is_available + matches!(parent_resolution, CodexParentResolution::Safe(_)) } } @@ -130,51 +141,101 @@ fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { .is_some_and(|state| state.locally_resolved) } -/// Return a parent cumulative baseline only when exactly one cached session -/// identity is current, complete, timestamp-ordered, and safe to trust. -fn codex_parent_baseline( +/// Resolve one parent identity through the persisted cache graph. Absence is +/// deliberately distinct from ambiguity or transitive unsafety so copied +/// prefixes may infer only when no owner exists at all. +fn codex_parent_resolution( + cache: &CostUsageCache, + parent_session_id: &str, + child_fork_timestamp: Option<&str>, +) -> CodexParentResolution { + codex_parent_resolution_inner( + cache, + parent_session_id, + child_fork_timestamp, + &mut HashSet::new(), + ) +} + +fn codex_parent_resolution_inner( cache: &CostUsageCache, parent_session_id: &str, child_fork_timestamp: Option<&str>, + visiting: &mut HashSet, +) -> CodexParentResolution { + let mut owners = cache + .files + .iter() + .filter(|(_, usage)| usage.codex_session_id.as_deref() == Some(parent_session_id)); + let Some((path_key, usage)) = owners.next() else { + return CodexParentResolution::Absent; + }; + if owners.next().is_some() || !visiting.insert(path_key.clone()) { + return CodexParentResolution::Unsafe; + } + + let resolution = + codex_parent_owner_baseline(cache, path_key, usage, child_fork_timestamp, visiting) + .map_or(CodexParentResolution::Unsafe, CodexParentResolution::Safe); + visiting.remove(path_key); + resolution +} + +fn codex_parent_owner_baseline( + cache: &CostUsageCache, + path_key: &str, + usage: &CostUsageFileUsage, + child_fork_timestamp: Option<&str>, + visiting: &mut HashSet, ) -> Option { - let mut baseline = None; - for (path_key, usage) in &cache.files { - if usage.codex_session_id.as_deref() != Some(parent_session_id) { - continue; - } - if usage.codex_unresolved_fork_parent - || usage.codex_token_timestamps_monotonic != Some(true) - { - return None; - } - let metadata = fs::metadata(path_key).ok()?; - if let (Some(expected), Some(actual)) = ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata), - ) && expected != &actual - { - return None; - } - #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] - let size = metadata.len().min(i64::MAX as u64) as i64; - if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) - || usage.size != size - || usage.parsed_bytes.unwrap_or(0) < size - { - return None; - } - let last_totals = usage.last_totals.clone()?; - let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; - let child_fork_timestamp = child_fork_timestamp?; - if !JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) { - return None; - } - if baseline.replace(last_totals).is_some() { - // Duplicate identities make the dependency ambiguous. - return None; + if usage.codex_unresolved_fork_parent + || usage.codex_token_timestamps_monotonic != Some(true) + || codex_fork_uses_local_inference(usage) + { + return None; + } + + if codex_usage_uses_parent(usage) { + let parent_id = usage.codex_forked_from_id.as_deref()?; + let inherited = usage + .codex_fork_accounting_state + .as_ref()? + .inherited_totals + .as_ref()?; + match codex_parent_resolution_inner( + cache, + parent_id, + usage.codex_fork_timestamp.as_deref(), + visiting, + ) { + CodexParentResolution::Safe(baseline) if &baseline == inherited => {} + CodexParentResolution::Absent + | CodexParentResolution::Safe(_) + | CodexParentResolution::Unsafe => return None, } } - baseline + + let metadata = fs::metadata(path_key).ok()?; + if let (Some(expected), Some(actual)) = ( + usage.codex_file_identity.as_ref(), + JsonlScanner::codex_file_identity(Path::new(path_key), &metadata), + ) && expected != &actual + { + return None; + } + #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] + let size = metadata.len().min(i64::MAX as u64) as i64; + if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) + || usage.size != size + || usage.parsed_bytes.unwrap_or(0) < size + { + return None; + } + let last_totals = usage.last_totals.clone()?; + let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; + let child_fork_timestamp = child_fork_timestamp?; + JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) + .then_some(last_totals) } fn is_codex_path_in_scan_window( @@ -565,11 +626,11 @@ impl CostScanner { let matching_cached_fork_state = cached_fork_accounting_state .as_ref() .filter(|_| cached_fork_state_matches); - let parent_fork_baseline = is_fork + let parent_resolution = is_fork .then_some(codex_forked_from_id.as_deref()) .flatten() - .and_then(|parent_id| { - codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) + .map_or(CodexParentResolution::Unsafe, |parent_id| { + codex_parent_resolution(cache, parent_id, codex_fork_timestamp.as_deref()) }); let paginated_continuation = is_fork && codex_forked_from_id.is_some() @@ -582,15 +643,15 @@ impl CostScanner { CodexAccountingMode::Unresolved } else if !is_fork { CodexAccountingMode::Standard - } else if let Some(baseline) = parent_fork_baseline { + } else if let CodexParentResolution::Safe(baseline) = &parent_resolution { let reparse_cached_file = matching_cached_fork_state.is_some_and(|state| { - state.locally_resolved || state.inherited_totals.as_ref() != Some(&baseline) + state.locally_resolved || state.inherited_totals.as_ref() != Some(baseline) }); let cached_parent_state = matching_cached_fork_state.filter(|state| { - !state.locally_resolved && state.inherited_totals.as_ref() == Some(&baseline) + !state.locally_resolved && state.inherited_totals.as_ref() == Some(baseline) }); CodexAccountingMode::Baseline { - baseline, + baseline: baseline.clone(), paginated_continuation, remaining_inherited_totals: cached_parent_state .and_then(|state| state.remaining_inherited_totals.clone()), @@ -598,7 +659,8 @@ impl CostScanner { replaces_cached_state: reparse_cached_file, }, } - } else if let Some(state) = matching_cached_fork_state + } else if matches!(&parent_resolution, CodexParentResolution::Absent) + && let Some(state) = matching_cached_fork_state && let Some(baseline) = state.inherited_totals.clone() { CodexAccountingMode::Baseline { @@ -611,7 +673,9 @@ impl CostScanner { CodexBaselineProvenance::CachedValidatedParent }, } - } else if session_metadata.is_subagent { + } else if matches!(&parent_resolution, CodexParentResolution::Absent) + && session_metadata.is_subagent + { CodexAccountingMode::InferSubagent { start_ordinal: session_metadata.subagent_history_start_ordinal, } diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index b22b7ee6c3..25eb78d8e8 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -78,18 +78,77 @@ pub(super) fn defer_codex_locally_inferred_candidates( candidates.extend(other); } -/// Order one bounded work set so every uniquely identified parent is parsed -/// before its children. Duplicate identities, cycles, and every dependent -/// candidate are marked unsafe so parsing cannot accept or infer a baseline -/// from ambiguous lineage. -pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec) { +struct CodexLineageNode { + path: String, + session_id: Option, + parent_id: Option, + candidate_index: Option, + may_infer_missing_parent: bool, + may_author_parent: bool, + initially_unsafe: bool, +} + +/// Order one bounded work set against both its admitted metadata and the +/// persisted cache graph. The returned cache paths became structurally unsafe +/// and must be invalidated even when the candidate limit deferred them. +pub(super) fn order_codex_candidates_by_lineage( + cache: &CostUsageCache, + candidates: &mut Vec, +) -> Vec { if candidates.is_empty() { - return; + return Vec::new(); + } + + let candidate_paths = candidates + .iter() + .map(|candidate| candidate.path.to_string_lossy().to_string()) + .collect::>(); + let mut cached_paths = cache + .files + .keys() + .filter(|path| !candidate_paths.contains(*path)) + .cloned() + .collect::>(); + cached_paths.sort(); + + let mut nodes = Vec::with_capacity(cached_paths.len() + candidates.len()); + for path in cached_paths { + let usage = &cache.files[&path]; + let uses_parent = super::codex_usage_uses_parent(usage); + let locally_inferred = super::codex_fork_uses_local_inference(usage); + nodes.push(CodexLineageNode { + path, + session_id: usage.codex_session_id.clone(), + parent_id: uses_parent + .then(|| usage.codex_forked_from_id.clone()) + .flatten(), + candidate_index: None, + may_infer_missing_parent: locally_inferred, + may_author_parent: !locally_inferred && !usage.codex_unresolved_fork_parent, + initially_unsafe: usage.codex_unresolved_fork_parent, + }); + } + let mut candidate_node_indices = Vec::with_capacity(candidates.len()); + for (candidate_index, candidate) in candidates.iter().enumerate() { + let uses_parent = candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some(); + nodes.push(CodexLineageNode { + path: candidate.path.to_string_lossy().to_string(), + session_id: candidate.session_metadata.session_id.clone(), + parent_id: uses_parent + .then(|| candidate.session_metadata.forked_from_id.clone()) + .flatten(), + candidate_index: Some(candidate_index), + may_infer_missing_parent: candidate.session_metadata.is_subagent, + may_author_parent: true, + initially_unsafe: false, + }); + candidate_node_indices.push(nodes.len() - 1); } let mut session_owners = HashMap::>::new(); - for (index, candidate) in candidates.iter().enumerate() { - let Some(session_id) = candidate.session_metadata.session_id.as_ref() else { + for (index, node) in nodes.iter().enumerate() { + let Some(session_id) = node.session_id.as_ref() else { continue; }; session_owners @@ -97,42 +156,47 @@ pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec>(); for owners in session_owners.values().filter(|owners| owners.len() > 1) { for &index in owners { unsafe_lineage[index] = true; } } - let mut parent_indices = vec![None; candidates.len()]; - for (index, candidate) in candidates.iter().enumerate() { - let Some(parent_id) = candidate.session_metadata.forked_from_id.as_ref() else { + let mut parent_indices = vec![None; nodes.len()]; + for (index, node) in nodes.iter().enumerate() { + let Some(parent_id) = node.parent_id.as_ref() else { continue; }; match session_owners.get(parent_id).map(Vec::as_slice) { Some([parent_index]) => parent_indices[index] = Some(*parent_index), - Some([]) | None => {} + Some([]) | None if node.may_infer_missing_parent => {} + Some([]) | None => unsafe_lineage[index] = true, Some(_) => unsafe_lineage[index] = true, } } - let mut remaining = candidates.drain(..).map(Some).collect::>(); - let mut ordered = Vec::with_capacity(remaining.len()); - let mut completed = vec![false; remaining.len()]; + let mut completed = vec![false; nodes.len()]; + let mut ordered_indices = Vec::with_capacity(candidates.len()); loop { let mut progressed = false; - for index in 0..remaining.len() { - if remaining[index].is_none() || unsafe_lineage[index] { + for index in 0..nodes.len() { + if completed[index] || unsafe_lineage[index] { continue; } let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { - completed[parent_index] && !unsafe_lineage[parent_index] + completed[parent_index] + && !unsafe_lineage[parent_index] + && nodes[parent_index].may_author_parent }); if parent_is_ready { - let mut candidate = remaining[index].take().expect("candidate checked above"); - candidate.lineage_disposition = CodexLineageDisposition::Ready; - ordered.push(candidate); completed[index] = true; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_indices.push(candidate_index); + } progressed = true; } } @@ -141,11 +205,58 @@ pub(super) fn order_codex_candidates_by_lineage(candidates: &mut Vec>(); + for candidate_index in ordered_indices { + let node_index = candidate_node_indices[candidate_index]; + let mut candidate = remaining[candidate_index] + .take() + .expect("candidate is ordered once"); + candidate.lineage_disposition = if unsafe_lineage[node_index] { + CodexLineageDisposition::AmbiguousOrCyclic + } else { + CodexLineageDisposition::Ready + }; + candidates.push(candidate); + } + + nodes + .iter() + .zip(unsafe_lineage) + .filter(|(node, unsafe_lineage)| { + *unsafe_lineage + && cache + .files + .get(&node.path) + .is_some_and(|usage| !usage.codex_unresolved_fork_parent) + }) + .map(|(node, _)| node.path.clone()) + .collect() +} + +pub(super) fn invalidate_codex_unsafe_lineage(cache: &mut CostUsageCache, paths: &[String]) { + for path in paths { + let Some(usage) = cache.files.get_mut(path) else { + continue; + }; + usage.days.clear(); + usage.parsed_bytes = Some(0); + usage.codex_scan_target_size = None; + usage.last_model = None; + usage.last_totals = None; + usage.codex_token_timestamps_monotonic = None; + usage.codex_last_token_timestamp = None; + usage.codex_fork_accounting_state = None; + usage.codex_unresolved_fork_parent = true; } - candidates.extend(ordered); } /// Give paths already in the durable queue their saved turn before newly diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 2353b80dad..6da7fb5188 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -197,6 +197,7 @@ pub(super) fn scan_codex_detailed_with_cache( let mut bytes_read_this_refresh = 0_i64; let mut pending_next = cache.codex_pending_paths.clone(); let pending_paths_before_pass = cache.codex_pending_paths.clone(); + let mut invalidated_unsafe_lineage = false; prioritize_codex_pending_candidates(&mut candidates, &pending_paths_before_pass); defer_codex_locally_inferred_candidates(&mut candidates, &cache); if discovery_complete && !is_cancelled(cancel) { @@ -236,7 +237,17 @@ pub(super) fn scan_codex_detailed_with_cache( unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); unprocessed.extend(cancelled_during_preparation); } else { - order_codex_candidates_by_lineage(&mut work_queue); + let unsafe_cached_paths = order_codex_candidates_by_lineage(&cache, &mut work_queue); + invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); + if invalidated_unsafe_lineage { + cache.previous_report = None; + } + invalidate_codex_unsafe_lineage(&mut cache, &unsafe_cached_paths); + for path in unsafe_cached_paths { + if !pending_next.contains(&path) { + pending_next.push(path); + } + } } let mut incomplete_processed = Vec::new(); @@ -369,7 +380,7 @@ pub(super) fn scan_codex_detailed_with_cache( // the range so unchanged files stay on the cache fast path. cache.scan_since_key = Some(scan_range.scan_since_key.clone()); cache.scan_until_key = Some(scan_range.scan_until_key.clone()); - } else if cache.previous_report.is_none() { + } else if cache.previous_report.is_none() && !invalidated_unsafe_lineage { cache.previous_report = established_report_before_scan; } if !is_cancelled(cancel) { diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index fb592e7180..d8dc136047 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -3040,5 +3040,8 @@ fn incomplete_or_buffered_empty_codex_fragment_is_not_marked_complete() { #[path = "tests/copied_prefix.rs"] mod copied_prefix; #[cfg(test)] +#[path = "tests/lineage_cache.rs"] +mod lineage_cache; +#[cfg(test)] #[path = "tests/paginated.rs"] mod paginated; diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs new file mode 100644 index 0000000000..18c742ba72 --- /dev/null +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -0,0 +1,212 @@ +use super::*; + +fn write_subagent( + sessions_root: &Path, + name: &str, + session_id: &str, + parent_id: &str, + timestamp: DateTime, +) -> PathBuf { + let day = timestamp.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let rows = [ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": timestamp.to_rfc3339(), + "payload": { + "id": session_id, + "forked_from_id": parent_id, + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} + } + }), + lineage_token_row(timestamp, 2, 1_000, 0), + serde_json::json!({ + "type": "turn_context", "ordinal": 10, "timestamp": timestamp.to_rfc3339(), + "payload": {"model": "gpt-5.6-sol"} + }), + lineage_token_row(timestamp, 12, 1_000, 1_000), + lineage_token_row(timestamp + Duration::seconds(1), 20, 1_050, 50), + ]; + let body = rows + .into_iter() + .map(|row| row.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn lineage_token_row( + timestamp: DateTime, + ordinal: i64, + total_input: i64, + last_input: i64, +) -> serde_json::Value { + serde_json::json!({ + "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": { + "input_tokens": total_input, "cached_input_tokens": 0, "output_tokens": 5 + }, + "last_token_usage": { + "input_tokens": last_input, "cached_input_tokens": 0, "output_tokens": 5 + } + }} + }) +} + +fn bounded_scanner(sessions: &Path, cache_root: &Path) -> CostScanner { + let mut options = CostScanOptions::app_driven(); + options.codex_candidate_limit = 1; + options.prefer_newest_codex_sessions_first = false; + CostScanner::new(7) + .with_options(options) + .with_cache_root(cache_root) + .with_sessions_dirs(vec![sessions.to_path_buf()]) +} + +fn assert_locally_inferred(cache: &CostUsageCache, path: &Path) { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(!usage.codex_unresolved_fork_parent); + assert!( + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + ); +} + +fn assert_unresolved(cache: &CostUsageCache, path: &Path) { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + assert!(usage.codex_fork_accounting_state.is_none()); +} + +#[test] +fn bounded_refresh_detects_duplicate_parent_owners_across_cache_and_candidate() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent(&sessions, "child.jsonl", "child-id", "parent-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (_, first_stats, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(first_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(first_stats.codex_read_receipt.history_reads, 1); + assert_locally_inferred(&first_cache, &child); + + let first_parent = write_codex_fork_session_fixture( + &sessions, + "parent-a.jsonl", + "parent-id", + None, + base - Duration::seconds(2), + base - Duration::seconds(2), + &[1_000], + ); + let (_, second_stats, second_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(second_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(second_stats.codex_read_receipt.history_reads, 1); + assert_locally_inferred(&second_cache, &child); + + let second_parent = write_codex_fork_session_fixture( + &sessions, + "parent-b.jsonl", + "parent-id", + None, + base - Duration::seconds(1), + base - Duration::seconds(1), + &[2_000], + ); + let (summary, third_stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(third_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(third_stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &first_parent); + assert_unresolved(&cache, &second_parent); + assert_unresolved(&cache, &child); +} + +#[test] +fn bounded_refresh_detects_equal_timestamp_two_node_cycle() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first = write_subagent(&sessions, "first.jsonl", "first-id", "second-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &first); + + let second = write_subagent(&sessions, "second.jsonl", "second-id", "first-id", base); + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &first); + assert_unresolved(&cache, &second); +} + +#[test] +fn bounded_refresh_rejects_self_cycle_migration() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let session = write_subagent(&sessions, "self.jsonl", "self-id", "missing-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &session); + + write_subagent( + &sessions, + "self.jsonl", + "self-id", + "self-id", + base + Duration::seconds(1), + ); + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &session); +} + +#[test] +fn bounded_refresh_rejects_dependent_of_locally_inferred_parent() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_subagent(&sessions, "parent.jsonl", "parent-id", "missing-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &parent); + + let dependent = write_subagent( + &sessions, + "dependent.jsonl", + "dependent-id", + "parent-id", + base + Duration::seconds(1), + ); + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_locally_inferred(&cache, &parent); + assert_unresolved(&cache, &dependent); +} From 4558dc2efb5d83e88ca348f62c000487cc010931 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 05:23:56 +0700 Subject: [PATCH 35/62] Harden Codex fork lineage validation --- rust/src/core/jsonl_scanner/codex/parser.rs | 30 +- rust/src/core/jsonl_scanner/tests.rs | 40 +++ rust/src/cost_scanner/codex.rs | 231 ++------------- rust/src/cost_scanner/codex/logical_target.rs | 270 ++++++++++++++++-- rust/src/cost_scanner/codex/reconciliation.rs | 39 ++- rust/src/cost_scanner/codex/scan.rs | 4 +- rust/src/cost_scanner/tests/lineage_cache.rs | 195 +++++++++++++ 7 files changed, 562 insertions(+), 247 deletions(-) diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index f33316e378..728bd16f76 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -58,6 +58,7 @@ struct ForkBaselineInference { baseline: Option, boundary_open: bool, inherited_opening: bool, + missing_explicit_ordinal: bool, locally_confirmed: bool, resolved: bool, } @@ -79,11 +80,18 @@ impl ForkBaselineInference { }), boundary_open: false, inherited_opening: false, - locally_confirmed: explicit_start_ordinal.is_some(), + missing_explicit_ordinal: false, + locally_confirmed: false, resolved: false, } } + fn confirm_local_resolution(&mut self) { + if !self.missing_explicit_ordinal { + self.locally_confirmed = true; + } + } + fn observe_non_token(&mut self, obj: &Value) { if obj.get("type").and_then(Value::as_str) == Some("turn_context") && self.inherited_opening { @@ -108,13 +116,19 @@ impl ForkBaselineInference { let last = read_token_totals(last_usage); let ordinal = obj.get("ordinal").and_then(Value::as_i64); + if self.explicit_start_ordinal.is_some() && ordinal.is_none() { + self.missing_explicit_ordinal = true; + self.locally_confirmed = false; + if !self.boundary_open { + self.baseline = Some(total); + } + return ForkBaselineDecision::SkipCopiedPrefix; + } + if let Some(start) = self.explicit_start_ordinal && !self.boundary_open { - let Some(ordinal) = ordinal else { - self.baseline = Some(total); - return ForkBaselineDecision::SkipCopiedPrefix; - }; + let ordinal = ordinal.expect("missing explicit ordinals return above"); if ordinal < start { self.baseline = Some(total); return ForkBaselineDecision::SkipCopiedPrefix; @@ -124,7 +138,7 @@ impl ForkBaselineInference { if totals_contain_usage(&total) && !totals_contain_usage(&last) { self.baseline = Some(total); self.inherited_opening = true; - self.locally_confirmed = true; + self.confirm_local_resolution(); } return ForkBaselineDecision::SkipCopiedPrefix; } else if !self.boundary_open { @@ -152,13 +166,13 @@ impl ForkBaselineInference { totals_contain_usage(&baseline) && total == last && totals_at_least(&total, &baseline); if copied_snapshot { self.baseline = Some(total); - self.locally_confirmed = true; + self.confirm_local_resolution(); return ForkBaselineDecision::SkipCopiedPrefix; } let owned_baseline = totals_delta(&last, &total); self.baseline = Some(owned_baseline.clone()); - self.locally_confirmed = true; + self.confirm_local_resolution(); self.resolved = true; ForkBaselineDecision::ProcessWithBaseline(owned_baseline) } diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index 2614ebefd3..4abfce1c93 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -132,6 +132,7 @@ fn inferred_fork_waits_for_present_explicit_start_ordinal() { assert!(state.records.is_empty()); assert!(state.fork_baseline.is_none()); + assert!(!state.fork_baseline_locally_resolved()); state.process_line( r#"{"ordinal":10,"timestamp":"2026-09-22T10:00:01Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.6-sol","total_token_usage":{"input_tokens":110,"cached_input_tokens":22,"output_tokens":11},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":1}}}}"#, @@ -142,6 +143,45 @@ fn inferred_fork_waits_for_present_explicit_start_ordinal() { assert_eq!(state.records[0].0.input, 10); assert_eq!(state.records[0].0.cached, 2); assert_eq!(state.records[0].0.output, 1); + assert!(!state.fork_baseline_locally_resolved()); +} + +#[test] +fn inferred_fork_keeps_missing_ordinal_unresolved_after_boundary_opens() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + let token_line = |ordinal: Option, total: i64, last: i64| { + let mut value = serde_json::json!({ + "timestamp": "2026-09-22T10:00:00Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": {"input_tokens": total, "cached_input_tokens": 0, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last, "cached_input_tokens": 0, "output_tokens": 0} + }} + }); + if let Some(ordinal) = ordinal { + value["ordinal"] = serde_json::json!(ordinal); + } + value.to_string() + }; + + state.process_line(&token_line(Some(9), 100, 0), &range); + state.process_line(&token_line(Some(10), 100, 0), &range); + state.process_line(&token_line(Some(11), 110, 110), &range); + assert!(state.fork_baseline_locally_resolved()); + state.process_line(&token_line(None, 120, 10), &range); + assert!(!state.fork_baseline_locally_resolved()); + state.process_line(&token_line(Some(12), 130, 10), &range); + + assert_eq!(state.records.len(), 1); + assert_eq!(state.records[0].0.input, 10); + assert!(!state.fork_baseline_locally_resolved()); } #[test] diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 23022e2226..7e522c9da9 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -95,13 +95,6 @@ fn summary_from_cached_report( } } -#[derive(Debug, Clone, PartialEq, Eq)] -enum CodexParentResolution { - Absent, - Safe(crate::core::CodexTotals), - Unsafe, -} - fn codex_usage_uses_parent(usage: &CostUsageFileUsage) -> bool { usage.codex_lineage.uses_parent_baseline() || (matches!(usage.codex_lineage, CodexSessionLineage::Root) @@ -109,29 +102,7 @@ fn codex_usage_uses_parent(usage: &CostUsageFileUsage) -> bool { } fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { - let locally_resolved = usage - .codex_fork_accounting_state - .as_ref() - .is_some_and(|state| state.locally_resolved); - if !codex_usage_uses_parent(usage) { - return true; - } - let parent_resolution = - usage - .codex_forked_from_id - .as_deref() - .map_or(CodexParentResolution::Unsafe, |parent_id| { - codex_parent_resolution(cache, parent_id, usage.codex_fork_timestamp.as_deref()) - }); - - // Local inference is safe only while the parent is genuinely absent. - // An owner that is ambiguous, stale, locally inferred, cyclic, or - // transitively unsafe must fail closed instead of looking absent. - if locally_resolved { - matches!(parent_resolution, CodexParentResolution::Absent) - } else { - matches!(parent_resolution, CodexParentResolution::Safe(_)) - } + CodexLineagePlanner::new(cache).cached_usage_is_safe(usage) } fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { @@ -141,103 +112,6 @@ fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { .is_some_and(|state| state.locally_resolved) } -/// Resolve one parent identity through the persisted cache graph. Absence is -/// deliberately distinct from ambiguity or transitive unsafety so copied -/// prefixes may infer only when no owner exists at all. -fn codex_parent_resolution( - cache: &CostUsageCache, - parent_session_id: &str, - child_fork_timestamp: Option<&str>, -) -> CodexParentResolution { - codex_parent_resolution_inner( - cache, - parent_session_id, - child_fork_timestamp, - &mut HashSet::new(), - ) -} - -fn codex_parent_resolution_inner( - cache: &CostUsageCache, - parent_session_id: &str, - child_fork_timestamp: Option<&str>, - visiting: &mut HashSet, -) -> CodexParentResolution { - let mut owners = cache - .files - .iter() - .filter(|(_, usage)| usage.codex_session_id.as_deref() == Some(parent_session_id)); - let Some((path_key, usage)) = owners.next() else { - return CodexParentResolution::Absent; - }; - if owners.next().is_some() || !visiting.insert(path_key.clone()) { - return CodexParentResolution::Unsafe; - } - - let resolution = - codex_parent_owner_baseline(cache, path_key, usage, child_fork_timestamp, visiting) - .map_or(CodexParentResolution::Unsafe, CodexParentResolution::Safe); - visiting.remove(path_key); - resolution -} - -fn codex_parent_owner_baseline( - cache: &CostUsageCache, - path_key: &str, - usage: &CostUsageFileUsage, - child_fork_timestamp: Option<&str>, - visiting: &mut HashSet, -) -> Option { - if usage.codex_unresolved_fork_parent - || usage.codex_token_timestamps_monotonic != Some(true) - || codex_fork_uses_local_inference(usage) - { - return None; - } - - if codex_usage_uses_parent(usage) { - let parent_id = usage.codex_forked_from_id.as_deref()?; - let inherited = usage - .codex_fork_accounting_state - .as_ref()? - .inherited_totals - .as_ref()?; - match codex_parent_resolution_inner( - cache, - parent_id, - usage.codex_fork_timestamp.as_deref(), - visiting, - ) { - CodexParentResolution::Safe(baseline) if &baseline == inherited => {} - CodexParentResolution::Absent - | CodexParentResolution::Safe(_) - | CodexParentResolution::Unsafe => return None, - } - } - - let metadata = fs::metadata(path_key).ok()?; - if let (Some(expected), Some(actual)) = ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata), - ) && expected != &actual - { - return None; - } - #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] - let size = metadata.len().min(i64::MAX as u64) as i64; - if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) - || usage.size != size - || usage.parsed_bytes.unwrap_or(0) < size - { - return None; - } - let last_totals = usage.last_totals.clone()?; - let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; - let child_fork_timestamp = child_fork_timestamp?; - JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) - .then_some(last_totals) -} - fn is_codex_path_in_scan_window( path: &Path, sessions_dirs: &[PathBuf], @@ -268,14 +142,7 @@ struct CodexScanCandidate { struct CodexPreparedCandidate { path: PathBuf, session_metadata: CodexSessionMetadata, - lineage_disposition: CodexLineageDisposition, -} - -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -enum CodexLineageDisposition { - #[default] - Ready, - AmbiguousOrCyclic, + lineage_gate: CodexLineageGate, } #[derive(Debug, Clone, Copy, Default)] @@ -511,21 +378,19 @@ impl CostScanner { let cache_entry_is_fresh = |entry: &CostUsageFileUsage| { cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) }; - let identity_matches_cached = |entry: &CostUsageFileUsage| match ( - entry.codex_file_identity.as_ref(), - file_identity.as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - _ => false, + let identity_matches_cached = |entry: &CostUsageFileUsage| { + codex_file_identity_matches( + entry.codex_file_identity.as_deref(), + file_identity.as_deref(), + ) }; // The compact cache is authoritative for an unchanged file. Do this // before reading even the bounded metadata prefix; raw token history // is only needed after freshness fails or a fork needs reconciliation. if let Some(entry) = cached.as_ref() - && prepared_candidate.is_none_or(|candidate| { - candidate.lineage_disposition == CodexLineageDisposition::Ready - }) + && prepared_candidate + .is_none_or(|candidate| candidate.lineage_gate == CodexLineageGate::Eligible) && cache_entry_is_fresh(entry) && identity_matches_cached(entry) { @@ -550,9 +415,9 @@ impl CostScanner { stats.codex_read_receipt.metadata_reads.saturating_add(1); JsonlScanner::read_codex_session_metadata(path).unwrap_or_default() }; - let cached_identity_matches = cached - .as_ref() - .is_some_and(|entry| entry.mtime_unix_ms == mtime_ms && entry.size == size); + // Cached lineage metadata belongs to a physical file, not merely a + // path/size/mtime tuple. Missing identity evidence fails closed. + let cached_identity_matches = cached.as_ref().is_some_and(identity_matches_cached); let codex_session_id = session_metadata.session_id.clone().or_else(|| { cached_identity_matches .then(|| cached.as_ref()?.codex_session_id.clone()) @@ -626,62 +491,25 @@ impl CostScanner { let matching_cached_fork_state = cached_fork_accounting_state .as_ref() .filter(|_| cached_fork_state_matches); - let parent_resolution = is_fork - .then_some(codex_forked_from_id.as_deref()) - .flatten() - .map_or(CodexParentResolution::Unsafe, |parent_id| { - codex_parent_resolution(cache, parent_id, codex_fork_timestamp.as_deref()) - }); let paginated_continuation = is_fork && codex_forked_from_id.is_some() && history_base_thread_id .as_deref() .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); - let accounting_mode = if prepared_candidate.is_some_and(|candidate| { - candidate.lineage_disposition == CodexLineageDisposition::AmbiguousOrCyclic - }) { - CodexAccountingMode::Unresolved - } else if !is_fork { - CodexAccountingMode::Standard - } else if let CodexParentResolution::Safe(baseline) = &parent_resolution { - let reparse_cached_file = matching_cached_fork_state.is_some_and(|state| { - state.locally_resolved || state.inherited_totals.as_ref() != Some(baseline) - }); - let cached_parent_state = matching_cached_fork_state.filter(|state| { - !state.locally_resolved && state.inherited_totals.as_ref() == Some(baseline) - }); - CodexAccountingMode::Baseline { - baseline: baseline.clone(), - paginated_continuation, - remaining_inherited_totals: cached_parent_state - .and_then(|state| state.remaining_inherited_totals.clone()), - provenance: CodexBaselineProvenance::ValidatedParent { - replaces_cached_state: reparse_cached_file, - }, - } - } else if matches!(&parent_resolution, CodexParentResolution::Absent) - && let Some(state) = matching_cached_fork_state - && let Some(baseline) = state.inherited_totals.clone() - { - CodexAccountingMode::Baseline { - baseline, - paginated_continuation, - remaining_inherited_totals: state.remaining_inherited_totals.clone(), - provenance: if state.locally_resolved { - CodexBaselineProvenance::CachedLocalInference - } else { - CodexBaselineProvenance::CachedValidatedParent - }, - } - } else if matches!(&parent_resolution, CodexParentResolution::Absent) - && session_metadata.is_subagent - { - CodexAccountingMode::InferSubagent { - start_ordinal: session_metadata.subagent_history_start_ordinal, - } - } else { - CodexAccountingMode::Unresolved - }; + let lineage_gate = prepared_candidate + .map(|candidate| candidate.lineage_gate) + .unwrap_or_default(); + let lineage_decision = CodexLineagePlanner::new(cache).decision_for_scan( + is_fork, + lineage_gate, + codex_forked_from_id.as_deref(), + codex_fork_timestamp.as_deref(), + ); + let accounting_mode = lineage_decision.accounting_mode( + matching_cached_fork_state, + &session_metadata, + paginated_continuation, + ); if accounting_mode.is_unresolved() { cache.files.insert( @@ -714,7 +542,7 @@ impl CostScanner { if let Some(entry) = &cached && cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) - && (entry.codex_file_identity.is_none() || identity_matches_cached(entry)) + && identity_matches_cached(entry) && !cached_identity_changed && !accounting_mode.requires_cached_reparse() { @@ -724,11 +552,6 @@ impl CostScanner { summary.total_cost_usd += session_cost; summary.sessions_count += 1; } - if entry.codex_file_identity != file_identity { - let mut refreshed = entry.clone(); - refreshed.codex_file_identity = file_identity.clone(); - cache.files.insert(path_key.clone(), refreshed); - } stats.files_skipped = stats.files_skipped.saturating_add(1); return CodexFileScanOutcome { bytes_read: 0, diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 25eb78d8e8..6245ea6c3b 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -1,5 +1,221 @@ use super::*; +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(super) enum CodexLineageGate { + #[default] + Eligible, + Unsafe, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum CodexLineageDecision { + Root, + ParentAbsent, + ParentReady(crate::core::CodexTotals), + Unsafe, +} + +pub(super) struct CodexLineagePlanner<'a> { + cache: &'a CostUsageCache, +} + +impl<'a> CodexLineagePlanner<'a> { + pub(super) fn new(cache: &'a CostUsageCache) -> Self { + Self { cache } + } + + pub(super) fn decision_for_scan( + &self, + uses_parent: bool, + gate: CodexLineageGate, + parent_id: Option<&str>, + fork_timestamp: Option<&str>, + ) -> CodexLineageDecision { + if gate == CodexLineageGate::Unsafe { + return CodexLineageDecision::Unsafe; + } + if !uses_parent { + return CodexLineageDecision::Root; + } + parent_id.map_or(CodexLineageDecision::Unsafe, |parent_id| { + self.resolve_parent(parent_id, fork_timestamp) + }) + } + + pub(super) fn decision_for_usage(&self, usage: &CostUsageFileUsage) -> CodexLineageDecision { + if usage.codex_unresolved_fork_parent { + return CodexLineageDecision::Unsafe; + } + if !super::codex_usage_uses_parent(usage) { + return CodexLineageDecision::Root; + } + usage + .codex_forked_from_id + .as_deref() + .map_or(CodexLineageDecision::Unsafe, |parent_id| { + self.resolve_parent(parent_id, usage.codex_fork_timestamp.as_deref()) + }) + } + + pub(super) fn cached_usage_is_safe(&self, usage: &CostUsageFileUsage) -> bool { + let locally_resolved = super::codex_fork_uses_local_inference(usage); + match self.decision_for_usage(usage) { + CodexLineageDecision::Root => true, + CodexLineageDecision::ParentAbsent => locally_resolved, + CodexLineageDecision::ParentReady(_) => !locally_resolved, + CodexLineageDecision::Unsafe => false, + } + } + + /// Resolve one parent identity through the persisted graph. Absence is + /// distinct from ambiguity and transitive unsafety so local inference is + /// allowed only when no owner exists at all. + fn resolve_parent( + &self, + parent_session_id: &str, + child_fork_timestamp: Option<&str>, + ) -> CodexLineageDecision { + self.resolve_parent_inner(parent_session_id, child_fork_timestamp, &mut HashSet::new()) + } + + fn resolve_parent_inner( + &self, + parent_session_id: &str, + child_fork_timestamp: Option<&str>, + visiting: &mut HashSet, + ) -> CodexLineageDecision { + let mut owners = self + .cache + .files + .iter() + .filter(|(_, usage)| usage.codex_session_id.as_deref() == Some(parent_session_id)); + let Some((path_key, usage)) = owners.next() else { + return CodexLineageDecision::ParentAbsent; + }; + if owners.next().is_some() || !visiting.insert(path_key.clone()) { + return CodexLineageDecision::Unsafe; + } + + let decision = self + .parent_owner_baseline(path_key, usage, child_fork_timestamp, visiting) + .map_or( + CodexLineageDecision::Unsafe, + CodexLineageDecision::ParentReady, + ); + visiting.remove(path_key); + decision + } + + fn parent_owner_baseline( + &self, + path_key: &str, + usage: &CostUsageFileUsage, + child_fork_timestamp: Option<&str>, + visiting: &mut HashSet, + ) -> Option { + if usage.codex_unresolved_fork_parent + || usage.codex_token_timestamps_monotonic != Some(true) + || super::codex_fork_uses_local_inference(usage) + { + return None; + } + + if super::codex_usage_uses_parent(usage) { + let parent_id = usage.codex_forked_from_id.as_deref()?; + let inherited = usage + .codex_fork_accounting_state + .as_ref()? + .inherited_totals + .as_ref()?; + match self.resolve_parent_inner( + parent_id, + usage.codex_fork_timestamp.as_deref(), + visiting, + ) { + CodexLineageDecision::ParentReady(baseline) if &baseline == inherited => {} + CodexLineageDecision::Root + | CodexLineageDecision::ParentAbsent + | CodexLineageDecision::ParentReady(_) + | CodexLineageDecision::Unsafe => return None, + } + } + + let metadata = fs::metadata(path_key).ok()?; + let expected_identity = usage.codex_file_identity.as_ref()?; + let actual_identity = JsonlScanner::codex_file_identity(Path::new(path_key), &metadata)?; + if expected_identity != &actual_identity { + return None; + } + #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] + let size = metadata.len().min(i64::MAX as u64) as i64; + if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) + || usage.size != size + || usage.parsed_bytes.unwrap_or(0) < size + { + return None; + } + let last_totals = usage.last_totals.clone()?; + let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; + let child_fork_timestamp = child_fork_timestamp?; + JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) + .then_some(last_totals) + } +} + +impl CodexLineageDecision { + pub(super) fn accounting_mode( + &self, + matching_cached_state: Option<&CodexForkAccountingState>, + metadata: &CodexSessionMetadata, + paginated_continuation: bool, + ) -> CodexAccountingMode { + match self { + Self::Root => CodexAccountingMode::Standard, + Self::Unsafe => CodexAccountingMode::Unresolved, + Self::ParentReady(baseline) => { + let replaces_cached_state = matching_cached_state.is_some_and(|state| { + state.locally_resolved || state.inherited_totals.as_ref() != Some(baseline) + }); + let cached_parent_state = matching_cached_state.filter(|state| { + !state.locally_resolved && state.inherited_totals.as_ref() == Some(baseline) + }); + CodexAccountingMode::Baseline { + baseline: baseline.clone(), + paginated_continuation, + remaining_inherited_totals: cached_parent_state + .and_then(|state| state.remaining_inherited_totals.clone()), + provenance: CodexBaselineProvenance::ValidatedParent { + replaces_cached_state, + }, + } + } + Self::ParentAbsent => { + if let Some(state) = matching_cached_state + && let Some(baseline) = state.inherited_totals.clone() + { + return CodexAccountingMode::Baseline { + baseline, + paginated_continuation, + remaining_inherited_totals: state.remaining_inherited_totals.clone(), + provenance: if state.locally_resolved { + CodexBaselineProvenance::CachedLocalInference + } else { + CodexBaselineProvenance::CachedValidatedParent + }, + }; + } + if metadata.is_subagent { + CodexAccountingMode::InferSubagent { + start_ordinal: metadata.subagent_history_start_ordinal, + } + } else { + CodexAccountingMode::Unresolved + } + } + } + } +} + pub(super) fn cached_codex_file_is_fresh( cache: &CostUsageCache, entry: &CostUsageFileUsage, @@ -16,6 +232,12 @@ pub(super) fn cached_codex_file_is_fresh( && super::codex_fork_parent_is_safe(cache, entry) } +pub(super) fn codex_file_identity_matches(expected: Option<&str>, actual: Option<&str>) -> bool { + expected + .zip(actual) + .is_some_and(|(expected, actual)| expected == actual) +} + pub(super) fn cached_codex_file_is_complete_for_range( cache: &CostUsageCache, path_key: &str, @@ -26,14 +248,10 @@ pub(super) fn cached_codex_file_is_complete_for_range( let Ok(metadata) = fs::metadata(path_key) else { return false; }; - let identity_matches = match ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - (Some(_), None) => false, - (None, _) => true, - }; + let identity_matches = codex_file_identity_matches( + usage.codex_file_identity.as_deref(), + JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_deref(), + ); #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] let size = metadata.len().min(i64::MAX as u64) as i64; identity_matches @@ -91,7 +309,7 @@ struct CodexLineageNode { /// Order one bounded work set against both its admitted metadata and the /// persisted cache graph. The returned cache paths became structurally unsafe /// and must be invalidated even when the candidate limit deferred them. -pub(super) fn order_codex_candidates_by_lineage( +pub(super) fn plan_codex_candidates_by_lineage( cache: &CostUsageCache, candidates: &mut Vec, ) -> Vec { @@ -156,13 +374,19 @@ pub(super) fn order_codex_candidates_by_lineage( .or_default() .push(index); } - let mut unsafe_lineage = nodes + let mut lineage_gates = nodes .iter() - .map(|node| node.initially_unsafe) + .map(|node| { + if node.initially_unsafe { + CodexLineageGate::Unsafe + } else { + CodexLineageGate::Eligible + } + }) .collect::>(); for owners in session_owners.values().filter(|owners| owners.len() > 1) { for &index in owners { - unsafe_lineage[index] = true; + lineage_gates[index] = CodexLineageGate::Unsafe; } } let mut parent_indices = vec![None; nodes.len()]; @@ -173,8 +397,8 @@ pub(super) fn order_codex_candidates_by_lineage( match session_owners.get(parent_id).map(Vec::as_slice) { Some([parent_index]) => parent_indices[index] = Some(*parent_index), Some([]) | None if node.may_infer_missing_parent => {} - Some([]) | None => unsafe_lineage[index] = true, - Some(_) => unsafe_lineage[index] = true, + Some([]) | None => lineage_gates[index] = CodexLineageGate::Unsafe, + Some(_) => lineage_gates[index] = CodexLineageGate::Unsafe, } } @@ -184,12 +408,12 @@ pub(super) fn order_codex_candidates_by_lineage( loop { let mut progressed = false; for index in 0..nodes.len() { - if completed[index] || unsafe_lineage[index] { + if completed[index] || lineage_gates[index] == CodexLineageGate::Unsafe { continue; } let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { completed[parent_index] - && !unsafe_lineage[parent_index] + && lineage_gates[parent_index] == CodexLineageGate::Eligible && nodes[parent_index].may_author_parent }); if parent_is_ready { @@ -207,7 +431,7 @@ pub(super) fn order_codex_candidates_by_lineage( for index in 0..nodes.len() { if !completed[index] { - unsafe_lineage[index] = true; + lineage_gates[index] = CodexLineageGate::Unsafe; if let Some(candidate_index) = nodes[index].candidate_index { ordered_indices.push(candidate_index); } @@ -220,19 +444,15 @@ pub(super) fn order_codex_candidates_by_lineage( let mut candidate = remaining[candidate_index] .take() .expect("candidate is ordered once"); - candidate.lineage_disposition = if unsafe_lineage[node_index] { - CodexLineageDisposition::AmbiguousOrCyclic - } else { - CodexLineageDisposition::Ready - }; + candidate.lineage_gate = lineage_gates[node_index]; candidates.push(candidate); } nodes .iter() - .zip(unsafe_lineage) - .filter(|(node, unsafe_lineage)| { - *unsafe_lineage + .zip(lineage_gates) + .filter(|(node, gate)| { + *gate == CodexLineageGate::Unsafe && cache .files .get(&node.path) diff --git a/rust/src/cost_scanner/codex/reconciliation.rs b/rust/src/cost_scanner/codex/reconciliation.rs index b16d211b3b..8451c3f72b 100644 --- a/rust/src/cost_scanner/codex/reconciliation.rs +++ b/rust/src/cost_scanner/codex/reconciliation.rs @@ -108,14 +108,10 @@ fn codex_pending_path_affects_current_window( if codex_logical_target_has_unconsumed_tail(observed_size, usage) { return true; } - let identity_matches = match ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - (Some(_), None) => false, - (None, _) => true, - }; + let identity_matches = super::codex_file_identity_matches( + usage.codex_file_identity.as_deref(), + JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_deref(), + ); if !identity_matches || usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) || usage.size != observed_size @@ -270,6 +266,7 @@ mod tests { let old_usage = cache.files.get_mut(&old_key).unwrap(); old_usage.mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); old_usage.size = i64::try_from(metadata.len()).unwrap(); + old_usage.codex_file_identity = JsonlScanner::codex_file_identity(&old_path, &metadata); let range = active_range(); assert!(codex_current_window_is_established(&cache, &range)); @@ -279,6 +276,32 @@ mod tests { assert_eq!(report.sessions_count, 1); } + #[test] + fn historical_pending_entry_without_or_mismatched_identity_blocks_publication() { + let root = tempfile::tempdir().unwrap(); + let old_path = root.path().join("old.jsonl"); + let current_path = root.path().join("current.jsonl"); + std::fs::write(&old_path, vec![0_u8; 100]).unwrap(); + std::fs::write(¤t_path, vec![0_u8; 100]).unwrap(); + let old_key = old_path.to_string_lossy().into_owned(); + let current_key = current_path.to_string_lossy().into_owned(); + let metadata = std::fs::metadata(&old_path).unwrap(); + let range = active_range(); + + for cached_identity in [None, Some("different-file".to_string())] { + let mut cache = historical_pending_cache(&old_key, ¤t_key); + let old_usage = cache.files.get_mut(&old_key).unwrap(); + old_usage.mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); + old_usage.size = i64::try_from(metadata.len()).unwrap(); + old_usage.codex_file_identity = cached_identity; + + assert!(codex_pending_path_affects_current_window( + &cache, &old_key, &range + )); + assert!(!codex_current_window_is_established(&cache, &range)); + } + } + #[test] fn metadata_failure_blocks_historical_pending_publication() { let old_path = r"C:\sessions\missing-old.jsonl"; diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 6da7fb5188..8b95e31f0d 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -229,7 +229,7 @@ pub(super) fn scan_codex_detailed_with_cache( session_metadata: JsonlScanner::read_codex_session_metadata(&candidate.path) .unwrap_or_default(), path: candidate.path, - lineage_disposition: CodexLineageDisposition::Ready, + lineage_gate: CodexLineageGate::Eligible, }); } let mut unprocessed = Vec::new(); @@ -237,7 +237,7 @@ pub(super) fn scan_codex_detailed_with_cache( unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); unprocessed.extend(cancelled_during_preparation); } else { - let unsafe_cached_paths = order_codex_candidates_by_lineage(&cache, &mut work_queue); + let unsafe_cached_paths = plan_codex_candidates_by_lineage(&cache, &mut work_queue); invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); if invalidated_unsafe_lineage { cache.previous_report = None; diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs index 18c742ba72..8ff35834d3 100644 --- a/rust/src/cost_scanner/tests/lineage_cache.rs +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -63,6 +63,52 @@ fn lineage_token_row( }) } +fn write_missing_ordinal_subagent( + sessions_root: &Path, + name: &str, + base: DateTime, + include_owned_usage: bool, +) -> PathBuf { + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let mut missing_ordinal = lineage_token_row(base, 11, 100, 0); + missing_ordinal.as_object_mut().unwrap().remove("ordinal"); + let tail = if include_owned_usage { + lineage_token_row(base, 12, 120, 10) + } else { + lineage_token_row(base, 12, 100, 0) + }; + let rows = [ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), + "payload": { + "id": "child-id", + "forked_from_id": "absent-parent-id", + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": "absent-parent-id"}}} + } + }), + lineage_token_row(base, 9, 100, 0), + lineage_token_row(base, 10, 100, 0), + missing_ordinal, + tail, + ]; + let body = rows + .into_iter() + .map(|row| row.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + fn bounded_scanner(sessions: &Path, cache_root: &Path) -> CostScanner { let mut options = CostScanOptions::app_driven(); options.codex_candidate_limit = 1; @@ -91,6 +137,155 @@ fn assert_unresolved(cache: &CostUsageCache, path: &Path) { assert!(usage.codex_fork_accounting_state.is_none()); } +#[test] +fn replaced_parent_with_same_path_size_and_mtime_cannot_author_lineage() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_subagent( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (_, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let parent_key = parent.to_string_lossy().to_string(); + let child_usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(matches!( + CodexLineagePlanner::new(&cache).decision_for_usage(child_usage), + CodexLineageDecision::ParentReady(_) + )); + + let old_identity = cache.files[&parent_key] + .codex_file_identity + .clone() + .expect("parent identity persisted"); + let old_metadata = std::fs::metadata(&parent).unwrap(); + let old_mtime = old_metadata.modified().unwrap(); + let old_size = old_metadata.len(); + let rotated = parent.with_extension("old"); + std::fs::rename(&parent, &rotated).unwrap(); + let replacement = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[2_000], + ); + std::fs::OpenOptions::new() + .write(true) + .open(&replacement) + .unwrap() + .set_modified(old_mtime) + .unwrap(); + let replacement_metadata = std::fs::metadata(&replacement).unwrap(); + assert_eq!(replacement_metadata.len(), old_size); + let replacement_identity = + JsonlScanner::codex_file_identity(&replacement, &replacement_metadata) + .expect("replacement identity available"); + assert_ne!(replacement_identity, old_identity); + + assert_eq!( + CodexLineagePlanner::new(&cache).decision_for_usage(child_usage), + CodexLineageDecision::Unsafe + ); +} + +#[test] +fn missing_explicit_ordinal_keeps_subagent_cache_unresolved() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_missing_ordinal_subagent(&sessions, "child.jsonl", base, true); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + +#[test] +fn missing_ordinal_cannot_complete_zero_usage_subagent_cache() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_missing_ordinal_subagent(&sessions, "child.jsonl", base, false); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + +#[test] +fn legacy_cache_without_file_identity_is_reparsed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let session = write_codex_fork_session_fixture( + &sessions, + "session.jsonl", + "root-session-id", + None, + base, + base, + &[1_000], + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, _, _) = scanner.scan_codex_detailed_with_cache(None); + let session_key = session.to_string_lossy().to_string(); + let mut legacy_cache = JsonlScanner::load_cache(ProviderId::Codex, Some(&cache_root)); + legacy_cache + .files + .get_mut(&session_key) + .unwrap() + .codex_file_identity = None; + JsonlScanner::save_cache(ProviderId::Codex, &mut legacy_cache, Some(&cache_root)); + + let (_, stats, refreshed_cache) = scanner.scan_codex_detailed_with_cache(None); + + assert!(stats.codex_history_read_paths.contains(&session_key)); + assert!( + refreshed_cache.files[&session_key] + .codex_file_identity + .is_some() + ); +} + #[test] fn bounded_refresh_detects_duplicate_parent_owners_across_cache_and_candidate() { let root = tempfile::tempdir().unwrap(); From a7f685ae4939d3253e1f0473c4fa260971be62d3 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 05:56:21 +0700 Subject: [PATCH 36/62] Consolidate Codex lineage validation graph --- rust/src/cost_scanner/codex.rs | 4 + rust/src/cost_scanner/codex/logical_target.rs | 448 ++++++++++-------- rust/src/cost_scanner/codex/scan.rs | 4 +- 3 files changed, 245 insertions(+), 211 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 7e522c9da9..a64eb92a39 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -143,6 +143,7 @@ struct CodexPreparedCandidate { path: PathBuf, session_metadata: CodexSessionMetadata, lineage_gate: CodexLineageGate, + parent_owner_expected: bool, } #[derive(Debug, Clone, Copy, Default)] @@ -499,11 +500,14 @@ impl CostScanner { let lineage_gate = prepared_candidate .map(|candidate| candidate.lineage_gate) .unwrap_or_default(); + let parent_owner_expected = + prepared_candidate.is_some_and(|candidate| candidate.parent_owner_expected); let lineage_decision = CodexLineagePlanner::new(cache).decision_for_scan( is_fork, lineage_gate, codex_forked_from_id.as_deref(), codex_fork_timestamp.as_deref(), + parent_owner_expected, ); let accounting_mode = lineage_decision.accounting_mode( matching_cached_fork_state, diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 6245ea6c3b..8382c0a098 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -15,13 +15,227 @@ pub(super) enum CodexLineageDecision { Unsafe, } +struct CodexLineageNode { + path: String, + session_id: Option, + parent_id: Option, + candidate_index: Option, + may_infer_missing_parent: bool, + may_author_parent: bool, + initially_unsafe: bool, +} + +struct CodexLineageGraph { + nodes: Vec, + session_owners: HashMap>, + parent_indices: Vec>, + gates: Vec, + candidate_node_indices: Vec, + ordered_candidate_indices: Vec, +} + +impl CodexLineageGraph { + fn new(cache: &CostUsageCache, candidates: Option<&[CodexPreparedCandidate]>) -> Self { + let candidate_paths = candidates + .into_iter() + .flatten() + .map(|candidate| candidate.path.to_string_lossy().to_string()) + .collect::>(); + let mut cached_paths = cache + .files + .keys() + .filter(|path| !candidate_paths.contains(*path)) + .cloned() + .collect::>(); + cached_paths.sort(); + + let mut nodes = Vec::with_capacity(cached_paths.len() + candidate_paths.len()); + for path in cached_paths { + let usage = &cache.files[&path]; + let uses_parent = codex_usage_uses_parent(usage); + let locally_inferred = codex_fork_uses_local_inference(usage); + nodes.push(CodexLineageNode { + path, + session_id: usage.codex_session_id.clone(), + parent_id: uses_parent + .then(|| usage.codex_forked_from_id.clone()) + .flatten(), + candidate_index: None, + may_infer_missing_parent: locally_inferred, + may_author_parent: !locally_inferred && !usage.codex_unresolved_fork_parent, + initially_unsafe: usage.codex_unresolved_fork_parent, + }); + } + + let mut candidate_node_indices = Vec::new(); + if let Some(candidates) = candidates { + candidate_node_indices.reserve(candidates.len()); + for (candidate_index, candidate) in candidates.iter().enumerate() { + let uses_parent = candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some(); + nodes.push(CodexLineageNode { + path: candidate.path.to_string_lossy().to_string(), + session_id: candidate.session_metadata.session_id.clone(), + parent_id: uses_parent + .then(|| candidate.session_metadata.forked_from_id.clone()) + .flatten(), + candidate_index: Some(candidate_index), + may_infer_missing_parent: candidate.session_metadata.is_subagent, + may_author_parent: true, + initially_unsafe: false, + }); + candidate_node_indices.push(nodes.len() - 1); + } + } + + let mut session_owners = HashMap::>::new(); + for (index, node) in nodes.iter().enumerate() { + if let Some(session_id) = node.session_id.as_ref() { + session_owners + .entry(session_id.clone()) + .or_default() + .push(index); + } + } + + let mut gates = nodes + .iter() + .map(|node| { + if node.initially_unsafe { + CodexLineageGate::Unsafe + } else { + CodexLineageGate::Eligible + } + }) + .collect::>(); + for owners in session_owners.values().filter(|owners| owners.len() > 1) { + for &index in owners { + gates[index] = CodexLineageGate::Unsafe; + } + } + + let mut parent_indices = vec![None; nodes.len()]; + for (index, node) in nodes.iter().enumerate() { + let Some(parent_id) = node.parent_id.as_ref() else { + continue; + }; + match session_owners.get(parent_id) { + Some(owners) if owners.len() == 1 => parent_indices[index] = Some(owners[0]), + Some(_) => { + gates[index] = CodexLineageGate::Unsafe; + } + None if !node.may_infer_missing_parent => { + gates[index] = CodexLineageGate::Unsafe; + } + None => {} + } + } + + // This single topological pass both rejects cycles/unsafe ancestry and + // orders candidates. Cached-parent validation consumes the same gates. + let mut completed = vec![false; nodes.len()]; + let mut ordered_candidate_indices = Vec::with_capacity(candidate_node_indices.len()); + loop { + let mut progressed = false; + for index in 0..nodes.len() { + if completed[index] || gates[index] == CodexLineageGate::Unsafe { + continue; + } + let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { + completed[parent_index] + && gates[parent_index] == CodexLineageGate::Eligible + && nodes[parent_index].may_author_parent + }); + if parent_is_ready { + completed[index] = true; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_candidate_indices.push(candidate_index); + } + progressed = true; + } + } + if !progressed { + break; + } + } + + for index in 0..nodes.len() { + if !completed[index] { + gates[index] = CodexLineageGate::Unsafe; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_candidate_indices.push(candidate_index); + } + } + } + + Self { + nodes, + session_owners, + parent_indices, + gates, + candidate_node_indices, + ordered_candidate_indices, + } + } + + fn unique_owner(&self, session_id: &str) -> Result, ()> { + match self.session_owners.get(session_id).map(Vec::as_slice) { + None | Some([]) => Ok(None), + Some([index]) => Ok(Some(*index)), + Some(_) => Err(()), + } + } + + fn apply_candidate_plan(&self, candidates: &mut Vec) -> Vec { + if candidates.is_empty() { + return Vec::new(); + } + for (candidate_index, candidate) in candidates.iter_mut().enumerate() { + let node_index = self.candidate_node_indices[candidate_index]; + candidate.lineage_gate = self.gates[node_index]; + candidate.parent_owner_expected = self.parent_indices[node_index].is_some(); + } + + let mut remaining = candidates.drain(..).map(Some).collect::>(); + for candidate_index in &self.ordered_candidate_indices { + candidates.push( + remaining[*candidate_index] + .take() + .expect("candidate is ordered once"), + ); + } + + self.nodes + .iter() + .zip(&self.gates) + .filter(|(node, gate)| { + node.candidate_index.is_none() + && **gate == CodexLineageGate::Unsafe + && !node.initially_unsafe + }) + .map(|(node, _)| node.path.clone()) + .collect() + } +} + pub(super) struct CodexLineagePlanner<'a> { cache: &'a CostUsageCache, + graph: CodexLineageGraph, } impl<'a> CodexLineagePlanner<'a> { pub(super) fn new(cache: &'a CostUsageCache) -> Self { - Self { cache } + Self { + cache, + graph: CodexLineageGraph::new(cache, None), + } + } + + pub(super) fn plan_candidates_by_lineage( + cache: &CostUsageCache, + candidates: &mut Vec, + ) -> Vec { + CodexLineageGraph::new(cache, Some(candidates)).apply_candidate_plan(candidates) } pub(super) fn decision_for_scan( @@ -30,6 +244,7 @@ impl<'a> CodexLineagePlanner<'a> { gate: CodexLineageGate, parent_id: Option<&str>, fork_timestamp: Option<&str>, + parent_owner_expected: bool, ) -> CodexLineageDecision { if gate == CodexLineageGate::Unsafe { return CodexLineageDecision::Unsafe; @@ -38,7 +253,7 @@ impl<'a> CodexLineagePlanner<'a> { return CodexLineageDecision::Root; } parent_id.map_or(CodexLineageDecision::Unsafe, |parent_id| { - self.resolve_parent(parent_id, fork_timestamp) + self.resolve_parent(parent_id, fork_timestamp, parent_owner_expected) }) } @@ -53,7 +268,7 @@ impl<'a> CodexLineagePlanner<'a> { .codex_forked_from_id .as_deref() .map_or(CodexLineageDecision::Unsafe, |parent_id| { - self.resolve_parent(parent_id, usage.codex_fork_timestamp.as_deref()) + self.resolve_parent(parent_id, usage.codex_fork_timestamp.as_deref(), false) }) } @@ -74,45 +289,30 @@ impl<'a> CodexLineagePlanner<'a> { &self, parent_session_id: &str, child_fork_timestamp: Option<&str>, + parent_owner_expected: bool, ) -> CodexLineageDecision { - self.resolve_parent_inner(parent_session_id, child_fork_timestamp, &mut HashSet::new()) - } - - fn resolve_parent_inner( - &self, - parent_session_id: &str, - child_fork_timestamp: Option<&str>, - visiting: &mut HashSet, - ) -> CodexLineageDecision { - let mut owners = self - .cache - .files - .iter() - .filter(|(_, usage)| usage.codex_session_id.as_deref() == Some(parent_session_id)); - let Some((path_key, usage)) = owners.next() else { - return CodexLineageDecision::ParentAbsent; + let node_index = match self.graph.unique_owner(parent_session_id) { + Ok(None) if !parent_owner_expected => return CodexLineageDecision::ParentAbsent, + Ok(None) | Err(()) => return CodexLineageDecision::Unsafe, + Ok(Some(index)) => index, }; - if owners.next().is_some() || !visiting.insert(path_key.clone()) { - return CodexLineageDecision::Unsafe; - } - - let decision = self - .parent_owner_baseline(path_key, usage, child_fork_timestamp, visiting) + self.parent_owner_baseline(node_index, child_fork_timestamp) .map_or( CodexLineageDecision::Unsafe, CodexLineageDecision::ParentReady, - ); - visiting.remove(path_key); - decision + ) } fn parent_owner_baseline( &self, - path_key: &str, - usage: &CostUsageFileUsage, + node_index: usize, child_fork_timestamp: Option<&str>, - visiting: &mut HashSet, ) -> Option { + let node = self.graph.nodes.get(node_index)?; + if self.graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { + return None; + } + let usage = self.cache.files.get(&node.path)?; if usage.codex_unresolved_fork_parent || usage.codex_token_timestamps_monotonic != Some(true) || super::codex_fork_uses_local_inference(usage) @@ -121,28 +321,22 @@ impl<'a> CodexLineagePlanner<'a> { } if super::codex_usage_uses_parent(usage) { - let parent_id = usage.codex_forked_from_id.as_deref()?; let inherited = usage .codex_fork_accounting_state .as_ref()? .inherited_totals .as_ref()?; - match self.resolve_parent_inner( - parent_id, - usage.codex_fork_timestamp.as_deref(), - visiting, - ) { - CodexLineageDecision::ParentReady(baseline) if &baseline == inherited => {} - CodexLineageDecision::Root - | CodexLineageDecision::ParentAbsent - | CodexLineageDecision::ParentReady(_) - | CodexLineageDecision::Unsafe => return None, + let parent_index = self.graph.parent_indices[node_index]?; + let baseline = + self.parent_owner_baseline(parent_index, usage.codex_fork_timestamp.as_deref())?; + if &baseline != inherited { + return None; } } - let metadata = fs::metadata(path_key).ok()?; + let metadata = fs::metadata(&node.path).ok()?; let expected_identity = usage.codex_file_identity.as_ref()?; - let actual_identity = JsonlScanner::codex_file_identity(Path::new(path_key), &metadata)?; + let actual_identity = JsonlScanner::codex_file_identity(Path::new(&node.path), &metadata)?; if expected_identity != &actual_identity { return None; } @@ -296,172 +490,6 @@ pub(super) fn defer_codex_locally_inferred_candidates( candidates.extend(other); } -struct CodexLineageNode { - path: String, - session_id: Option, - parent_id: Option, - candidate_index: Option, - may_infer_missing_parent: bool, - may_author_parent: bool, - initially_unsafe: bool, -} - -/// Order one bounded work set against both its admitted metadata and the -/// persisted cache graph. The returned cache paths became structurally unsafe -/// and must be invalidated even when the candidate limit deferred them. -pub(super) fn plan_codex_candidates_by_lineage( - cache: &CostUsageCache, - candidates: &mut Vec, -) -> Vec { - if candidates.is_empty() { - return Vec::new(); - } - - let candidate_paths = candidates - .iter() - .map(|candidate| candidate.path.to_string_lossy().to_string()) - .collect::>(); - let mut cached_paths = cache - .files - .keys() - .filter(|path| !candidate_paths.contains(*path)) - .cloned() - .collect::>(); - cached_paths.sort(); - - let mut nodes = Vec::with_capacity(cached_paths.len() + candidates.len()); - for path in cached_paths { - let usage = &cache.files[&path]; - let uses_parent = super::codex_usage_uses_parent(usage); - let locally_inferred = super::codex_fork_uses_local_inference(usage); - nodes.push(CodexLineageNode { - path, - session_id: usage.codex_session_id.clone(), - parent_id: uses_parent - .then(|| usage.codex_forked_from_id.clone()) - .flatten(), - candidate_index: None, - may_infer_missing_parent: locally_inferred, - may_author_parent: !locally_inferred && !usage.codex_unresolved_fork_parent, - initially_unsafe: usage.codex_unresolved_fork_parent, - }); - } - let mut candidate_node_indices = Vec::with_capacity(candidates.len()); - for (candidate_index, candidate) in candidates.iter().enumerate() { - let uses_parent = candidate.session_metadata.lineage.uses_parent_baseline() - || candidate.session_metadata.forked_from_id.is_some(); - nodes.push(CodexLineageNode { - path: candidate.path.to_string_lossy().to_string(), - session_id: candidate.session_metadata.session_id.clone(), - parent_id: uses_parent - .then(|| candidate.session_metadata.forked_from_id.clone()) - .flatten(), - candidate_index: Some(candidate_index), - may_infer_missing_parent: candidate.session_metadata.is_subagent, - may_author_parent: true, - initially_unsafe: false, - }); - candidate_node_indices.push(nodes.len() - 1); - } - - let mut session_owners = HashMap::>::new(); - for (index, node) in nodes.iter().enumerate() { - let Some(session_id) = node.session_id.as_ref() else { - continue; - }; - session_owners - .entry(session_id.clone()) - .or_default() - .push(index); - } - let mut lineage_gates = nodes - .iter() - .map(|node| { - if node.initially_unsafe { - CodexLineageGate::Unsafe - } else { - CodexLineageGate::Eligible - } - }) - .collect::>(); - for owners in session_owners.values().filter(|owners| owners.len() > 1) { - for &index in owners { - lineage_gates[index] = CodexLineageGate::Unsafe; - } - } - let mut parent_indices = vec![None; nodes.len()]; - for (index, node) in nodes.iter().enumerate() { - let Some(parent_id) = node.parent_id.as_ref() else { - continue; - }; - match session_owners.get(parent_id).map(Vec::as_slice) { - Some([parent_index]) => parent_indices[index] = Some(*parent_index), - Some([]) | None if node.may_infer_missing_parent => {} - Some([]) | None => lineage_gates[index] = CodexLineageGate::Unsafe, - Some(_) => lineage_gates[index] = CodexLineageGate::Unsafe, - } - } - - let mut completed = vec![false; nodes.len()]; - let mut ordered_indices = Vec::with_capacity(candidates.len()); - - loop { - let mut progressed = false; - for index in 0..nodes.len() { - if completed[index] || lineage_gates[index] == CodexLineageGate::Unsafe { - continue; - } - let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { - completed[parent_index] - && lineage_gates[parent_index] == CodexLineageGate::Eligible - && nodes[parent_index].may_author_parent - }); - if parent_is_ready { - completed[index] = true; - if let Some(candidate_index) = nodes[index].candidate_index { - ordered_indices.push(candidate_index); - } - progressed = true; - } - } - if !progressed { - break; - } - } - - for index in 0..nodes.len() { - if !completed[index] { - lineage_gates[index] = CodexLineageGate::Unsafe; - if let Some(candidate_index) = nodes[index].candidate_index { - ordered_indices.push(candidate_index); - } - } - } - - let mut remaining = candidates.drain(..).map(Some).collect::>(); - for candidate_index in ordered_indices { - let node_index = candidate_node_indices[candidate_index]; - let mut candidate = remaining[candidate_index] - .take() - .expect("candidate is ordered once"); - candidate.lineage_gate = lineage_gates[node_index]; - candidates.push(candidate); - } - - nodes - .iter() - .zip(lineage_gates) - .filter(|(node, gate)| { - *gate == CodexLineageGate::Unsafe - && cache - .files - .get(&node.path) - .is_some_and(|usage| !usage.codex_unresolved_fork_parent) - }) - .map(|(node, _)| node.path.clone()) - .collect() -} - pub(super) fn invalidate_codex_unsafe_lineage(cache: &mut CostUsageCache, paths: &[String]) { for path in paths { let Some(usage) = cache.files.get_mut(path) else { diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 8b95e31f0d..3fb183eabb 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -230,6 +230,7 @@ pub(super) fn scan_codex_detailed_with_cache( .unwrap_or_default(), path: candidate.path, lineage_gate: CodexLineageGate::Eligible, + parent_owner_expected: false, }); } let mut unprocessed = Vec::new(); @@ -237,7 +238,8 @@ pub(super) fn scan_codex_detailed_with_cache( unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); unprocessed.extend(cancelled_during_preparation); } else { - let unsafe_cached_paths = plan_codex_candidates_by_lineage(&cache, &mut work_queue); + let unsafe_cached_paths = + CodexLineagePlanner::plan_candidates_by_lineage(&cache, &mut work_queue); invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); if invalidated_unsafe_lineage { cache.previous_report = None; From ac88694a92fd5aa1093b949252c916d30d06f3ae Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 06:19:26 +0700 Subject: [PATCH 37/62] Reject late ordinal gaps in inferred forks --- rust/src/core/jsonl_scanner/codex/parser.rs | 29 +++++++++++++++-- rust/src/core/jsonl_scanner/tests.rs | 34 ++++++++++++++++++++ rust/src/cost_scanner/tests/lineage_cache.rs | 32 ++++++++++++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index 728bd16f76..bdce9c68d8 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -92,6 +92,11 @@ impl ForkBaselineInference { } } + fn mark_missing_explicit_ordinal(&mut self) { + self.missing_explicit_ordinal = true; + self.locally_confirmed = false; + } + fn observe_non_token(&mut self, obj: &Value) { if obj.get("type").and_then(Value::as_str) == Some("turn_context") && self.inherited_opening { @@ -117,8 +122,7 @@ impl ForkBaselineInference { let ordinal = obj.get("ordinal").and_then(Value::as_i64); if self.explicit_start_ordinal.is_some() && ordinal.is_none() { - self.missing_explicit_ordinal = true; - self.locally_confirmed = false; + self.mark_missing_explicit_ordinal(); if !self.boundary_open { self.baseline = Some(total); } @@ -338,6 +342,27 @@ impl CodexParserState { } let event_candidate = is_candidate_codex_line(line); + if event_candidate + && self + .fork_baseline_inference + .as_ref() + .is_some_and(|inference| { + inference.resolved && inference.explicit_start_ordinal.is_some() + }) + { + let Ok(obj) = serde_json::from_str::(line) else { + return; + }; + if token_count_payload(&obj).is_some() + && obj.get("ordinal").and_then(Value::as_i64).is_none() + { + self.fork_baseline_inference + .as_mut() + .expect("resolved inference exists") + .mark_missing_explicit_ordinal(); + return; + } + } let bare_candidate = !event_candidate && line.contains("\"usage\""); if !event_candidate && !bare_candidate { return; diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index 4abfce1c93..1aebd89e20 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -184,6 +184,40 @@ fn inferred_fork_keeps_missing_ordinal_unresolved_after_boundary_opens() { assert!(!state.fork_baseline_locally_resolved()); } +#[test] +fn inferred_fork_keeps_missing_ordinal_unresolved_after_local_resolution() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + let token_line = |ordinal: Option, total: i64, last: i64| { + let mut value = serde_json::json!({ + "timestamp": "2026-09-22T10:00:00Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": {"input_tokens": total, "cached_input_tokens": 0, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last, "cached_input_tokens": 0, "output_tokens": 0} + }} + }); + if let Some(ordinal) = ordinal { + value["ordinal"] = serde_json::json!(ordinal); + } + value.to_string() + }; + + state.process_line(&token_line(Some(9), 100, 0), &range); + state.process_line(&token_line(Some(10), 100, 0), &range); + state.process_line(&token_line(Some(11), 110, 10), &range); + assert!(state.fork_baseline_locally_resolved()); + state.process_line(&token_line(None, 120, 10), &range); + + assert!(!state.fork_baseline_locally_resolved()); +} + #[test] fn codex_token_pipeline_preserves_counts_above_i32_max() { let parsed = read_token_totals(&serde_json::json!({ diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs index 8ff35834d3..facfae31cf 100644 --- a/rust/src/cost_scanner/tests/lineage_cache.rs +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -246,6 +246,38 @@ fn missing_ordinal_cannot_complete_zero_usage_subagent_cache() { assert_unresolved(&cache, &child); } +#[test] +fn missing_ordinal_after_local_resolution_keeps_subagent_cache_unresolved() { + use std::io::Write as _; + + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent( + &sessions, + "child.jsonl", + "child-id", + "missing-parent-id", + base, + ); + let mut missing_ordinal = lineage_token_row(base + Duration::seconds(2), 21, 1_060, 10); + missing_ordinal.as_object_mut().unwrap().remove("ordinal"); + std::fs::OpenOptions::new() + .append(true) + .open(&child) + .unwrap() + .write_all(format!("{missing_ordinal}\n").as_bytes()) + .unwrap(); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + #[test] fn legacy_cache_without_file_identity_is_reparsed() { let root = tempfile::tempdir().unwrap(); From 042c88b7dfa26d01f15a96c4cd5876b74a96b842 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:41:11 +0700 Subject: [PATCH 38/62] Bound Antigravity local history scans --- rust/src/providers/antigravity/cost.rs | 50 ++ .../providers/antigravity/local_history.rs | 218 +++++++ .../providers/antigravity/local_sessions.rs | 588 +----------------- .../antigravity/local_sessions_reader.rs | 583 +++++++++++++++++ .../src/providers/antigravity/local_sqlite.rs | 91 ++- rust/src/providers/antigravity/mod.rs | 6 +- 6 files changed, 943 insertions(+), 593 deletions(-) create mode 100644 rust/src/providers/antigravity/cost.rs create mode 100644 rust/src/providers/antigravity/local_history.rs create mode 100644 rust/src/providers/antigravity/local_sessions_reader.rs diff --git a/rust/src/providers/antigravity/cost.rs b/rust/src/providers/antigravity/cost.rs new file mode 100644 index 0000000000..f8eb38f24a --- /dev/null +++ b/rust/src/providers/antigravity/cost.rs @@ -0,0 +1,50 @@ +use crate::core::CostUsagePricing; + +pub(super) fn estimate_cost_usd( + model: Option<&str>, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> Option { + let model = model.map(str::trim).filter(|value| !value.is_empty())?; + let input = i32::try_from(input).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let output = i32::try_from(output).ok()?; + let resolve = |candidate: &str| { + CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + }; + resolve(model).or_else(|| { + ["-tiered", "-low", "-thinking"] + .iter() + .find_map(|suffix| model.strip_suffix(suffix)) + .filter(|base| !base.is_empty()) + .and_then(resolve) + }) +} + +#[cfg(test)] +mod tests { + use super::estimate_cost_usd; + + #[test] + fn prices_known_models_and_provider_local_routing_variants() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) + .expect("routing suffix uses the base public price"); + assert!(direct > 0.0); + assert_eq!(direct, routed); + } + + #[test] + fn unknown_or_oversized_pricing_inputs_fail_closed() { + assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), + None + ); + } +} diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs new file mode 100644 index 0000000000..055f9b8d97 --- /dev/null +++ b/rust/src/providers/antigravity/local_history.rs @@ -0,0 +1,218 @@ +use super::{local_sessions_reader as local_sessions, local_sqlite}; +use std::fs; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; + +use crate::spend_contract::LocalTokenHistorySummary; + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +fn configured_database_roots(home: &Path) -> [PathBuf; 3] { + let gemini_cli_home = std::env::var("GEMINI_CLI_HOME").ok(); + configured_database_roots_from_values(home, gemini_cli_home.as_deref()) +} + +fn configured_database_roots_from_values( + home: &Path, + gemini_cli_home: Option<&str>, +) -> [PathBuf; 3] { + let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); + local_sqlite::database_roots(&gemini_base) +} + +fn summarize_local_usage_from( + roots: &[PathBuf], + now: DateTime, + days: u32, + jsonl_fallback: impl FnOnce() -> LocalTokenHistorySummary, +) -> LocalTokenHistorySummary { + match local_sqlite::summarize(roots, now, days) { + local_sqlite::SQLiteScan::Summary(summary) => summary, + local_sqlite::SQLiteScan::NoDatabases | local_sqlite::SQLiteScan::Unsupported => { + jsonl_fallback() + } + } +} + +pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { + let now = Utc::now(); + let Some(home) = dirs::home_dir() else { + return LocalTokenHistorySummary::default(); + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + summarize_local_usage_from(&roots, now, days, || { + local_sessions::summarize_jsonl_at(&tokscale_sessions, now, days) + }) +} + +/// Count local Antigravity conversation artifacts for the quota provider's +/// offline fallback. Mirrors upstream #3119 without opening SQLite files. +pub fn offline_conversation_count() -> usize { + let Some(home) = dirs::home_dir() else { + return 0; + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +fn offline_conversation_count_in(home: &Path) -> usize { + let roots = local_sqlite::database_roots(&home.join(".gemini")); + let tokscale_sessions = local_sessions::tokscale_sessions_from_values(home, None); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +fn offline_conversation_count_with_roots( + database_roots: &[PathBuf], + tokscale_sessions: &Path, +) -> usize { + let db_count = database_roots + .iter() + .map(|root| count_extension(root, "db")) + .sum::(); + if db_count > 0 { + return db_count; + } + local_sessions::count_jsonl_sessions_at(tokscale_sessions) +} + +fn count_extension(root: &Path, extension: &str) -> usize { + fs::read_dir(root) + .ok() + .into_iter() + .flatten() + .flatten() + .map(|entry| entry.path()) + .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) + .count() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spend_contract::LocalHistoryCoverage; + use chrono::TimeZone; + use rusqlite::Connection; + + #[test] + fn foreign_database_preserves_valid_tokscale_history() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute( + "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", + [], + ) + .unwrap(); + + let tokscale_sessions = dir + .path() + .join(".config/tokscale/antigravity-cache/sessions"); + fs::create_dir_all(&tokscale_sessions).unwrap(); + let session_path = tokscale_sessions.join("session-a.jsonl"); + fs::write( + &session_path, + b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", + ) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, || { + local_sessions::summarize_jsonl_paths( + std::slice::from_ref(&session_path), + now, + 7, + false, + ) + }); + + assert_eq!(summary.total_tokens, 120); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, LocalTokenHistorySummary::default); + + assert_eq!(summary, LocalTokenHistorySummary::default()); + assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); + } + + #[test] + fn scan_context_honors_non_empty_root_overrides() { + let home = Path::new(r"C:\Users\test"); + let tokscale_sessions = + local_sessions::tokscale_sessions_from_values(home, Some(r"E:\tokscale-root")); + let roots = configured_database_roots_from_values(home, Some(r"D:\gemini-root")); + assert_eq!( + roots[0], + PathBuf::from(r"D:\gemini-root") + .join("antigravity-cli") + .join("conversations") + ); + assert_eq!( + tokscale_sessions, + PathBuf::from(r"E:\tokscale-root") + .join("antigravity-cache") + .join("sessions") + ); + let defaults = local_sessions::tokscale_sessions_from_values(home, Some("")); + let default_roots = configured_database_roots_from_values(home, Some(" ")); + assert_eq!(default_roots[1], home.join(".gemini").join("antigravity")); + assert_eq!( + defaults, + home.join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions") + ); + } + + #[test] + fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { + let dir = tempfile::tempdir().unwrap(); + let app = dir + .path() + .join(".gemini") + .join("antigravity") + .join("conversations"); + fs::create_dir_all(&app).unwrap(); + fs::write(app.join("a.db"), b"").unwrap(); + fs::write(app.join("a.db-wal"), b"").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + + fs::remove_file(app.join("a.db")).unwrap(); + let cache = dir + .path() + .join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions"); + fs::create_dir_all(&cache).unwrap(); + fs::write(cache.join("one.jsonl"), b"{}\n").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + } +} diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 3191e88810..4efb9af063 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -1,584 +1,4 @@ -use std::collections::HashSet; -use std::fs::{self, File}; -use std::io::{BufRead, BufReader}; -use std::path::{Path, PathBuf}; - -use chrono::{DateTime, Duration, Local, TimeZone, Utc}; -use serde_json::Value; - -use crate::core::CostUsagePricing; - -const MAX_SESSION_FILES: usize = 2048; -const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; -const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; -const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; - -pub use crate::spend_contract::LocalHistoryCoverage; -pub type LocalSessionSummary = crate::spend_contract::LocalTokenHistorySummary; - -#[derive(Debug, Clone, PartialEq, Eq)] -struct ScanContext { - database_roots: [PathBuf; 3], - tokscale_sessions: PathBuf, -} - -impl ScanContext { - fn from_values( - home: &Path, - gemini_cli_home: Option<&str>, - tokscale_config_dir: Option<&str>, - ) -> Self { - let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); - let tokscale_base = clean_env_path(tokscale_config_dir) - .unwrap_or_else(|| home.join(".config").join("tokscale")); - Self { - database_roots: super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: tokscale_base.join("antigravity-cache").join("sessions"), - } - } - - fn capture() -> Option { - let home = dirs::home_dir()?; - let gemini = std::env::var("GEMINI_CLI_HOME").ok(); - let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); - Some(Self::from_values( - &home, - gemini.as_deref(), - tokscale.as_deref(), - )) - } -} - -fn clean_env_path(value: Option<&str>) -> Option { - value - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) -} - -pub fn summarize(days: u32) -> LocalSessionSummary { - let now = Utc::now(); - let Some(context) = ScanContext::capture() else { - return LocalSessionSummary::default(); - }; - summarize_context(&context, now, days) -} - -fn summarize_context(context: &ScanContext, now: DateTime, days: u32) -> LocalSessionSummary { - match super::local_sqlite::summarize(&context.database_roots, now, days) { - super::local_sqlite::SQLiteScan::Summary(summary) => summary, - super::local_sqlite::SQLiteScan::NoDatabases - | super::local_sqlite::SQLiteScan::Unsupported => { - let (paths, truncated) = tokscale_paths(&context.tokscale_sessions); - if paths.is_empty() { - LocalSessionSummary::default() - } else { - summarize_paths(&paths, now, days, truncated) - } - } - } -} - -/// Count local Antigravity conversation artifacts for the quota provider's -/// offline fallback. Mirrors upstream #3119 without opening SQLite files. -pub fn offline_conversation_count() -> usize { - let Some(context) = ScanContext::capture() else { - return 0; - }; - offline_conversation_count_context(&context) -} - -fn offline_conversation_count_in(home: &Path) -> usize { - offline_conversation_count_context(&ScanContext::from_values(home, None, None)) -} - -fn offline_conversation_count_context(context: &ScanContext) -> usize { - let db_count = context - .database_roots - .iter() - .map(|root| count_extension(root, "db")) - .sum::(); - if db_count > 0 { - return db_count; - } - tokscale_paths(&context.tokscale_sessions).0.len() -} - -fn count_extension(root: &Path, extension: &str) -> usize { - fs::read_dir(root) - .ok() - .into_iter() - .flatten() - .flatten() - .map(|entry| entry.path()) - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) - .count() -} - -fn tokscale_paths(base: &Path) -> (Vec, bool) { - let Ok(entries) = fs::read_dir(base) else { - return (Vec::new(), false); - }; - let mut paths: Vec<_> = entries - .flatten() - .map(|entry| entry.path()) - .filter(|path| { - path.extension() - .and_then(|value| value.to_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")) - }) - .collect(); - paths.sort(); - let truncated = paths.len() > MAX_SESSION_FILES; - if truncated { - paths.drain(..paths.len() - MAX_SESSION_FILES); - } - (paths, truncated) -} - -fn summarize_paths( - paths: &[PathBuf], - now: DateTime, - days: u32, - truncated: bool, -) -> LocalSessionSummary { - let first_day = now.with_timezone(&Local).date_naive() - - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); - let mut total_tokens = 0_u64; - let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); - let mut sessions_with_usage = HashSet::new(); - let mut seen_response_ids = HashSet::new(); - let mut complete = !truncated; - - for path in paths.iter().take(MAX_SESSION_FILES) { - let file = match File::open(path) { - Ok(file) => file, - Err(_) => { - complete = false; - continue; - } - }; - match file.metadata() { - Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, - Ok(_) => {} - Err(_) => complete = false, - } - let mut reader = BufReader::new(file); - let mut remaining = MAX_SESSION_FILE_BYTES; - let mut path_had_usage = false; - let mut model = None::; - loop { - let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { - Ok(Some(line)) => line, - Ok(None) => break, - Err(_) => { - complete = false; - break; - } - }; - if line.is_empty() { - continue; - } - let Ok(value) = serde_json::from_slice::(&line) else { - continue; - }; - let kind = value.get("type").and_then(Value::as_str); - if kind == Some("session_meta") { - model = value - .get("modelId") - .or_else(|| value.get("model_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string); - continue; - } - if kind != Some("usage") && value.get("input").is_none() { - continue; - } - if let Some(response_id) = value - .get("responseId") - .or_else(|| value.get("response_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - && !seen_response_ids.insert(response_id.to_string()) - { - continue; - } - - let timestamp_ms = value - .get("timestamp") - .and_then(Value::as_i64) - .unwrap_or_default(); - let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { - continue; - }; - if at > now || at.with_timezone(&Local).date_naive() < first_day { - continue; - } - - let input = token_field(&value, &["input"]); - let output = token_field(&value, &["output"]); - let cache_read = token_field(&value, &["cacheRead", "cache_read"]); - let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); - let reasoning = token_field( - &value, - &["reasoning", "reasoningTokens", "reasoning_tokens"], - ); - let total = input - .saturating_add(output) - .saturating_add(cache_read) - .saturating_add(cache_write) - .saturating_add(reasoning); - if total == 0 { - continue; - } - total_tokens = total_tokens.saturating_add(total); - cost_estimate.record_list_price(estimate_cost_usd( - model.as_deref(), - input, - cache_read, - cache_write, - output.saturating_add(reasoning), - )); - path_had_usage = true; - } - if path_had_usage { - sessions_with_usage.insert(path.clone()); - } - } - - LocalSessionSummary { - total_tokens, - session_count: sessions_with_usage.len(), - coverage: if paths.is_empty() { - LocalHistoryCoverage::Unavailable - } else if complete { - LocalHistoryCoverage::Complete - } else { - LocalHistoryCoverage::Partial - }, - cost_estimate, - } -} - -pub(super) fn estimate_cost_usd( - model: Option<&str>, - input: u64, - cache_read: u64, - cache_write: u64, - output: u64, -) -> Option { - let model = model.map(str::trim).filter(|value| !value.is_empty())?; - let input = i32::try_from(input).ok()?; - let cache_read = i32::try_from(cache_read).ok()?; - let cache_write = i32::try_from(cache_write).ok()?; - let output = i32::try_from(output).ok()?; - let resolve = |candidate: &str| { - CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) - .filter(|cost| cost.is_finite() && *cost >= 0.0) - }; - resolve(model).or_else(|| { - ["-tiered", "-low", "-thinking"] - .iter() - .find_map(|suffix| model.strip_suffix(suffix)) - .filter(|base| !base.is_empty()) - .and_then(resolve) - }) -} - -fn read_bounded_jsonl_line( - reader: &mut R, - remaining_file_bytes: &mut usize, -) -> std::io::Result>> { - if *remaining_file_bytes == 0 { - return Ok(None); - } - let mut line = Vec::new(); - let mut saw_input = false; - let mut discarding = false; - - loop { - let chunk = reader.fill_buf()?; - if chunk.is_empty() { - return Ok(saw_input.then_some(if discarding { Vec::new() } else { line })); - } - let bounded_len = chunk.len().min(*remaining_file_bytes); - if bounded_len == 0 { - return Ok(None); - } - let bounded = &chunk[..bounded_len]; - let newline = bounded.iter().position(|byte| *byte == b'\n'); - let segment_end = newline.unwrap_or(bounded.len()); - let segment = &bounded[..segment_end]; - saw_input = saw_input || !segment.is_empty() || newline.is_some(); - if !discarding { - if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { - line.extend_from_slice(segment); - } else { - line.clear(); - discarding = true; - } - } - let consumed = segment_end + usize::from(newline.is_some()); - reader.consume(consumed); - *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); - if newline.is_some() { - return Ok(Some(if discarding { Vec::new() } else { line })); - } - if *remaining_file_bytes == 0 { - return Ok(Some(Vec::new())); - } - } -} - -fn token_field(value: &Value, keys: &[&str]) -> u64 { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_u64)) - .unwrap_or(0) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn prices_known_models_and_provider_local_routing_variants() { - let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) - .expect("known public price"); - let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) - .expect("routing suffix uses the base public price"); - assert!(direct > 0.0); - assert_eq!(direct, routed); - } - - #[test] - fn unknown_or_oversized_pricing_inputs_fail_closed() { - assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); - assert_eq!( - estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), - None - ); - } - - #[test] - fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { - let dir = tempfile::tempdir().unwrap(); - let known = dir.path().join("known.jsonl"); - let unknown = dir.path().join("unknown.jsonl"); - fs::write( - &known, - concat!( - "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", - "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" - ), - ) - .unwrap(); - fs::write( - &unknown, - concat!( - "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", - "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_paths(&[known, unknown], now, 7, false); - - assert_eq!(summary.cost_estimate.coverage.estimated, 1); - assert_eq!(summary.cost_estimate.coverage.unpriced, 1); - assert!(summary.cost_estimate.known_subtotal_usd.is_some()); - assert_eq!(summary.total_usd(), None); - } - use rusqlite::Connection; - - #[test] - fn scan_context_honors_non_empty_root_overrides() { - let home = Path::new(r"C:\Users\test"); - let context = - ScanContext::from_values(home, Some(r"D:\gemini-root"), Some(r"E:\tokscale-root")); - assert_eq!( - context.database_roots[0], - PathBuf::from(r"D:\gemini-root") - .join("antigravity-cli") - .join("conversations") - ); - assert_eq!( - context.tokscale_sessions, - PathBuf::from(r"E:\tokscale-root") - .join("antigravity-cache") - .join("sessions") - ); - let defaults = ScanContext::from_values(home, Some(" "), Some("")); - assert_eq!( - defaults.database_roots[1], - home.join(".gemini").join("antigravity") - ); - assert_eq!( - defaults.tokscale_sessions, - home.join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions") - ); - } - - #[test] - fn foreign_database_preserves_valid_tokscale_history() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute( - "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", - [], - ) - .unwrap(); - - let tokscale_sessions = dir - .path() - .join(".config/tokscale/antigravity-cache/sessions"); - fs::create_dir_all(&tokscale_sessions).unwrap(); - fs::write( - tokscale_sessions.join("session-a.jsonl"), - b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", - ) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions, - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary.total_tokens, 120); - assert_eq!(summary.session_count, 1); - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - } - - #[test] - fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: dir.path().join("missing-tokscale-sessions"), - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary, LocalSessionSummary::default()); - assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); - } - - #[test] - fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write(&path, concat!( - "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", - "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", - "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" - )).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 135); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn truncated_or_unreadable_tokscale_history_is_partial() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); - assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); - - let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); - assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); - } - #[test] - fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { - let dir = tempfile::tempdir().unwrap(); - let app = dir - .path() - .join(".gemini") - .join("antigravity") - .join("conversations"); - fs::create_dir_all(&app).unwrap(); - fs::write(app.join("a.db"), b"").unwrap(); - fs::write(app.join("a.db-wal"), b"").unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - - fs::remove_file(app.join("a.db")).unwrap(); - let cache = dir - .path() - .join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions"); - fs::create_dir_all(&cache).unwrap(); - fs::write( - cache.join("one.jsonl"), - b"{} -", - ) - .unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - } - - #[test] - fn excludes_usage_outside_requested_window() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", - "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn oversized_jsonl_line_is_discarded_and_next_usage_row_is_counted() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - let mut text = format!( - r#"{{"type":"usage","padding":"{}"}}"#, - "x".repeat(MAX_JSONL_LINE_BYTES + 32) - ); - text.push('\n'); - text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); - text.push('\n'); - fs::write(&path, text).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } -} +pub use super::local_history::{offline_conversation_count, summarize_local_usage as summarize}; +pub use crate::spend_contract::{ + LocalHistoryCoverage, LocalTokenHistorySummary as LocalSessionSummary, +}; diff --git a/rust/src/providers/antigravity/local_sessions_reader.rs b/rust/src/providers/antigravity/local_sessions_reader.rs new file mode 100644 index 0000000000..cbc3297cb4 --- /dev/null +++ b/rust/src/providers/antigravity/local_sessions_reader.rs @@ -0,0 +1,583 @@ +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashSet}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Duration, Local, TimeZone, Utc}; +use serde_json::Value; + +use super::cost::estimate_cost_usd; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + +const MAX_SESSION_FILES: usize = 2048; +const MAX_SESSION_DISCOVERY_ENTRIES: usize = 16 * 1024; +const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; +const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; +const MAX_TOTAL_SESSION_BYTES: usize = 128 * 1024 * 1024; +const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; + +enum BoundedJsonlLine { + Record(Vec), + Oversized, + Truncated, +} + +pub(super) fn tokscale_sessions_from_values( + home: &Path, + tokscale_config_dir: Option<&str>, +) -> PathBuf { + let tokscale_base = clean_env_path(tokscale_config_dir) + .unwrap_or_else(|| home.join(".config").join("tokscale")); + tokscale_base.join("antigravity-cache").join("sessions") +} + +pub(super) fn configured_tokscale_sessions(home: &Path) -> PathBuf { + let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); + tokscale_sessions_from_values(home, tokscale.as_deref()) +} + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +pub(super) fn summarize_jsonl_at( + tokscale_sessions: &Path, + now: DateTime, + days: u32, +) -> LocalTokenHistorySummary { + let (paths, truncated) = tokscale_paths(tokscale_sessions); + summarize_jsonl_paths(&paths, now, days, truncated) +} + +pub(super) fn summarize_jsonl_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + if paths.is_empty() { + LocalTokenHistorySummary::default() + } else { + summarize_paths(paths, now, days, truncated) + } +} + +fn tokscale_paths(base: &Path) -> (Vec, bool) { + let Ok(entries) = fs::read_dir(base) else { + return (Vec::new(), false); + }; + bounded_tokscale_paths(entries, MAX_SESSION_DISCOVERY_ENTRIES, MAX_SESSION_FILES) +} + +fn bounded_tokscale_paths( + entries: fs::ReadDir, + max_entries: usize, + max_files: usize, +) -> (Vec, bool) { + let mut paths = BinaryHeap::>::with_capacity(max_files); + let mut truncated = false; + + for (entries_examined, entry) in entries.enumerate() { + if entries_examined == max_entries { + truncated = true; + break; + } + + let Ok(entry) = entry else { + truncated = true; + continue; + }; + let path = entry.path(); + let is_jsonl = path + .extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")); + if !is_jsonl { + continue; + } + + if paths.len() < max_files { + paths.push(Reverse(path)); + } else { + truncated = true; + if let Some(Reverse(smallest)) = paths.peek() + && path > *smallest + { + paths.pop(); + paths.push(Reverse(path)); + } + } + } + + let mut paths: Vec<_> = paths.into_iter().map(|Reverse(path)| path).collect(); + paths.sort(); + (paths, truncated) +} + +pub(super) fn count_jsonl_sessions_at(base: &Path) -> usize { + tokscale_paths(base).0.len() +} + +fn summarize_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + summarize_paths_with_budget(paths, now, days, truncated, MAX_TOTAL_SESSION_BYTES) +} + +fn summarize_paths_with_budget( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, + total_byte_budget: usize, +) -> LocalTokenHistorySummary { + let first_day = now.with_timezone(&Local).date_naive() + - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + let mut total_tokens = 0_u64; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); + let mut sessions_with_usage = HashSet::new(); + let mut seen_response_ids = HashSet::new(); + let mut complete = !truncated; + let mut remaining_total_bytes = total_byte_budget; + + for path in paths.iter().take(MAX_SESSION_FILES) { + if remaining_total_bytes == 0 { + complete = false; + break; + } + let file = match File::open(path) { + Ok(file) => file, + Err(_) => { + complete = false; + continue; + } + }; + match file.metadata() { + Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, + Ok(_) => {} + Err(_) => complete = false, + } + let mut reader = BufReader::new(file); + let mut remaining = MAX_SESSION_FILE_BYTES; + let mut path_had_usage = false; + let mut model = None::; + loop { + let line = match read_bounded_jsonl_line( + &mut reader, + &mut remaining, + &mut remaining_total_bytes, + ) { + Ok(Some(BoundedJsonlLine::Record(line))) => line, + Ok(Some(BoundedJsonlLine::Oversized)) => { + complete = false; + continue; + } + Ok(Some(BoundedJsonlLine::Truncated)) => { + complete = false; + break; + } + Ok(None) => break, + Err(_) => { + complete = false; + break; + } + }; + if line.is_empty() { + continue; + } + let Ok(value) = serde_json::from_slice::(&line) else { + complete = false; + continue; + }; + if !value.is_object() { + complete = false; + continue; + } + let kind = value.get("type").and_then(Value::as_str); + if kind == Some("session_meta") { + model = value + .get("modelId") + .or_else(|| value.get("model_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + continue; + } + if kind != Some("usage") && value.get("input").is_none() { + continue; + } + if !has_valid_token_fields(&value) { + complete = false; + continue; + } + + let Some(timestamp_ms) = value.get("timestamp").and_then(Value::as_i64) else { + complete = false; + continue; + }; + let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { + complete = false; + continue; + }; + + if let Some(response_id) = value + .get("responseId") + .or_else(|| value.get("response_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + && !seen_response_ids.insert(response_id.to_string()) + { + continue; + } + + if at > now || at.with_timezone(&Local).date_naive() < first_day { + continue; + } + + let input = token_field(&value, &["input"]); + let output = token_field(&value, &["output"]); + let cache_read = token_field(&value, &["cacheRead", "cache_read"]); + let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); + let reasoning = token_field( + &value, + &["reasoning", "reasoningTokens", "reasoning_tokens"], + ); + let Some(total) = [input, output, cache_read, cache_write, reasoning] + .into_iter() + .try_fold(0_u64, u64::checked_add) + else { + complete = false; + continue; + }; + if total == 0 { + continue; + } + let Some(next_total_tokens) = total_tokens.checked_add(total) else { + complete = false; + continue; + }; + total_tokens = next_total_tokens; + cost_estimate.record_list_price(estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + )); + path_had_usage = true; + } + if path_had_usage { + sessions_with_usage.insert(path.clone()); + } + } + + LocalTokenHistorySummary { + total_tokens, + session_count: sessions_with_usage.len(), + coverage: if paths.is_empty() { + LocalHistoryCoverage::Unavailable + } else if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, + cost_estimate, + } +} + +fn read_bounded_jsonl_line( + reader: &mut R, + remaining_file_bytes: &mut usize, + remaining_total_bytes: &mut usize, +) -> std::io::Result> { + if *remaining_file_bytes == 0 { + return Ok(None); + } + if *remaining_total_bytes == 0 { + return Ok(if reader.fill_buf()?.is_empty() { + None + } else { + Some(BoundedJsonlLine::Truncated) + }); + } + let mut line = Vec::new(); + let mut saw_input = false; + let mut discarding = false; + + loop { + let chunk = reader.fill_buf()?; + if chunk.is_empty() { + return Ok(saw_input.then_some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + let bounded_len = chunk + .len() + .min(*remaining_file_bytes) + .min(*remaining_total_bytes); + if bounded_len == 0 { + return Ok(None); + } + let bounded = &chunk[..bounded_len]; + let newline = bounded.iter().position(|byte| *byte == b'\n'); + let segment_end = newline.unwrap_or(bounded.len()); + let segment = &bounded[..segment_end]; + saw_input = saw_input || !segment.is_empty() || newline.is_some(); + if !discarding { + if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { + line.extend_from_slice(segment); + } else { + line.clear(); + discarding = true; + } + } + let consumed = segment_end + usize::from(newline.is_some()); + reader.consume(consumed); + *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); + *remaining_total_bytes = remaining_total_bytes.saturating_sub(consumed); + if newline.is_some() { + return Ok(Some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + if *remaining_file_bytes == 0 || *remaining_total_bytes == 0 { + let at_eof = reader.fill_buf()?.is_empty(); + return Ok(Some(if !at_eof { + BoundedJsonlLine::Truncated + } else if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + } +} + +fn token_field(value: &Value, keys: &[&str]) -> u64 { + keys.iter() + .find_map(|key| value.get(*key).and_then(Value::as_u64)) + .unwrap_or(0) +} + +fn has_valid_token_fields(value: &Value) -> bool { + let mut has_token_field = false; + for key in [ + "input", + "output", + "cacheRead", + "cache_read", + "cacheWrite", + "cache_write", + "reasoning", + "reasoningTokens", + "reasoning_tokens", + ] { + if let Some(field) = value.get(key) { + if field.as_u64().is_none() { + return false; + } + has_token_field = true; + } + } + has_token_field +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokscale_discovery_bounds_entries_and_retained_paths() { + let dir = tempfile::tempdir().unwrap(); + for index in 0..5 { + fs::write(dir.path().join(format!("session-{index}.jsonl")), "").unwrap(); + } + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 2, 10); + assert_eq!(paths.len(), 2); + assert!(truncated); + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 10, 2); + assert_eq!(paths.len(), 2); + assert!(truncated); + assert_eq!(paths[0].file_name().unwrap(), "session-3.jsonl"); + assert_eq!(paths[1].file_name().unwrap(), "session-4.jsonl"); + } + + #[test] + fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { + let dir = tempfile::tempdir().unwrap(); + let known = dir.path().join("known.jsonl"); + let unknown = dir.path().join("unknown.jsonl"); + fs::write( + &known, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", + "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" + ), + ) + .unwrap(); + fs::write( + &unknown, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", + "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[known, unknown], now, 7, false); + + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert!(summary.cost_estimate.known_subtotal_usd.is_some()); + assert_eq!(summary.total_usd(), None); + } + #[test] + fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write(&path, concat!( + "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", + "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", + "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" + )).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 135); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn truncated_or_unreadable_tokscale_history_is_partial() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); + assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); + + let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); + assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); + } + #[test] + fn excludes_usage_outside_requested_window() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", + "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + } + + #[test] + fn oversized_jsonl_line_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + let mut text = format!( + r#"{{"type":"usage","padding":"{}"}}"#, + "x".repeat(MAX_JSONL_LINE_BYTES + 32) + ); + text.push('\n'); + text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); + text.push('\n'); + fs::write(&path, text).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn malformed_jsonl_record_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{malformed json}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":\"invalid\"}\n", + "{\"type\":\"usage\",\"input\":10}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn total_scan_byte_budget_stops_later_records_and_marks_partial() { + let dir = tempfile::tempdir().unwrap(); + let first_path = dir.path().join("session-a.jsonl"); + let second_path = dir.path().join("session-b.jsonl"); + let first = "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n"; + fs::write(&first_path, first).unwrap(); + fs::write( + &second_path, + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":20}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = + summarize_paths_with_budget(&[first_path, second_path], now, 7, false, first.len()); + + assert_eq!(summary.total_tokens, 10); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn token_sum_overflow_marks_coverage_partial_without_saturation() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615,\"output\":1}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":1}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, u64::MAX); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } +} diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index 8e83eb8e33..4c5c2e5192 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -10,9 +10,12 @@ use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use rusqlite::{Connection, OpenFlags, TransactionBehavior, types::ValueRef}; use self::local_bot_id::{ExactStepTimestamp, embedded_timestamps_agree, record_exact_bot_id}; +use super::cost::estimate_cost_usd; use super::local_proto::{ParsedTurn, parse_step_metadata, parse_turn}; -use super::local_sessions::{LocalHistoryCoverage, LocalSessionSummary}; use super::local_step_resolver::{StepOccurrence, resolve_step_timestamps}; +#[cfg(test)] +use crate::spend_contract::LocalTokenHistorySummary as LocalSessionSummary; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; const MAX_DATABASES: usize = 500; const MAX_DIRECTORY_ENTRIES: usize = 10_000; @@ -34,7 +37,7 @@ pub(super) enum SQLiteScan { /// This is non-authoritative: callers may continue with another local /// history source instead of treating the scan as known-empty history. Unsupported, - Summary(LocalSessionSummary), + Summary(LocalTokenHistorySummary), } #[derive(Debug)] @@ -262,7 +265,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL let input = usage.system_prompt.checked_add(usage.new_input); let output = usage.output.checked_add(usage.reasoning); if let (Some(input), Some(output)) = (input, output) { - super::local_sessions::estimate_cost_usd(model, input, usage.cache_read, 0, output) + estimate_cost_usd(model, input, usage.cache_read, 0, output) } else { None } @@ -271,7 +274,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL sessions.insert(event.session); } - SQLiteScan::Summary(LocalSessionSummary { + SQLiteScan::Summary(LocalTokenHistorySummary { total_tokens, session_count: sessions.len(), coverage: if complete { @@ -389,10 +392,10 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } + } + + fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes + } + + fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes + } + + fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut timestamp = field_varint(1, timestamp_seconds); + timestamp.extend(field_varint(2, 0)); + let mut chat = field_bytes(4, &usage); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + field_bytes(1, &chat) + } + #[test] fn missing_databases_falls_through() { let dir = tempfile::tempdir().unwrap(); @@ -892,6 +937,36 @@ mod tests { assert_eq!(summary.session_count, 0); } + #[test] + fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { + let dir = tempfile::tempdir().unwrap(); + let first_root = dir.path().join("first"); + let second_root = dir.path().join("second"); + fs::create_dir_all(&first_root).unwrap(); + fs::create_dir_all(&second_root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + + for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { + let conn = Connection::open(root.join("session.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) + else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 2); + } + #[test] fn non_blob_rows_make_coverage_partial() { let dir = tempfile::tempdir().unwrap(); diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 5f35764663..d24eb8c637 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -4,9 +4,12 @@ //! Uses Windows process detection to find CSRF token mod cli_fallback; +mod cost; mod legacy_status; +mod local_history; mod local_proto; pub mod local_sessions; +mod local_sessions_reader; mod local_sqlite; mod local_step_resolver; mod quota_summary; @@ -50,6 +53,7 @@ const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; + /// Serialize task-owned `agy` launches so concurrent app surfaces never start /// multiple interactive CLI servers at the same time. #[cfg(windows)] @@ -623,7 +627,7 @@ impl AntigravityProvider { } fn offline_usage_result() -> Option { - let count = local_sessions::offline_conversation_count(); + let count = local_history::offline_conversation_count(); if count == 0 { return None; } From cce5b27c5eef2b97350bd648792b60e46a162e28 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:45:50 +0700 Subject: [PATCH 39/62] Preserve uncertainty in Claude history scans --- rust/src/core/claude_routed_pricing.rs | 102 ++++++++---- rust/src/core/cost_pricing/claude.rs | 138 ++++++++++------ rust/src/cost_scanner.rs | 213 +++++++++++++++---------- rust/src/cost_scanner/tests.rs | 97 +++++++++++ 4 files changed, 386 insertions(+), 164 deletions(-) diff --git a/rust/src/core/claude_routed_pricing.rs b/rust/src/core/claude_routed_pricing.rs index 1076992c9f..1f6a750b69 100644 --- a/rust/src/core/claude_routed_pricing.rs +++ b/rust/src/core/claude_routed_pricing.rs @@ -147,13 +147,52 @@ pub fn cost_usd_from_pricing_with_threshold( cache_write: i32, output: i32, ) -> f64 { - let input = input.max(0); - let cache_read = cache_read.max(0); - let cache_write = cache_write.max(0); - let output = output.max(0); + cost_usd_from_u64_counts_with_threshold( + pricing, + threshold_tokens, + input.max(0) as u64, + cache_read.max(0) as u64, + cache_write.max(0) as u64, + output.max(0) as u64, + ) +} + +/// Calculate routed cost for local history counters without narrowing them to +/// the signed API token-count type. +pub(crate) fn cost_usd_from_u64_counts_with_threshold( + pricing: models_dev_pricing::DynamicModelPricing, + threshold_tokens: Option, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> f64 { let use_tier = threshold_tokens.is_some_and(|threshold| { - (input as u64) + (cache_read as u64) + (cache_write as u64) > threshold + input + .checked_add(cache_read) + .and_then(|total| total.checked_add(cache_write)) + .is_none_or(|total| total > threshold) }); + let rates = selected_cost_rates(pricing, use_tier); + + (input as f64) * rates.input + + (cache_read as f64) * rates.cache_read + + (cache_write as f64) * rates.cache_write + + (output as f64) * rates.output +} + +#[derive(Debug, Clone, Copy)] +struct SelectedCostRates { + input: f64, + cache_read: f64, + cache_write: f64, + output: f64, +} + +fn selected_cost_rates( + pricing: models_dev_pricing::DynamicModelPricing, + use_tier: bool, +) -> SelectedCostRates { let pick = |base: f64, above: Option| { if use_tier { above.unwrap_or(base) @@ -161,39 +200,34 @@ pub fn cost_usd_from_pricing_with_threshold( base } }; - let input_rate = pick( + let input = pick( pricing.input_cost_per_token, pricing.input_cost_per_token_above_threshold, ); - let cache_read_rate = if use_tier { - pricing - .cache_read_input_cost_per_token_above_threshold - .or(pricing.cache_read_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_read_input_cost_per_token - .unwrap_or(input_rate) - }; - let cache_write_rate = if use_tier { - pricing - .cache_write_input_cost_per_token_above_threshold - .or(pricing.cache_write_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_write_input_cost_per_token - .unwrap_or(input_rate) - }; - let output_rate = pick( - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - ); - (input as f64) * input_rate - + (cache_read as f64) * cache_read_rate - + (cache_write as f64) * cache_write_rate - + (output as f64) * output_rate + SelectedCostRates { + input, + cache_read: if use_tier { + pricing + .cache_read_input_cost_per_token_above_threshold + .or(pricing.cache_read_input_cost_per_token) + .unwrap_or(input) + } else { + pricing.cache_read_input_cost_per_token.unwrap_or(input) + }, + cache_write: if use_tier { + pricing + .cache_write_input_cost_per_token_above_threshold + .or(pricing.cache_write_input_cost_per_token) + .unwrap_or(input) + } else { + pricing.cache_write_input_cost_per_token.unwrap_or(input) + }, + output: pick( + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + ), + } } fn effective_threshold(provider: &str, model: &str, catalog_threshold: Option) -> Option { diff --git a/rust/src/core/cost_pricing/claude.rs b/rust/src/core/cost_pricing/claude.rs index 05429d57a9..fd06061eeb 100644 --- a/rust/src/core/cost_pricing/claude.rs +++ b/rust/src/core/cost_pricing/claude.rs @@ -149,54 +149,14 @@ impl CostUsagePricing { ClaudePricingResolution::ModelsDev { pricing, threshold_tokens, - } => { - let use_tier = threshold_tokens.is_some_and(|threshold| { - input_tokens - .checked_add(cache_read_input_tokens) - .and_then(|value| value.checked_add(cache_creation_input_tokens)) - .is_none_or(|total| total > threshold) - }); - let pick = |base: f64, above: Option| { - if use_tier { - above.unwrap_or(base) - } else { - base - } - }; - let input_rate = pick( - pricing.input_cost_per_token, - pricing.input_cost_per_token_above_threshold, - ); - let cache_read_rate = if use_tier { - pricing - .cache_read_input_cost_per_token_above_threshold - .or(pricing.cache_read_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_read_input_cost_per_token - .unwrap_or(input_rate) - }; - let cache_write_rate = if use_tier { - pricing - .cache_write_input_cost_per_token_above_threshold - .or(pricing.cache_write_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_write_input_cost_per_token - .unwrap_or(input_rate) - }; - let output_rate = pick( - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - ); - - (input_tokens as f64) * input_rate - + (cache_read_input_tokens as f64) * cache_read_rate - + (cache_creation_input_tokens as f64) * cache_write_rate - + (output_tokens as f64) * output_rate - } + } => claude_routed_pricing::cost_usd_from_u64_counts_with_threshold( + pricing, + threshold_tokens, + input_tokens, + cache_read_input_tokens, + cache_creation_input_tokens, + output_tokens, + ), } } @@ -254,3 +214,85 @@ impl CostUsagePricing { claude_routed_pricing::input_cost_per_token(model, Self::normalize_claude_model(model)) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn routed_pricing() -> (models_dev_pricing::DynamicModelPricing, Option) { + let snapshot = models_dev_pricing::ModelsDevPricingSnapshot::from_catalog_json_for_tests( + r#"{ + "anthropic": {"models": {"threshold-fixture": {"id": "threshold-fixture", "cost": { + "input": 2, "output": 4, "cache_read": 0.25, "cache_write": 3, + "context_over_200k": {"input": 7, "output": 11, "cache_read": 0.5, "cache_write": 9} + }}}} + }"#, + ) + .expect("pricing fixture"); + let resolution = CostUsagePricing::resolve_claude_pricing( + "anthropic/threshold-fixture", + "anthropic/threshold-fixture", + Some(&snapshot), + ) + .expect("Models.dev pricing"); + let ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens, + } = resolution + else { + panic!("expected Models.dev pricing"); + }; + (pricing, threshold_tokens) + } + + #[test] + fn models_dev_u64_cost_uses_routed_rates_for_every_token_field() { + let (pricing, threshold) = routed_pricing(); + assert_eq!(threshold, Some(200_000)); + + for (input, cache_read, cache_write, output, expected_usd) in [ + (10_000, 2_000, 1_000, 500, 0.0255), + (199_999, 1, 0, 25, 0.400_098_25), + (200_000, 1, 0, 25, 1.400_275_5), + (220_000, 10_000, 2_000, 50, 1.563_55), + ] { + let actual = CostUsagePricing::claude_cost_usd_u64_from_resolution( + ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens: threshold, + }, + input, + cache_read, + cache_write, + output, + ); + let routed = claude_routed_pricing::cost_usd_from_pricing_with_threshold( + pricing, + threshold, + input as i32, + cache_read as i32, + cache_write as i32, + output as i32, + ); + assert!((actual - expected_usd).abs() < 1e-12); + assert!((actual - routed).abs() < 1e-12); + } + } + + #[test] + fn models_dev_u64_counter_overflow_selects_above_threshold_rates() { + let (pricing, threshold) = routed_pricing(); + let actual = CostUsagePricing::claude_cost_usd_u64_from_resolution( + ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens: threshold, + }, + u64::MAX, + 1, + 0, + 0, + ); + let expected = (u64::MAX as f64) * 7e-6 + 0.5e-6; + assert!((actual - expected).abs() < 1e-6); + } +} diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 2b730cdc01..3046996972 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -608,6 +608,7 @@ impl CostScanner { cancel, &mut pricing, |record| { + aggregation_complete &= record.timestamp.is_some(); aggregation_complete &= add_claude_record_to_summary(&mut summary, record); }, ); @@ -620,7 +621,12 @@ impl CostScanner { } claude_scan.absorb(file_result); }; - self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file); + let traversal_read_failures = + self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file); + drop(handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); } // OMP / pi-compatible anthropic rows, deduped across shared files. @@ -675,42 +681,45 @@ impl CostScanner { let mut quota_records = Vec::new(); let mut scan_result = ClaudeFileScanResult::default(); - let mut missing_timestamp = false; if projects_dir.exists() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut |path| { - let mut file_has_usage = false; - let mut aggregation_complete = true; - let mut file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - cancel, - &mut pricing, - |record| { - file_has_usage = true; - aggregation_complete &= add_claude_record_to_summary(&mut summary, record); - aggregation_complete &= - add_claude_record_to_daily_costs(&mut daily_cost, record); - aggregation_complete &= - add_claude_record_to_daily_tokens(&mut daily_tokens, record); - if let Some(quota_record) = quota_history_record_from_usage(record) { - quota_records.push(quota_record); - } else { - missing_timestamp = true; - } - }, - ); - if !aggregation_complete { - file_result.aggregation_failures = - file_result.aggregation_failures.saturating_add(1); - } - if file_has_usage { - summary.sessions_count += 1; - } - scan_result.absorb(file_result); - }); + let traversal_read_failures = + self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut |path| { + let mut file_has_usage = false; + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + cancel, + &mut pricing, + |record| { + file_has_usage = true; + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_summary(&mut summary, record); + aggregation_complete &= + add_claude_record_to_daily_costs(&mut daily_cost, record); + aggregation_complete &= + add_claude_record_to_daily_tokens(&mut daily_tokens, record); + if let Some(quota_record) = quota_history_record_from_usage(record) { + quota_records.push(quota_record); + } + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + if file_has_usage { + summary.sessions_count += 1; + } + scan_result.absorb(file_result); + }); + scan_result.read_failures = scan_result + .read_failures + .saturating_add(traversal_read_failures); } crate::pi_session_cost::scan_pi_compatible_into( @@ -721,10 +730,7 @@ impl CostScanner { &mut HashSet::new(), ); - let complete = projects_dir.exists() - && !is_cancelled(cancel) - && scan_result.is_complete() - && !missing_timestamp; + let complete = projects_dir.exists() && !is_cancelled(cancel) && scan_result.is_complete(); finalize_claude_summary( &mut summary, projects_dir.exists(), @@ -822,36 +828,52 @@ impl CostScanner { cutoff: &DateTime, cancel: Option<&AtomicBool>, on_file: &mut F, - ) where + ) -> u32 + where F: FnMut(&Path), { if is_cancelled(cancel) { - return; + return 0; } let entries = match fs::read_dir(dir) { Ok(e) => e, - Err(_) => return, + Err(_) => return 1, }; - for entry in entries.flatten() { + let mut read_failures = 0u32; + for entry in entries { if is_cancelled(cancel) { break; } + let entry = match entry { + Ok(entry) => entry, + Err(_) => { + read_failures = read_failures.saturating_add(1); + continue; + } + }; let path = entry.path(); - if path.is_dir() { - self.walk_claude_files(&path, cutoff, cancel, on_file); - } else if path.extension().is_some_and(|e| e == "jsonl") { - // Check file modification time - if let Ok(metadata) = fs::metadata(&path) - && let Ok(modified) = metadata.modified() - { - let modified_dt: DateTime = modified.into(); - if modified_dt >= *cutoff { - on_file(&path); + match fs::metadata(&path) { + Ok(metadata) if metadata.is_dir() => { + read_failures = read_failures + .saturating_add(self.walk_claude_files(&path, cutoff, cancel, on_file)); + } + Ok(metadata) if path.extension().is_some_and(|e| e == "jsonl") => { + match metadata.modified() { + Ok(modified) => { + let modified_dt: DateTime = modified.into(); + if modified_dt >= *cutoff { + on_file(&path); + } + } + Err(_) => read_failures = read_failures.saturating_add(1), } } + Ok(_) => {} + Err(_) => read_failures = read_failures.saturating_add(1), } } + read_failures } } @@ -875,20 +897,6 @@ where scan_claude_file_with_pricing(path, cutoff, seen, cancel, &mut pricing, on_record).counted } -fn for_each_claude_usage_record_with_pricing( - path: &Path, - cutoff: &DateTime, - seen: &mut HashSet, - cancel: Option<&AtomicBool>, - pricing: &mut ClaudeScanPricingResolver, - on_record: F, -) -> usize -where - F: FnMut(&ClaudeUsageRecord), -{ - scan_claude_file_with_pricing(path, cutoff, seen, cancel, pricing, on_record).counted -} - fn scan_claude_file_with_pricing( path: &Path, cutoff: &DateTime, @@ -1248,19 +1256,31 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< let mut pricing = ClaudeScanPricingResolver::default(); let mut claude_scan = ClaudeFileScanResult::default(); let mut handle_file = |path: &Path| { - let file_result = scan_claude_file_with_pricing( + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( path, &cutoff, &mut seen, None, &mut pricing, |record| { - add_claude_record_to_daily_costs(&mut daily_costs, record); + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_daily_costs(&mut daily_costs, record); }, ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } claude_scan.absorb(file_result); }; - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + let traversal_read_failures = + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + drop(handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); if claude_scan.is_complete() { for slot in daily_costs.values_mut() { if slot.is_none() { @@ -1353,26 +1373,24 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - let mut aggregation_complete = true; + let mut claude_scan = ClaudeFileScanResult::default(); let mut handle_file = |path: &Path| { - for_each_claude_usage_record_with_pricing( + let file_result = scan_claude_file_for_daily_tokens( path, &cutoff, &mut seen, - None, &mut pricing, - |record| { - aggregation_complete &= - add_claude_record_to_daily_tokens(&mut daily_tokens, record); - }, + &mut daily_tokens, ); + claude_scan.absorb(file_result); }; - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); - if !aggregation_complete { - covered_days.clear(); - } else { - covered_days.extend(daily_tokens.keys().cloned()); - } + let traversal_read_failures = + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + drop(handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); + mark_claude_daily_token_coverage(&mut covered_days, &daily_tokens, claude_scan); } } "pi" => { @@ -1434,3 +1452,34 @@ fn add_claude_record_to_daily_tokens( } true } + +fn scan_claude_file_for_daily_tokens( + path: &Path, + cutoff: &DateTime, + seen: &mut HashSet, + pricing: &mut ClaudeScanPricingResolver, + daily_tokens: &mut HashMap, +) -> ClaudeFileScanResult { + let mut aggregation_failures = 0u32; + let mut result = scan_claude_file_with_pricing(path, cutoff, seen, None, pricing, |record| { + if record.timestamp.is_none() || !add_claude_record_to_daily_tokens(daily_tokens, record) { + aggregation_failures = aggregation_failures.saturating_add(1); + } + }); + result.aggregation_failures = result + .aggregation_failures + .saturating_add(aggregation_failures); + result +} + +fn mark_claude_daily_token_coverage( + covered_days: &mut HashSet, + daily_tokens: &HashMap, + scan_result: ClaudeFileScanResult, +) { + if scan_result.is_complete() { + covered_days.extend(daily_tokens.keys().cloned()); + } else { + covered_days.clear(); + } +} diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index d20e770b26..221f9d2b8a 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -495,6 +495,103 @@ fn malformed_claude_history_stays_unknown_while_valid_empty_history_is_known_zer assert!(!malformed_summary.known_zero); } +#[test] +fn claude_daily_token_coverage_requires_a_complete_valid_scan() { + let root = tempfile::tempdir().unwrap(); + let cutoff = Utc::now() - Duration::days(1); + let valid_path = root.path().join("valid.jsonl"); + let timestamp = Utc::now() - Duration::hours(1); + let today = timestamp + .with_timezone(&Local) + .date_naive() + .format("%Y-%m-%d") + .to_string(); + std::fs::write( + &valid_path, + format!( + "{}\n", + claude_transcript_line( + ×tamp.to_rfc3339(), + "requestId", + "req_valid", + "msg_valid" + ) + ), + ) + .unwrap(); + + let mut valid_tokens = HashMap::from([(today.clone(), 0)]); + let valid_result = scan_claude_file_for_daily_tokens( + &valid_path, + &cutoff, + &mut HashSet::new(), + &mut ClaudeScanPricingResolver::default(), + &mut valid_tokens, + ); + assert!(valid_result.is_complete()); + let mut covered_days = HashSet::new(); + mark_claude_daily_token_coverage(&mut covered_days, &valid_tokens, valid_result); + assert!(covered_days.contains(&today)); + + let assert_uncovered = |path: &Path| { + let mut daily_tokens = HashMap::from([(today.clone(), 0)]); + let result = scan_claude_file_for_daily_tokens( + path, + &cutoff, + &mut HashSet::new(), + &mut ClaudeScanPricingResolver::default(), + &mut daily_tokens, + ); + assert!(!result.is_complete()); + let mut covered_days = HashSet::from(["stale-coverage".to_string()]); + mark_claude_daily_token_coverage(&mut covered_days, &daily_tokens, result); + assert!(covered_days.is_empty()); + result + }; + + let malformed_path = root.path().join("malformed.jsonl"); + std::fs::write(&malformed_path, b"{malformed\n").unwrap(); + assert_eq!(assert_uncovered(&malformed_path).malformed_lines, 1); + + let incomplete_path = root.path().join("incomplete.jsonl"); + std::fs::write( + &incomplete_path, + r#"{"type":"assistant","message":{"id":"msg_preliminary","model":"gpt-5.6-sol","stop_reason":null,"usage":{"input_tokens":1000}}}"#, + ) + .unwrap(); + assert_eq!(assert_uncovered(&incomplete_path).incomplete_requests, 1); + + let missing_timestamp_path = root.path().join("missing-timestamp.jsonl"); + std::fs::write( + &missing_timestamp_path, + r#"{"type":"assistant","requestId":"req_no_timestamp","message":{"id":"msg_no_timestamp","model":"claude-sonnet-4-6","usage":{"input_tokens":1000,"output_tokens":500}}}"#, + ) + .unwrap(); + assert_eq!( + assert_uncovered(&missing_timestamp_path).aggregation_failures, + 1 + ); + + let unreadable_path = root.path().join("missing.jsonl"); + assert_eq!(assert_uncovered(&unreadable_path).read_failures, 1); + + let scanner = CostScanner::new(1); + let missing_directory = root.path().join("missing-directory"); + let traversal_read_failures = + scanner.walk_claude_files(&missing_directory, &cutoff, None, &mut |_| {}); + assert_eq!(traversal_read_failures, 1); + let mut covered_days = HashSet::from([today]); + mark_claude_daily_token_coverage( + &mut covered_days, + &valid_tokens, + ClaudeFileScanResult { + read_failures: traversal_read_failures, + ..ClaudeFileScanResult::default() + }, + ); + assert!(covered_days.is_empty()); +} + #[test] fn classifies_vertex_ai_claude_metadata_without_changing_anthropic_rows() { let cases = [ From 69398db13302815d4b847803a1dfba6fff3490ba Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:07:01 +0700 Subject: [PATCH 40/62] Avoid Clippy warnings in Claude history scanning --- rust/src/core/cost_pricing/claude.rs | 10 +-- rust/src/cost_scanner.rs | 117 ++++++++++++++------------- 2 files changed, 64 insertions(+), 63 deletions(-) diff --git a/rust/src/core/cost_pricing/claude.rs b/rust/src/core/cost_pricing/claude.rs index fd06061eeb..32e4aa4e36 100644 --- a/rust/src/core/cost_pricing/claude.rs +++ b/rust/src/core/cost_pricing/claude.rs @@ -266,13 +266,13 @@ mod tests { cache_write, output, ); - let routed = claude_routed_pricing::cost_usd_from_pricing_with_threshold( + let routed = claude_routed_pricing::cost_usd_from_u64_counts_with_threshold( pricing, threshold, - input as i32, - cache_read as i32, - cache_write as i32, - output as i32, + input, + cache_read, + cache_write, + output, ); assert!((actual - expected_usd).abs() < 1e-12); assert!((actual - routed).abs() < 1e-12); diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 3046996972..f936ec34c7 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -599,31 +599,32 @@ impl CostScanner { if projects_dir.exists() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - let mut handle_file = |path: &Path| { - let mut aggregation_complete = true; - let mut file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - cancel, - &mut pricing, - |record| { - aggregation_complete &= record.timestamp.is_some(); - aggregation_complete &= add_claude_record_to_summary(&mut summary, record); - }, - ); - if !aggregation_complete { - file_result.aggregation_failures = - file_result.aggregation_failures.saturating_add(1); - } - if file_result.counted > 0 { - summary.sessions_count += 1; - } - claude_scan.absorb(file_result); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + cancel, + &mut pricing, + |record| { + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_summary(&mut summary, record); + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + if file_result.counted > 0 { + summary.sessions_count += 1; + } + claude_scan.absorb(file_result); + }; + self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file) }; - let traversal_read_failures = - self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file); - drop(handle_file); claude_scan.read_failures = claude_scan .read_failures .saturating_add(traversal_read_failures); @@ -1255,29 +1256,29 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let mut claude_scan = ClaudeFileScanResult::default(); - let mut handle_file = |path: &Path| { - let mut aggregation_complete = true; - let mut file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - None, - &mut pricing, - |record| { - aggregation_complete &= record.timestamp.is_some(); - aggregation_complete &= - add_claude_record_to_daily_costs(&mut daily_costs, record); - }, - ); - if !aggregation_complete { - file_result.aggregation_failures = - file_result.aggregation_failures.saturating_add(1); - } - claude_scan.absorb(file_result); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + None, + &mut pricing, + |record| { + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_daily_costs(&mut daily_costs, record); + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + claude_scan.absorb(file_result); + }; + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file) }; - let traversal_read_failures = - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); - drop(handle_file); claude_scan.read_failures = claude_scan .read_failures .saturating_add(traversal_read_failures); @@ -1374,19 +1375,19 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let mut claude_scan = ClaudeFileScanResult::default(); - let mut handle_file = |path: &Path| { - let file_result = scan_claude_file_for_daily_tokens( - path, - &cutoff, - &mut seen, - &mut pricing, - &mut daily_tokens, - ); - claude_scan.absorb(file_result); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let file_result = scan_claude_file_for_daily_tokens( + path, + &cutoff, + &mut seen, + &mut pricing, + &mut daily_tokens, + ); + claude_scan.absorb(file_result); + }; + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file) }; - let traversal_read_failures = - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); - drop(handle_file); claude_scan.read_failures = claude_scan .read_failures .saturating_add(traversal_read_failures); From 94884dbe3063c64fcf8833f3f6e5d09218d9d780 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:43:08 +0700 Subject: [PATCH 41/62] Stabilize tray panel sizing test --- .../hooks/useTrayPanelLayout.sizing.test.tsx | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx index 02a55b21ed..7dbf6d6f7d 100644 --- a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx +++ b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx @@ -185,11 +185,20 @@ describe("useTrayPanelLayout sizing", () => { tauriMocks.revealTrayPanelWindow.mock.calls.length; expect(settledRevealCount - revealsBeforeSettle).toBeLessThanOrEqual(1); - await act(async () => { - await new Promise((resolve) => window.setTimeout(resolve, 500)); - }); - expect(tauriMocks.revealTrayPanelWindow.mock.calls.length).toBe( - settledRevealCount, + let lastRevealCount = settledRevealCount; + let stableSince = Date.now(); + await waitFor( + () => { + const revealCount = + tauriMocks.revealTrayPanelWindow.mock.calls.length; + if (revealCount !== lastRevealCount) { + lastRevealCount = revealCount; + stableSince = Date.now(); + } + expect(revealCount - revealsBeforeSettle).toBeLessThanOrEqual(1); + expect(Date.now() - stableSince).toBeGreaterThanOrEqual(500); + }, + { timeout: 3000, interval: 50 }, ); }); From 0f759e624376f905d3047f5f39bc698b6f54939d Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:37:11 +0700 Subject: [PATCH 42/62] Remove shell startup timing from login exit tests --- rust/src/providers/claude/accounts/login.rs | 29 +++++++++++++-------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/rust/src/providers/claude/accounts/login.rs b/rust/src/providers/claude/accounts/login.rs index 8f0a6e002f..a3f64d00f4 100644 --- a/rust/src/providers/claude/accounts/login.rs +++ b/rust/src/providers/claude/accounts/login.rs @@ -582,23 +582,30 @@ mod tests { ) .unwrap(); std::fs::write(dir.path().join(".claude.json"), r#"{"oauthAccount":{"accountUuid":"test","organizationUuid":"org","emailAddress":"test@example.com"}}"#).unwrap(); + // This test covers exit handling and credential isolation, not shell + // startup speed. Reap each fixture before starting the login deadline; + // cancelled_and_timed_out_logins_reap_child covers a running process. + let mut successful_child = child(dir.path(), false, 0); + assert!(successful_child.wait().unwrap().success()); let login = wait_for_login( - &mut child(dir.path(), false, 0), + &mut successful_child, dir.path(), &AtomicBool::new(false), - Duration::from_secs(10), + Duration::ZERO, ) .unwrap(); assert_eq!(login.id().unwrap(), "test:org"); - assert!( - wait_for_login( - &mut child(dir.path(), false, 1), - dir.path(), - &AtomicBool::new(false), - Duration::from_secs(10) - ) - .is_err() - ); + let mut failed_child = child(dir.path(), false, 1); + assert_eq!(failed_child.wait().unwrap().code(), Some(1)); + let error = wait_for_login( + &mut failed_child, + dir.path(), + &AtomicBool::new(false), + Duration::ZERO, + ) + .err() + .expect("a failed login process must be rejected"); + assert!(error.to_string().contains("exit code 1")); } #[test] From a6517c387a23096baf3c3abe953850e27a1f5265 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:07:30 +0700 Subject: [PATCH 43/62] Require known Kimi legacy window duration --- rust/src/providers/kimi/code_api.rs | 57 +++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 7 deletions(-) diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 7e488594e8..8531eaad50 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -118,13 +118,8 @@ pub(super) fn snapshot_from_code_api_response( ) -> Result { let pools_present = response.usages.is_some(); let legacy_limit = response.limits.as_ref().and_then(|limits| limits.first()); - let legacy_session_minutes = legacy_limit.map(|limit| { - limit - .window - .as_ref() - .and_then(kimi_window_minutes) - .unwrap_or(300) - }); + let legacy_session_minutes = + legacy_limit.and_then(|limit| limit.window.as_ref().and_then(kimi_window_minutes)); let session_pool = response .usages .as_ref() @@ -589,6 +584,54 @@ mod tests { assert_eq!(weekly.window_minutes, Some(10_080)); } + fn snapshot_with_zero_session_ratio_and_legacy_window( + window: Option, + ) -> UsageSnapshot { + let mut legacy_limit = json!({ + "detail": { + "limit": "100", + "used": "1", + "resetTime": "2026-09-19T14:45:58Z" + } + }); + if let Some(window) = window { + legacy_limit["window"] = window; + } + + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "limits": [legacy_limit], + "usages": { + "limit_5h": { + "used_ratio": 0, + "reset_time": "2026-09-19T14:45:58Z" + }, + "limit_7d": { "used_ratio": 0 } + } + })) + .expect("fixture parses"); + + snapshot_from_code_api_response(response).expect("ratio pools are usable") + } + + #[test] + fn missing_legacy_window_does_not_override_zero_session_ratio() { + let snapshot = snapshot_with_zero_session_ratio_and_legacy_window(None); + + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!(snapshot.primary.used_percent, 0.0); + } + + #[test] + fn unrecognized_legacy_window_does_not_override_zero_session_ratio() { + let snapshot = snapshot_with_zero_session_ratio_and_legacy_window(Some(json!({ + "duration": 300, + "timeUnit": "TIME_UNIT_FORTNIGHT" + }))); + + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!(snapshot.primary.used_percent, 0.0); + } + #[test] fn zero_ratio_with_different_reset_stays_authoritative() { let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ From 4f5a6614e40b26cee7460c5ddc9607f250e04982 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:19:30 +0700 Subject: [PATCH 44/62] Test Claude daily history public coverage --- rust/src/cost_scanner/tests.rs | 54 ++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 63e5478025..219aed2e48 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -592,6 +592,60 @@ fn claude_daily_token_coverage_requires_a_complete_valid_scan() { assert!(covered_days.is_empty()); } +#[test] +fn public_claude_daily_token_dispatch_reports_incomplete_fixture_scans() { + const CHILD_MARKER: &str = "CODEXBAR_CLAUDE_DAILY_TOKEN_TEST_CHILD"; + const CHILD_DONE: &str = "isolated Claude daily-history fixture verified"; + if std::env::var_os(CHILD_MARKER).is_some() { + let config_dir = std::env::var_os("CLAUDE_CONFIG_DIR") + .map(PathBuf::from) + .expect("child receives isolated Claude config directory"); + let projects_dir = config_dir.join("projects"); + let project_dir = projects_dir.join("fixture-project"); + let (complete_history, incomplete) = get_daily_token_history("claude", 1); + assert!(!incomplete, "valid fixture scan should establish coverage"); + assert!(complete_history.iter().any(|(_, tokens)| *tokens > 0)); + + std::fs::write(project_dir.join("malformed.jsonl"), b"{malformed\n").unwrap(); + let (partial_history, incomplete) = get_daily_token_history("claude", 1); + assert!( + incomplete, + "malformed fixture should leave coverage incomplete" + ); + assert_eq!(partial_history, complete_history); + println!("{CHILD_DONE}"); + return; + } + + let config_dir = tempfile::tempdir().unwrap(); + let project_dir = config_dir.path().join("projects").join("fixture-project"); + std::fs::create_dir_all(&project_dir).unwrap(); + let timestamp = Utc::now().to_rfc3339(); + std::fs::write( + project_dir.join("valid.jsonl"), + format!( + "{}\n", + claude_transcript_line(×tamp, "requestId", "req_public", "msg_public") + ), + ) + .unwrap(); + + let test_thread = std::thread::current(); + let test_name = test_thread.name().expect("test harness names this thread"); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test_name, "--nocapture", "--test-threads=1"]) + .env(CHILD_MARKER, "1") + .env("CLAUDE_CONFIG_DIR", config_dir.path()) + .output() + .expect("spawn isolated exact-test child"); + assert!( + output.status.success() && String::from_utf8_lossy(&output.stdout).contains(CHILD_DONE), + "fixture child failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + #[test] fn classifies_vertex_ai_claude_metadata_without_changing_anthropic_rows() { let cases = [ From 5e5fd5a36e1e16f5cba9e7cd78f569e2520d3f4b Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:32:00 +0700 Subject: [PATCH 45/62] Bound tray sizing test timeout to async work --- apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx index 7dbf6d6f7d..1e70c015ff 100644 --- a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx +++ b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx @@ -258,7 +258,7 @@ describe("useTrayPanelLayout sizing", () => { await nudgePass(result, 417, "421px"); // → 421 → 526 phys expect(lastResize()).toEqual({ width: 328, height: 421 }); expect(surface.style.maxHeight).toBe("421px"); - }); + }, 30_000); // 8 bounded 3s settling passes + 3s readiness can exceed Vitest's 5s default. it("reconciles to the applied physical frame after an OS snap (no churn, no cycle)", async () => { // Deliberate 5-physical snap: requesting 539 logical (→674 phys) yields an From b35ec2941e99f730caa39eb495fd9cbc0e83e4de Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:09:39 +0700 Subject: [PATCH 46/62] Skip pricing zero-token Antigravity events --- .../src/providers/antigravity/local_sqlite.rs | 253 ++-------------- .../antigravity/local_sqlite_tests.rs | 280 ++++++++++++++++++ 2 files changed, 301 insertions(+), 232 deletions(-) create mode 100644 rust/src/providers/antigravity/local_sqlite_tests.rs diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index 4c5c2e5192..bc71f794d7 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -253,24 +253,26 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } - let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { - let inherited_model = event.turn.label.as_ref().and_then(|label| { - let key = (event.session.clone(), label.clone()); - (!conflicting_labels.contains(&key)) - .then(|| label_models.get(&key)) - .flatten() - .map(String::as_str) + if event.total > 0 { + let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { + let inherited_model = event.turn.label.as_ref().and_then(|label| { + let key = (event.session.clone(), label.clone()); + (!conflicting_labels.contains(&key)) + .then(|| label_models.get(&key)) + .flatten() + .map(String::as_str) + }); + let model = event.turn.model.as_deref().or(inherited_model); + let input = usage.system_prompt.checked_add(usage.new_input); + let output = usage.output.checked_add(usage.reasoning); + if let (Some(input), Some(output)) = (input, output) { + estimate_cost_usd(model, input, usage.cache_read, 0, output) + } else { + None + } }); - let model = event.turn.model.as_deref().or(inherited_model); - let input = usage.system_prompt.checked_add(usage.new_input); - let output = usage.output.checked_add(usage.reasoning); - if let (Some(input), Some(output)) = (input, output) { - estimate_cost_usd(model, input, usage.cache_read, 0, output) - } else { - None - } - }); - cost_estimate.record_list_price(estimated_cost); + cost_estimate.record_list_price(estimated_cost); + } sessions.insert(event.session); } @@ -848,218 +850,5 @@ fn has_stored_columns( mod synthetic_tests; #[cfg(test)] -mod tests { - use super::*; - use rusqlite::params; - - fn varint(mut value: u64) -> Vec { - let mut bytes = Vec::new(); - loop { - let mut byte = (value & 0x7f) as u8; - value >>= 7; - if value != 0 { - byte |= 0x80; - } - bytes.push(byte); - if value == 0 { - return bytes; - } - } - } - - fn field_varint(number: u64, value: u64) -> Vec { - let mut bytes = varint(number << 3); - bytes.extend(varint(value)); - bytes - } - - fn field_bytes(number: u64, value: &[u8]) -> Vec { - let mut bytes = varint((number << 3) | 2); - bytes.extend(varint(value.len() as u64)); - bytes.extend(value); - bytes - } - - fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { - let mut usage = field_varint(1, 11); - usage.extend(field_varint(2, input)); - usage.extend(field_varint(5, 50)); - usage.extend(field_varint(9, 30)); - usage.extend(field_varint(10, 7)); - - let mut timestamp = field_varint(1, timestamp_seconds); - timestamp.extend(field_varint(2, 0)); - let mut chat = field_bytes(4, &usage); - chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); - field_bytes(1, &chat) - } - - #[test] - fn missing_databases_falls_through() { - let dir = tempfile::tempdir().unwrap(); - assert!(matches!( - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), - SQLiteScan::NoDatabases - )); - } - - #[test] - fn foreign_database_is_non_authoritative() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) - .unwrap(); - drop(conn); - assert!(matches!( - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), - SQLiteScan::Unsupported - )); - } - - #[test] - fn empty_supported_database_is_confirmed_zero() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 0); - assert_eq!(summary.session_count, 0); - } - - #[test] - fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { - let dir = tempfile::tempdir().unwrap(); - let first_root = dir.path().join("first"); - let second_root = dir.path().join("second"); - fs::create_dir_all(&first_root).unwrap(); - fs::create_dir_all(&second_root).unwrap(); - let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); - - for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { - let conn = Connection::open(root.join("session.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob(input, timestamp)], - ) - .unwrap(); - } - - let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) - else { - panic!("supported databases should produce coverage"); - }; - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); - assert_eq!(summary.session_count, 2); - } - - #[test] - fn non_blob_rows_make_coverage_partial() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", - params![1_i64, "not-a-blob"], - ) - .unwrap(); - drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - } - #[test] - fn discovery_allows_exactly_500_databases_but_marks_501_partial() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join("dbs"); - fs::create_dir_all(&root).unwrap(); - for index in 0..MAX_DATABASES { - fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); - } - let mut budget = Budget::new(); - let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert_eq!(paths.len(), MAX_DATABASES); - assert!(complete); - - fs::write(root.join("overflow.db"), b"").unwrap(); - let mut budget = Budget::new(); - let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert_eq!(paths.len(), MAX_DATABASES); - assert!(!complete); - } - - #[test] - fn expired_budget_marks_discovery_incomplete() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join("dbs"); - fs::create_dir_all(&root).unwrap(); - fs::write(root.join("one.db"), b"").unwrap(); - let mut budget = Budget::with_deadline(Instant::now()); - let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert!(!complete); - } - - #[test] - fn extra_columns_and_without_rowid_schema_is_supported() { - let conn = Connection::open_in_memory().unwrap(); - conn.execute( - "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", - [], - ) - .unwrap(); - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Supported - ); - } - - #[test] - fn generated_columns_are_rejected() { - let conn = Connection::open_in_memory().unwrap(); - conn.execute( - "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", - [], - ) - .unwrap(); - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Unsupported - ); - } - - #[test] - fn schema_entry_budget_is_incomplete_not_foreign() { - let conn = Connection::open_in_memory().unwrap(); - for index in 0..=MAX_SCHEMA_ENTRIES { - conn.execute(&format!("CREATE TABLE unrelated_{index}(value TEXT)"), []) - .unwrap(); - } - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Incomplete - ); - } -} +#[path = "local_sqlite_tests.rs"] +mod tests; diff --git a/rust/src/providers/antigravity/local_sqlite_tests.rs b/rust/src/providers/antigravity/local_sqlite_tests.rs new file mode 100644 index 0000000000..1846bd5e26 --- /dev/null +++ b/rust/src/providers/antigravity/local_sqlite_tests.rs @@ -0,0 +1,280 @@ +use super::*; +use rusqlite::params; + +fn varint(mut value: u64) -> Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } +} + +fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes +} + +fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes +} + +fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { + valid_turn_blob_with_model(input, timestamp_seconds, None) +} + +fn valid_turn_blob_with_model(input: u64, timestamp_seconds: u64, model: Option<&str>) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut timestamp = field_varint(1, timestamp_seconds); + timestamp.extend(field_varint(2, 0)); + let mut chat = field_bytes(4, &usage); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + if let Some(model) = model { + chat.extend(field_bytes(19, model.as_bytes())); + } + field_bytes(1, &chat) +} + +fn zero_token_turn_blob(timestamp_seconds: u64, model: &str) -> Vec { + let timestamp = field_varint(1, timestamp_seconds); + let mut chat = field_bytes(4, &[]); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + chat.extend(field_bytes(19, model.as_bytes())); + field_bytes(1, &chat) +} + +#[test] +fn missing_databases_falls_through() { + let dir = tempfile::tempdir().unwrap(); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::NoDatabases + )); +} + +#[test] +fn foreign_database_is_non_authoritative() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::Unsupported + )); +} + +#[test] +fn empty_supported_database_is_confirmed_zero() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 0); + assert_eq!(summary.session_count, 0); +} + +#[test] +fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { + let dir = tempfile::tempdir().unwrap(); + let first_root = dir.path().join("first"); + let second_root = dir.path().join("second"); + fs::create_dir_all(&first_root).unwrap(); + fs::create_dir_all(&second_root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + + for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { + let conn = Connection::open(root.join("session.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 2); +} + +#[test] +fn zero_token_unknown_model_does_not_poison_priced_history() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let now = Utc::now(); + let timestamp = u64::try_from(now.timestamp()).unwrap(); + + let priced = Connection::open(root.join("priced.db")).unwrap(); + priced + .execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + priced + .execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob_with_model( + 100, + timestamp, + Some("claude-sonnet-4-6"), + )], + ) + .unwrap(); + + let zero = Connection::open(root.join("zero.db")).unwrap(); + zero.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + zero.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [zero_token_turn_blob(timestamp, "unknown-model")], + ) + .unwrap(); + + let SQLiteScan::Summary(summary) = summarize(&[root], now, 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 198); + assert_eq!(summary.session_count, 2); + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 0); + assert!( + summary + .cost_estimate + .known_subtotal_usd + .is_some_and(|cost| cost > 0.0) + ); + assert_eq!( + summary.total_usd(), + summary.cost_estimate.known_subtotal_usd + ); +} + +#[test] +fn non_blob_rows_make_coverage_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", + params![1_i64, "not-a-blob"], + ) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); +} +#[test] +fn discovery_allows_exactly_500_databases_but_marks_501_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + for index in 0..MAX_DATABASES { + fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); + } + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(complete); + + fs::write(root.join("overflow.db"), b"").unwrap(); + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(!complete); +} + +#[test] +fn expired_budget_marks_discovery_incomplete() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("one.db"), b"").unwrap(); + let mut budget = Budget::with_deadline(Instant::now()); + let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert!(!complete); +} + +#[test] +fn extra_columns_and_without_rowid_schema_is_supported() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Supported + ); +} + +#[test] +fn generated_columns_are_rejected() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Unsupported + ); +} + +#[test] +fn schema_entry_budget_is_incomplete_not_foreign() { + let conn = Connection::open_in_memory().unwrap(); + for index in 0..=MAX_SCHEMA_ENTRIES { + conn.execute(&format!("CREATE TABLE unrelated_{index}(value TEXT)"), []) + .unwrap(); + } + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Incomplete + ); +} From ab706db6ce0397e6655c761a2cde80bc5eafc40d Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:26:11 +0700 Subject: [PATCH 47/62] Canonicalize stacked provider preferences --- rust/src/settings/raw.rs | 8 ++++++-- rust/src/settings/tests.rs | 29 ++++++++++++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 1ac101f616..8fec108d41 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -624,8 +624,12 @@ impl From for Settings { ), merge_tray_icons: raw.merge_tray_icons, tray_icon_mode: raw.tray_icon_mode, - stacked_tray_top_provider: raw.stacked_tray_top_provider, - stacked_tray_bottom_provider: raw.stacked_tray_bottom_provider, + stacked_tray_top_provider: raw + .stacked_tray_top_provider + .and_then(|provider_id| canonical_provider_id(&provider_id)), + stacked_tray_bottom_provider: raw + .stacked_tray_bottom_provider + .and_then(|provider_id| canonical_provider_id(&provider_id)), switcher_shows_icons: raw.switcher_shows_icons, menu_bar_shows_highest_usage: raw.menu_bar_shows_highest_usage, menu_bar_shows_percent: raw.menu_bar_shows_percent, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index 9192b048ca..1ebf48fc2e 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -797,6 +797,17 @@ fn stacked_tray_mode_preserves_provider_preferences() { settings.stacked_tray_bottom_provider.as_deref(), Some("codex") ); + + let saved = serde_json::to_string(&settings).unwrap(); + let reloaded: Settings = serde_json::from_str(&saved).unwrap(); + assert_eq!( + reloaded.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + reloaded.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); } #[test] @@ -1080,6 +1091,8 @@ fn retired_provider_config_is_ignored_until_explicit_save() { "refresh_interval_secs": 300, "provider_metrics": { "codex": "weekly", "crof": "session" }, "float_bar_provider_ids": ["codex", "crof"], + "stacked_tray_top_provider": "crof", + "stacked_tray_bottom_provider": "crof", "provider_configs": { "crof": { "api_token": "retired-fixture-key" }, "codex": { "cookie_source": "manual", "openai_web_extras": false }, @@ -1100,6 +1113,8 @@ fn retired_provider_config_is_ignored_until_explicit_save() { ); assert_eq!(settings.provider_metrics.len(), 1); assert_eq!(settings.float_bar_provider_ids, ["codex"]); + assert_eq!(settings.stacked_tray_top_provider, None); + assert_eq!(settings.stacked_tray_bottom_provider, None); assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); assert_eq!( settings.manual_cookie_header(ProviderId::Alibaba), @@ -1128,7 +1143,10 @@ fn provider_aliases_are_canonicalized_at_the_load_boundary() { "CoDeX": "session", "not-a-provider": "weekly" }, - "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"] + "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"], + "enabled_providers": ["claude"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "ClAuDe" }"#, ) .expect("load settings containing provider aliases"); @@ -1143,6 +1161,15 @@ fn provider_aliases_are_canonicalized_at_the_load_boundary() { ); assert_eq!(settings.provider_metrics.len(), 1); assert_eq!(settings.float_bar_provider_ids, ["codex", "claude"]); + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("claude") + ); + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string()]) + ); } /// Default `Settings` should serialize WITHOUT a `provider_configs` From 775a4c4af9892a576d895e1fdd936cc50784ed4f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:30:16 +0700 Subject: [PATCH 48/62] Reuse Codex lineage plan per refresh --- rust/src/cost_scanner/codex.rs | 21 +- rust/src/cost_scanner/codex/logical_target.rs | 224 +++++++++++++----- rust/src/cost_scanner/codex/scan.rs | 27 ++- rust/src/cost_scanner/tests.rs | 6 + rust/src/cost_scanner/tests/lineage_cache.rs | 4 +- 5 files changed, 204 insertions(+), 78 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index a64eb92a39..23730eb28f 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -101,10 +101,6 @@ fn codex_usage_uses_parent(usage: &CostUsageFileUsage) -> bool { && usage.codex_forked_from_id.is_some()) } -fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { - CodexLineagePlanner::new(cache).cached_usage_is_safe(usage) -} - fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { usage .codex_fork_accounting_state @@ -202,6 +198,7 @@ impl CostScanner { sessions_dirs: &[PathBuf], range: &CostUsageDayRange, cache: &CostUsageCache, + planner: &CodexLineagePlanner, cancel: Option<&AtomicBool>, stats: &mut CostScanStats, ) -> (Vec, bool) { @@ -263,7 +260,7 @@ impl CostScanner { }; let mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); let unchanged_complete = - cached_codex_file_is_complete_for_range(cache, &path_key, range); + cached_codex_file_is_complete_for_range(cache, planner, &path_key, range); if unchanged_complete { stats.files_seen = stats.files_seen.saturating_add(1); stats.files_skipped = stats.files_skipped.saturating_add(1); @@ -323,8 +320,10 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, ) { - let _ = - self.parse_codex_file_bounded(path, range, summary, cache, cancel, stats, None, None); + let planner = CodexLineagePlanner::new(cache); + let _ = self.parse_codex_file_bounded( + path, range, summary, cache, cancel, stats, None, None, &planner, + ); } #[allow( @@ -341,6 +340,7 @@ impl CostScanner { stats: &mut CostScanStats, max_bytes_to_read: Option, prepared_candidate: Option<&CodexPreparedCandidate>, + planner: &CodexLineagePlanner, ) -> CodexFileScanOutcome { if is_cancelled(cancel) { return CodexFileScanOutcome::default(); @@ -377,7 +377,7 @@ impl CostScanner { }; } let cache_entry_is_fresh = |entry: &CostUsageFileUsage| { - cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) + cached_codex_file_is_fresh(cache, planner, entry, cache_covers_range, mtime_ms, size) }; let identity_matches_cached = |entry: &CostUsageFileUsage| { codex_file_identity_matches( @@ -502,7 +502,8 @@ impl CostScanner { .unwrap_or_default(); let parent_owner_expected = prepared_candidate.is_some_and(|candidate| candidate.parent_owner_expected); - let lineage_decision = CodexLineagePlanner::new(cache).decision_for_scan( + let lineage_decision = planner.decision_for_scan( + cache, is_fork, lineage_gate, codex_forked_from_id.as_deref(), @@ -545,7 +546,7 @@ impl CostScanner { } if let Some(entry) = &cached - && cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) + && cached_codex_file_is_fresh(cache, planner, entry, cache_covers_range, mtime_ms, size) && identity_matches_cached(entry) && !cached_identity_changed && !accounting_mode.requires_cached_reparse() diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 8382c0a098..a559588de8 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -1,4 +1,10 @@ use super::*; +#[cfg(test)] +use std::cell::Cell; +use std::collections::VecDeque; + +#[cfg(test)] +thread_local! { static CODEX_LINEAGE_GRAPH_BUILDS: Cell = const { Cell::new(0) }; } #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub(super) enum CodexLineageGate { @@ -36,6 +42,8 @@ struct CodexLineageGraph { impl CodexLineageGraph { fn new(cache: &CostUsageCache, candidates: Option<&[CodexPreparedCandidate]>) -> Self { + #[cfg(test)] + CODEX_LINEAGE_GRAPH_BUILDS.with(|builds| builds.set(builds.get() + 1)); let candidate_paths = candidates .into_iter() .flatten() @@ -71,18 +79,61 @@ impl CodexLineageGraph { if let Some(candidates) = candidates { candidate_node_indices.reserve(candidates.len()); for (candidate_index, candidate) in candidates.iter().enumerate() { - let uses_parent = candidate.session_metadata.lineage.uses_parent_baseline() - || candidate.session_metadata.forked_from_id.is_some(); + let cached = cache + .files + .get(&candidate.path.to_string_lossy().to_string()); + let cached_identity_matches = cached.is_some_and(|usage| { + let Ok(metadata) = fs::metadata(&candidate.path) else { + return false; + }; + let expected = usage.codex_file_identity.as_deref(); + let actual = JsonlScanner::codex_file_identity(&candidate.path, &metadata); + codex_file_identity_matches(expected, actual.as_deref()) + }); + let metadata_owns_identity = candidate.session_metadata.session_id.is_some(); + let uses_parent = if metadata_owns_identity { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + } else if cached_identity_matches { + cached.is_some_and(super::codex_usage_uses_parent) + } else { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + }; + let session_id = candidate.session_metadata.session_id.clone().or_else(|| { + cached_identity_matches + .then(|| cached.and_then(|usage| usage.codex_session_id.clone())) + .flatten() + }); + let parent_id = if metadata_owns_identity { + candidate.session_metadata.forked_from_id.clone() + } else { + candidate + .session_metadata + .forked_from_id + .clone() + .or_else(|| { + (cached_identity_matches && uses_parent) + .then(|| { + cached.and_then(|usage| usage.codex_forked_from_id.clone()) + }) + .flatten() + }) + }; + let cached_fallback = cached_identity_matches.then_some(cached).flatten(); nodes.push(CodexLineageNode { path: candidate.path.to_string_lossy().to_string(), - session_id: candidate.session_metadata.session_id.clone(), - parent_id: uses_parent - .then(|| candidate.session_metadata.forked_from_id.clone()) - .flatten(), + session_id, + parent_id: uses_parent.then_some(parent_id).flatten(), candidate_index: Some(candidate_index), - may_infer_missing_parent: candidate.session_metadata.is_subagent, - may_author_parent: true, - initially_unsafe: false, + may_infer_missing_parent: candidate.session_metadata.is_subagent + || cached_fallback.is_some_and(super::codex_fork_uses_local_inference), + may_author_parent: cached_fallback.is_none_or(|usage| { + !super::codex_fork_uses_local_inference(usage) + && !usage.codex_unresolved_fork_parent + }), + initially_unsafe: cached_fallback + .is_some_and(|usage| usage.codex_unresolved_fork_parent), }); candidate_node_indices.push(nodes.len() - 1); } @@ -135,28 +186,30 @@ impl CodexLineageGraph { // orders candidates. Cached-parent validation consumes the same gates. let mut completed = vec![false; nodes.len()]; let mut ordered_candidate_indices = Vec::with_capacity(candidate_node_indices.len()); - loop { - let mut progressed = false; - for index in 0..nodes.len() { - if completed[index] || gates[index] == CodexLineageGate::Unsafe { - continue; - } - let parent_is_ready = parent_indices[index].is_none_or(|parent_index| { - completed[parent_index] - && gates[parent_index] == CodexLineageGate::Eligible - && nodes[parent_index].may_author_parent - }); - if parent_is_ready { - completed[index] = true; - if let Some(candidate_index) = nodes[index].candidate_index { - ordered_candidate_indices.push(candidate_index); + let mut children = vec![Vec::new(); nodes.len()]; + let mut ready = VecDeque::new(); + for (index, parent) in parent_indices.iter().enumerate() { + match parent { + Some(parent_index) => children[*parent_index].push(index), + None if gates[index] == CodexLineageGate::Eligible => ready.push_back(index), + None => {} + } + } + while let Some(index) = ready.pop_front() { + if completed[index] || gates[index] == CodexLineageGate::Unsafe { + continue; + } + completed[index] = true; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_candidate_indices.push(candidate_index); + } + if nodes[index].may_author_parent { + for child in &children[index] { + if gates[*child] == CodexLineageGate::Eligible { + ready.push_back(*child); } - progressed = true; } } - if !progressed { - break; - } } for index in 0..nodes.len() { @@ -218,28 +271,69 @@ impl CodexLineageGraph { } } -pub(super) struct CodexLineagePlanner<'a> { - cache: &'a CostUsageCache, - graph: CodexLineageGraph, +pub(super) struct CodexLineagePlanner { + graph: Option, } -impl<'a> CodexLineagePlanner<'a> { - pub(super) fn new(cache: &'a CostUsageCache) -> Self { +impl CodexLineagePlanner { + pub(super) fn new(cache: &CostUsageCache) -> Self { Self { - cache, - graph: CodexLineageGraph::new(cache, None), + graph: Self::needs_graph(cache, None).then(|| CodexLineageGraph::new(cache, None)), } } pub(super) fn plan_candidates_by_lineage( cache: &CostUsageCache, candidates: &mut Vec, - ) -> Vec { - CodexLineageGraph::new(cache, Some(candidates)).apply_candidate_plan(candidates) + ) -> (Self, Vec) { + let graph = Self::needs_graph(cache, Some(candidates)) + .then(|| CodexLineageGraph::new(cache, Some(candidates))); + let unsafe_paths = graph + .as_ref() + .map_or_else(Vec::new, |graph| graph.apply_candidate_plan(candidates)); + (Self { graph }, unsafe_paths) + } + + fn needs_graph(cache: &CostUsageCache, candidates: Option<&[CodexPreparedCandidate]>) -> bool { + cache.files.values().any(super::codex_usage_uses_parent) + || candidates.is_some_and(|items| { + items.iter().any(|candidate| { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + }) + }) + } + + #[cfg(test)] + pub(crate) fn reset_graph_build_count() { + CODEX_LINEAGE_GRAPH_BUILDS.with(|count| count.set(0)); + } + #[cfg(test)] + pub(crate) fn graph_build_count() -> usize { + CODEX_LINEAGE_GRAPH_BUILDS.with(Cell::get) + } + + fn graph(&self) -> Option<&CodexLineageGraph> { + self.graph.as_ref() + } + + pub(super) fn cached_usage_is_safe( + &self, + cache: &CostUsageCache, + usage: &CostUsageFileUsage, + ) -> bool { + let locally_resolved = super::codex_fork_uses_local_inference(usage); + match self.decision_for_usage(cache, usage) { + CodexLineageDecision::Root => true, + CodexLineageDecision::ParentAbsent => locally_resolved, + CodexLineageDecision::ParentReady(_) => !locally_resolved, + CodexLineageDecision::Unsafe => false, + } } pub(super) fn decision_for_scan( &self, + cache: &CostUsageCache, uses_parent: bool, gate: CodexLineageGate, parent_id: Option<&str>, @@ -253,11 +347,15 @@ impl<'a> CodexLineagePlanner<'a> { return CodexLineageDecision::Root; } parent_id.map_or(CodexLineageDecision::Unsafe, |parent_id| { - self.resolve_parent(parent_id, fork_timestamp, parent_owner_expected) + self.resolve_parent(cache, parent_id, fork_timestamp, parent_owner_expected) }) } - pub(super) fn decision_for_usage(&self, usage: &CostUsageFileUsage) -> CodexLineageDecision { + pub(super) fn decision_for_usage( + &self, + cache: &CostUsageCache, + usage: &CostUsageFileUsage, + ) -> CodexLineageDecision { if usage.codex_unresolved_fork_parent { return CodexLineageDecision::Unsafe; } @@ -268,35 +366,34 @@ impl<'a> CodexLineagePlanner<'a> { .codex_forked_from_id .as_deref() .map_or(CodexLineageDecision::Unsafe, |parent_id| { - self.resolve_parent(parent_id, usage.codex_fork_timestamp.as_deref(), false) + self.resolve_parent( + cache, + parent_id, + usage.codex_fork_timestamp.as_deref(), + false, + ) }) } - pub(super) fn cached_usage_is_safe(&self, usage: &CostUsageFileUsage) -> bool { - let locally_resolved = super::codex_fork_uses_local_inference(usage); - match self.decision_for_usage(usage) { - CodexLineageDecision::Root => true, - CodexLineageDecision::ParentAbsent => locally_resolved, - CodexLineageDecision::ParentReady(_) => !locally_resolved, - CodexLineageDecision::Unsafe => false, - } - } - /// Resolve one parent identity through the persisted graph. Absence is /// distinct from ambiguity and transitive unsafety so local inference is /// allowed only when no owner exists at all. fn resolve_parent( &self, + cache: &CostUsageCache, parent_session_id: &str, child_fork_timestamp: Option<&str>, parent_owner_expected: bool, ) -> CodexLineageDecision { - let node_index = match self.graph.unique_owner(parent_session_id) { + let Some(graph) = self.graph() else { + return CodexLineageDecision::Unsafe; + }; + let node_index = match graph.unique_owner(parent_session_id) { Ok(None) if !parent_owner_expected => return CodexLineageDecision::ParentAbsent, Ok(None) | Err(()) => return CodexLineageDecision::Unsafe, Ok(Some(index)) => index, }; - self.parent_owner_baseline(node_index, child_fork_timestamp) + self.parent_owner_baseline(cache, node_index, child_fork_timestamp) .map_or( CodexLineageDecision::Unsafe, CodexLineageDecision::ParentReady, @@ -305,14 +402,16 @@ impl<'a> CodexLineagePlanner<'a> { fn parent_owner_baseline( &self, + cache: &CostUsageCache, node_index: usize, child_fork_timestamp: Option<&str>, ) -> Option { - let node = self.graph.nodes.get(node_index)?; - if self.graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { + let graph = self.graph()?; + let node = graph.nodes.get(node_index)?; + if graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { return None; } - let usage = self.cache.files.get(&node.path)?; + let usage = cache.files.get(&node.path)?; if usage.codex_unresolved_fork_parent || usage.codex_token_timestamps_monotonic != Some(true) || super::codex_fork_uses_local_inference(usage) @@ -326,9 +425,12 @@ impl<'a> CodexLineagePlanner<'a> { .as_ref()? .inherited_totals .as_ref()?; - let parent_index = self.graph.parent_indices[node_index]?; - let baseline = - self.parent_owner_baseline(parent_index, usage.codex_fork_timestamp.as_deref())?; + let parent_index = graph.parent_indices[node_index]?; + let baseline = self.parent_owner_baseline( + cache, + parent_index, + usage.codex_fork_timestamp.as_deref(), + )?; if &baseline != inherited { return None; } @@ -412,6 +514,7 @@ impl CodexLineageDecision { pub(super) fn cached_codex_file_is_fresh( cache: &CostUsageCache, + planner: &CodexLineagePlanner, entry: &CostUsageFileUsage, cache_covers_range: bool, mtime_unix_ms: i64, @@ -423,7 +526,7 @@ pub(super) fn cached_codex_file_is_fresh( && entry.size == size && codex_scan_target_size(entry) == size && entry.parsed_bytes.unwrap_or(0) >= size - && super::codex_fork_parent_is_safe(cache, entry) + && planner.cached_usage_is_safe(cache, entry) } pub(super) fn codex_file_identity_matches(expected: Option<&str>, actual: Option<&str>) -> bool { @@ -434,6 +537,7 @@ pub(super) fn codex_file_identity_matches(expected: Option<&str>, actual: Option pub(super) fn cached_codex_file_is_complete_for_range( cache: &CostUsageCache, + planner: &CodexLineagePlanner, path_key: &str, range: &CostUsageDayRange, ) -> bool { @@ -457,7 +561,7 @@ pub(super) fn cached_codex_file_is_complete_for_range( // Reconsider locally inferred children after this pass has // had a chance to discover and cache their parent. && !super::codex_fork_uses_local_inference(usage) - && super::codex_fork_parent_is_safe(cache, usage) + && planner.cached_usage_is_safe(cache, usage) }) } diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 3fb183eabb..7a60775d3c 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -177,8 +177,15 @@ pub(super) fn scan_codex_detailed_with_cache( cache.codex_pending_scan_root_paths = pending_scan.root_paths.clone(); cache.codex_pending_scan_timezone = Some(pending_scan.timezone.clone()); - let (mut candidates, discovery_complete) = - scanner.collect_codex_candidates(&sessions_dirs, scan_range, &cache, cancel, &mut stats); + let cached_lineage = CodexLineagePlanner::new(&cache); + let (mut candidates, discovery_complete) = scanner.collect_codex_candidates( + &sessions_dirs, + scan_range, + &cache, + &cached_lineage, + cancel, + &mut stats, + ); let candidate_limit = if scanner.options.codex_candidate_limit == 0 { usize::MAX } else { @@ -201,8 +208,9 @@ pub(super) fn scan_codex_detailed_with_cache( prioritize_codex_pending_candidates(&mut candidates, &pending_paths_before_pass); defer_codex_locally_inferred_candidates(&mut candidates, &cache); if discovery_complete && !is_cancelled(cancel) { - pending_next - .retain(|path| !cached_codex_file_is_complete_for_range(&cache, path, scan_range)); + pending_next.retain(|path| { + !cached_codex_file_is_complete_for_range(&cache, &cached_lineage, path, scan_range) + }); } // Admit one bounded set, inspect each admitted candidate once, and order @@ -234,12 +242,15 @@ pub(super) fn scan_codex_detailed_with_cache( }); } let mut unprocessed = Vec::new(); - if !cancelled_during_preparation.is_empty() || is_cancelled(cancel) { + let lineage_planner; + let cancelled_before_plan = !cancelled_during_preparation.is_empty() || is_cancelled(cancel); + if cancelled_before_plan { unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); unprocessed.extend(cancelled_during_preparation); } else { - let unsafe_cached_paths = + let (planner, unsafe_cached_paths) = CodexLineagePlanner::plan_candidates_by_lineage(&cache, &mut work_queue); + lineage_planner = planner; invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); if invalidated_unsafe_lineage { cache.previous_report = None; @@ -251,6 +262,9 @@ pub(super) fn scan_codex_detailed_with_cache( } } } + if cancelled_before_plan { + lineage_planner = cached_lineage; + } let mut incomplete_processed = Vec::new(); for (index, candidate) in work_queue.iter().enumerate() { @@ -274,6 +288,7 @@ pub(super) fn scan_codex_detailed_with_cache( &mut stats, Some(allowance), Some(candidate), + &lineage_planner, ); bytes_read_this_refresh = bytes_read_this_refresh.saturating_add(outcome.bytes_read.max(0)); stats.codex_bytes_read = stats diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index d8dc136047..9a930a3e7b 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -1449,10 +1449,16 @@ fn cost_scan_second_pass_skips_unchanged_files_via_cache() { // Second pass with default debounce still inspects files but skips re-parse. // Use app_driven so we exercise per-file mtime skip rather than whole-scan debounce. + CodexLineagePlanner::reset_graph_build_count(); let (summary2, stats2) = scanner.scan_codex_detailed(None); assert_eq!(stats2.files_seen, 2); assert_eq!(stats2.files_skipped, 2, "cache hit skips re-parse"); assert_eq!(stats2.files_parsed, 0); + assert_eq!( + CodexLineagePlanner::graph_build_count(), + 0, + "warm root-only scan must bypass lineage graph construction" + ); assert!(stats2.codex_metadata_read_paths.is_empty()); assert!(stats2.codex_history_read_paths.is_empty()); assert_eq!(stats2.codex_read_receipt, Default::default()); diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs index facfae31cf..3fbca1dbd1 100644 --- a/rust/src/cost_scanner/tests/lineage_cache.rs +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -167,7 +167,7 @@ fn replaced_parent_with_same_path_size_and_mtime_cannot_author_lineage() { let parent_key = parent.to_string_lossy().to_string(); let child_usage = &cache.files[&child.to_string_lossy().to_string()]; assert!(matches!( - CodexLineagePlanner::new(&cache).decision_for_usage(child_usage), + CodexLineagePlanner::new(&cache).decision_for_usage(&cache, child_usage), CodexLineageDecision::ParentReady(_) )); @@ -203,7 +203,7 @@ fn replaced_parent_with_same_path_size_and_mtime_cannot_author_lineage() { assert_ne!(replacement_identity, old_identity); assert_eq!( - CodexLineagePlanner::new(&cache).decision_for_usage(child_usage), + CodexLineagePlanner::new(&cache).decision_for_usage(&cache, child_usage), CodexLineageDecision::Unsafe ); } From bc1ac36c59d78ccf695694c5d82dbf8789f09ddb Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:41:31 +0700 Subject: [PATCH 49/62] Initialize lineage planner in one branch expression --- rust/src/cost_scanner/codex/scan.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 7a60775d3c..758f887e44 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -242,15 +242,14 @@ pub(super) fn scan_codex_detailed_with_cache( }); } let mut unprocessed = Vec::new(); - let lineage_planner; let cancelled_before_plan = !cancelled_during_preparation.is_empty() || is_cancelled(cancel); - if cancelled_before_plan { + let lineage_planner = if cancelled_before_plan { unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); unprocessed.extend(cancelled_during_preparation); + cached_lineage } else { let (planner, unsafe_cached_paths) = CodexLineagePlanner::plan_candidates_by_lineage(&cache, &mut work_queue); - lineage_planner = planner; invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); if invalidated_unsafe_lineage { cache.previous_report = None; @@ -261,10 +260,8 @@ pub(super) fn scan_codex_detailed_with_cache( pending_next.push(path); } } - } - if cancelled_before_plan { - lineage_planner = cached_lineage; - } + planner + }; let mut incomplete_processed = Vec::new(); for (index, candidate) in work_queue.iter().enumerate() { From b79e67ef0939190f6a96556cab63e7cbaac61cf8 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:45:02 +0700 Subject: [PATCH 50/62] Separate stacked preference enablement regression --- rust/src/settings/tests.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index 1ebf48fc2e..6ffec91c6a 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -1144,7 +1144,6 @@ fn provider_aliases_are_canonicalized_at_the_load_boundary() { "not-a-provider": "weekly" }, "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"], - "enabled_providers": ["claude"], "stacked_tray_top_provider": "OPENAI", "stacked_tray_bottom_provider": "ClAuDe" }"#, @@ -1166,6 +1165,21 @@ fn provider_aliases_are_canonicalized_at_the_load_boundary() { settings.stacked_tray_bottom_provider.as_deref(), Some("claude") ); +} + +#[test] +fn stacked_preferences_preserve_known_disabled_providers() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["claude"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "not-a-provider" + }"#, + ) + .expect("load stacked preferences independently of enablement"); + + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!(settings.stacked_tray_bottom_provider, None); assert_eq!( settings.enabled_providers, HashSet::from(["claude".to_string()]) From c186e197179c1a9234d7c12a122f7eefd5684774 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:50:47 +0700 Subject: [PATCH 51/62] Prefer current Codex lineage metadata --- rust/src/cost_scanner/codex/logical_target.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index a559588de8..567e1a44a9 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -120,7 +120,14 @@ impl CodexLineageGraph { .flatten() }) }; - let cached_fallback = cached_identity_matches.then_some(cached).flatten(); + // Freshly read identity-bearing metadata owns this candidate's + // current lineage state. Retain cached lineage flags only when + // the bounded metadata read could not establish an identity. + let cached_fallback = if metadata_owns_identity { + None + } else { + cached_identity_matches.then_some(cached).flatten() + }; nodes.push(CodexLineageNode { path: candidate.path.to_string_lossy().to_string(), session_id, From 2df956fa0d399de8bc41984e9501423b1972fa91 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:53:09 +0700 Subject: [PATCH 52/62] Allow proof harness to seed multiple providers --- apps/desktop-tauri/src-tauri/src/main.rs | 31 +-- .../src-tauri/src/proof_harness.rs | 177 ++++++++++++++++-- 2 files changed, 183 insertions(+), 25 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 80b8b37a34..27516787e9 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -132,17 +132,26 @@ fn main() { let mut initial_state = AppState::new(); initial_state.proof_config = proof_config; - // Proof-harness seed: CODEXBAR_SEED_USAGE_JSON plants one synthetic Codex - // ProviderUsageSnapshot before the event loop and any WebView read. The - // cache timestamp makes the seeded cache count as fresh so the first - // frontend refresh-if-stale call does not evict the synthetic data. - if let Some(snapshot) = proof_harness::seed_usage_snapshot_from_env() { - tracing::info!( - "proof-harness: seeded provider snapshot for '{}'", - snapshot.provider_id - ); - initial_state.provider_cache.push(snapshot); - initial_state.provider_cache_updated_at = Some(std::time::Instant::now()); + // Validate the complete proof seed before installing any snapshots, so an + // invalid multi-provider fixture cannot leave a partial cache behind. + if let Some(snapshots) = + proof_harness::seed_usage_snapshots_from_env(initial_state.proof_config.as_ref()) + { + let seeded_at = std::time::Instant::now(); + for snapshot in &snapshots { + tracing::info!( + "proof-harness: seeded provider snapshot for '{}'", + snapshot.provider_id + ); + if let Some(provider) = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id) + { + initial_state + .provider_cache_updated_at_by_provider + .insert(provider, seeded_at); + } + } + initial_state.provider_cache.extend(snapshots); + initial_state.provider_cache_updated_at = Some(seeded_at); } tauri::Builder::default() diff --git a/apps/desktop-tauri/src-tauri/src/proof_harness.rs b/apps/desktop-tauri/src-tauri/src/proof_harness.rs index 0830f8569f..7b5464e88e 100644 --- a/apps/desktop-tauri/src-tauri/src/proof_harness.rs +++ b/apps/desktop-tauri/src-tauri/src/proof_harness.rs @@ -16,12 +16,14 @@ //! and suppresses blur-dismiss so the window stays visible for automated //! screenshot capture. //! -//! `CODEXBAR_SEED_USAGE_JSON=` additionally seeds one synthetic, -//! bridge-shaped Codex [`ProviderUsageSnapshot`] into the provider cache at -//! launch (before the first event/WebView read) and pins it against refresh -//! eviction for the run. Malformed files log a warning and the shell -//! continues without seeding — proof runs must never crash on the seed. - +//! `CODEXBAR_SEED_USAGE_JSON=` additionally seeds either the legacy +//! synthetic Codex [`ProviderUsageSnapshot`] object or, in valid proof mode, +//! a nonempty array of unique supported provider snapshots. The validated set +//! is installed before the first event/WebView read and pinned against refresh +//! eviction. Malformed files log a warning and the shell continues without +//! seeding — proof runs must never crash on the seed. + +use std::collections::HashSet; use std::sync::Mutex; use serde::Serialize; @@ -256,22 +258,37 @@ pub fn is_proof_mode(app: &AppHandle) -> bool { // ── Provider-usage seed (CODEXBAR_SEED_USAGE_JSON) ─────────────────── /// Environment variable pointing at a JSON file with one synthetic, -/// bridge-shaped `ProviderUsageSnapshot` for the codex provider. +/// bridge-shaped Codex snapshot or a proof-only array of provider snapshots. pub const SEED_USAGE_ENV_VAR: &str = "CODEXBAR_SEED_USAGE_JSON"; /// Whether a seed path was configured at launch. While set, the provider /// cache is pinned fresh so the synthetic snapshot is never evicted by an /// automatic refresh during a proof/capture run. pub fn seed_usage_json_active() -> bool { - std::env::var_os(SEED_USAGE_ENV_VAR).is_some() + let Some(path) = std::env::var_os(SEED_USAGE_ENV_VAR) else { + return false; + }; + let Ok(raw) = std::fs::read_to_string(path) else { + return true; + }; + if raw.trim_start().starts_with('[') { + let proof_config = ProofConfig::from_env(); + proof_config.as_ref().is_some_and(is_valid_proof_config) + && parse_seed_usage_snapshots(&raw, proof_config.as_ref()).is_ok() + } else { + // Preserve the legacy object's existing pin behavior. + true + } } /// Read and validate the seed file referenced by `CODEXBAR_SEED_USAGE_JSON`. /// /// Returns `None` (with a warn, never a crash) when the variable is unset, -/// the file is unreadable, the JSON is malformed, or the snapshot is not -/// for the `codex` provider. -pub fn seed_usage_snapshot_from_env() -> Option { +/// the file is unreadable, or the seed does not satisfy the selected legacy +/// object or proof-only array contract. +pub fn seed_usage_snapshots_from_env( + proof_config: Option<&ProofConfig>, +) -> Option> { let path = std::env::var_os(SEED_USAGE_ENV_VAR)?; let path = std::path::PathBuf::from(path); let raw = match std::fs::read_to_string(&path) { @@ -284,8 +301,8 @@ pub fn seed_usage_snapshot_from_env() -> Option { return None; } }; - match parse_seed_usage_snapshot(&raw) { - Ok(snapshot) => Some(snapshot), + match parse_seed_usage_snapshots(&raw, proof_config) { + Ok(snapshots) => Some(snapshots), Err(msg) => { tracing::warn!("{SEED_USAGE_ENV_VAR}: {msg} in {}", path.display()); None @@ -315,6 +332,58 @@ pub fn parse_seed_usage_snapshot(json: &str) -> Result, +) -> Result, String> { + if !json.trim_start().starts_with('[') { + return parse_seed_usage_snapshot(json).map(|snapshot| vec![snapshot]); + } + if !proof_config.is_some_and(is_valid_proof_config) { + return Err("provider snapshot arrays require valid proof mode".into()); + } + + let values: Vec = + serde_json::from_str(json).map_err(|e| format!("malformed JSON: {e}"))?; + if values.is_empty() { + return Err("provider snapshot array must not be empty".into()); + } + + let mut providers = HashSet::with_capacity(values.len()); + let mut snapshots = Vec::with_capacity(values.len()); + for value in values { + if json_value_has_non_finite_number(&value) { + return Err("provider snapshot contains a non-finite number".into()); + } + let mut snapshot: ProviderUsageSnapshot = serde_json::from_value(value) + .map_err(|e| format!("malformed provider snapshot: {e}"))?; + if !is_supported_provider_id(&snapshot.provider_id) { + return Err(format!( + "unsupported snapshot providerId '{}', ignoring", + snapshot.provider_id + )); + } + if !providers.insert(snapshot.provider_id.clone()) { + return Err(format!( + "duplicate snapshot providerId '{}', ignoring", + snapshot.provider_id + )); + } + normalize_seed_snapshot(&mut snapshot); + snapshots.push(snapshot); + } + Ok(snapshots) +} + +fn normalize_seed_snapshot(snapshot: &mut ProviderUsageSnapshot) { normalize_rate_window(&mut snapshot.primary); snapshot.secondary.as_mut().map(normalize_rate_window); snapshot.model_specific.as_mut().map(normalize_rate_window); @@ -327,8 +396,17 @@ pub fn parse_seed_usage_snapshot(json: &str) -> Result bool { + match value { + serde_json::Value::Number(number) => { + number.as_f64().is_some_and(|value| !value.is_finite()) + } + serde_json::Value::Array(values) => values.iter().any(json_value_has_non_finite_number), + serde_json::Value::Object(values) => values.values().any(json_value_has_non_finite_number), + _ => false, + } } /// Recompute `remaining_percent` from `used_percent` (matching the canonical @@ -594,4 +672,75 @@ mod tests { assert_eq!(cost.period, "month"); assert_eq!(cost.formatted_used, "$12.50"); } + + fn valid_proof_config() -> ProofConfig { + ProofConfig { + target_surface: "trayPanel".into(), + settings_tab: None, + target_payload: None, + } + } + + fn seed_snapshot(provider_id: &str, used_percent: f64) -> serde_json::Value { + serde_json::json!({ + "providerId": provider_id, + "primary": { "usedPercent": used_percent, "windowMinutes": 300 } + }) + } + + #[test] + fn seed_snapshot_arrays_preserve_legacy_codex_object_behavior() { + let json = seed_snapshot("codex", 61.0).to_string(); + let snapshots = parse_seed_usage_snapshots(&json, None).expect("legacy seed parses"); + assert_eq!(snapshots.len(), 1); + assert_eq!(snapshots[0].provider_id, "codex"); + assert_eq!(snapshots[0].primary.remaining_percent, 39.0); + } + + #[test] + fn seed_snapshot_array_normalizes_multiple_supported_providers() { + let json = serde_json::json!([seed_snapshot("codex", 61.0), seed_snapshot("claude", 24.0)]) + .to_string(); + let snapshots = parse_seed_usage_snapshots(&json, Some(&valid_proof_config())) + .expect("supported snapshots parse in proof mode"); + + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0].provider_id, "codex"); + assert_eq!(snapshots[0].primary.remaining_percent, 39.0); + assert_eq!(snapshots[1].provider_id, "claude"); + assert_eq!(snapshots[1].primary.remaining_percent, 76.0); + assert!( + snapshots + .iter() + .all(|snapshot| !snapshot.updated_at.is_empty()) + ); + } + + #[test] + fn seed_snapshot_arrays_reject_empty_unknown_duplicate_and_non_finite_values() { + let proof = valid_proof_config(); + for json in [ + "[]".to_string(), + serde_json::json!([seed_snapshot("unknown-provider", 1.0)]).to_string(), + serde_json::json!([seed_snapshot("codex", 1.0), seed_snapshot("codex", 2.0)]) + .to_string(), + r#"[{"providerId":"codex","primary":{"usedPercent":1e400}}]"#.to_string(), + ] { + assert!( + parse_seed_usage_snapshots(&json, Some(&proof)).is_err(), + "{json}" + ); + } + } + + #[test] + fn seed_snapshot_arrays_are_rejected_without_valid_proof_config() { + let json = serde_json::json!([seed_snapshot("codex", 10.0)]).to_string(); + assert!(parse_seed_usage_snapshots(&json, None).is_err()); + } +} + +fn is_valid_proof_config(config: &ProofConfig) -> bool { + SurfaceMode::parse(&config.target_surface) + .is_some_and(|mode| proof_payload_is_supported(mode, config.target_payload.as_deref())) } From 1455c040ec9ea97dbf45b710911d087dc1056245 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:56:02 +0700 Subject: [PATCH 53/62] Keep validated proof seed state in memory --- .../src/commands/provider_refresh.rs | 34 +++++++++++- apps/desktop-tauri/src-tauri/src/main.rs | 1 + .../src-tauri/src/proof_harness.rs | 54 +++---------------- apps/desktop-tauri/src-tauri/src/state.rs | 3 ++ 4 files changed, 45 insertions(+), 47 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs index 80675e1eca..a92aa1bdc4 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs @@ -87,7 +87,7 @@ fn provider_cache_can_skip_refresh( .iter() .any(|snapshot| snapshot.provider_id == id.cli_name()) }); - if !force && crate::proof_harness::seed_usage_json_active() && cache_has_all { + if !force && state.provider_cache_seeded && cache_has_all { return true; } !force @@ -125,6 +125,38 @@ pub(super) fn complete_provider_refresh( mod tests { use super::*; + #[test] + fn validated_seed_pins_only_complete_nonforced_cache() { + let mut state = AppState::new(); + state.provider_cache.push( + crate::proof_harness::parse_seed_usage_snapshot( + r#"{"providerId":"codex","primary":{"usedPercent":25.0}}"#, + ) + .unwrap(), + ); + assert!(!provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex] + )); + state.provider_cache_seeded = true; + assert!(provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex] + )); + assert!(!provider_cache_can_skip_refresh( + &state, + true, + &[ProviderId::Codex] + )); + assert!(!provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex, ProviderId::Claude] + )); + } + #[test] fn stale_inputs_cannot_reserve_a_new_generation_after_invalidation() { let mut state = AppState::new(); diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 27516787e9..f7bbbef11f 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -151,6 +151,7 @@ fn main() { } } initial_state.provider_cache.extend(snapshots); + initial_state.provider_cache_seeded = true; initial_state.provider_cache_updated_at = Some(seeded_at); } diff --git a/apps/desktop-tauri/src-tauri/src/proof_harness.rs b/apps/desktop-tauri/src-tauri/src/proof_harness.rs index 7b5464e88e..4d05b39b78 100644 --- a/apps/desktop-tauri/src-tauri/src/proof_harness.rs +++ b/apps/desktop-tauri/src-tauri/src/proof_harness.rs @@ -261,26 +261,6 @@ pub fn is_proof_mode(app: &AppHandle) -> bool { /// bridge-shaped Codex snapshot or a proof-only array of provider snapshots. pub const SEED_USAGE_ENV_VAR: &str = "CODEXBAR_SEED_USAGE_JSON"; -/// Whether a seed path was configured at launch. While set, the provider -/// cache is pinned fresh so the synthetic snapshot is never evicted by an -/// automatic refresh during a proof/capture run. -pub fn seed_usage_json_active() -> bool { - let Some(path) = std::env::var_os(SEED_USAGE_ENV_VAR) else { - return false; - }; - let Ok(raw) = std::fs::read_to_string(path) else { - return true; - }; - if raw.trim_start().starts_with('[') { - let proof_config = ProofConfig::from_env(); - proof_config.as_ref().is_some_and(is_valid_proof_config) - && parse_seed_usage_snapshots(&raw, proof_config.as_ref()).is_ok() - } else { - // Preserve the legacy object's existing pin behavior. - true - } -} - /// Read and validate the seed file referenced by `CODEXBAR_SEED_USAGE_JSON`. /// /// Returns `None` (with a warn, never a crash) when the variable is unset, @@ -351,20 +331,14 @@ pub fn parse_seed_usage_snapshots( return Err("provider snapshot arrays require valid proof mode".into()); } - let values: Vec = + let mut snapshots: Vec = serde_json::from_str(json).map_err(|e| format!("malformed JSON: {e}"))?; - if values.is_empty() { + if snapshots.is_empty() { return Err("provider snapshot array must not be empty".into()); } - let mut providers = HashSet::with_capacity(values.len()); - let mut snapshots = Vec::with_capacity(values.len()); - for value in values { - if json_value_has_non_finite_number(&value) { - return Err("provider snapshot contains a non-finite number".into()); - } - let mut snapshot: ProviderUsageSnapshot = serde_json::from_value(value) - .map_err(|e| format!("malformed provider snapshot: {e}"))?; + let mut providers = HashSet::with_capacity(snapshots.len()); + for snapshot in &mut snapshots { if !is_supported_provider_id(&snapshot.provider_id) { return Err(format!( "unsupported snapshot providerId '{}', ignoring", @@ -377,8 +351,7 @@ pub fn parse_seed_usage_snapshots( snapshot.provider_id )); } - normalize_seed_snapshot(&mut snapshot); - snapshots.push(snapshot); + normalize_seed_snapshot(snapshot); } Ok(snapshots) } @@ -398,15 +371,9 @@ fn normalize_seed_snapshot(snapshot: &mut ProviderUsageSnapshot) { } } -fn json_value_has_non_finite_number(value: &serde_json::Value) -> bool { - match value { - serde_json::Value::Number(number) => { - number.as_f64().is_some_and(|value| !value.is_finite()) - } - serde_json::Value::Array(values) => values.iter().any(json_value_has_non_finite_number), - serde_json::Value::Object(values) => values.values().any(json_value_has_non_finite_number), - _ => false, - } +fn is_valid_proof_config(config: &ProofConfig) -> bool { + SurfaceMode::parse(&config.target_surface) + .is_some_and(|mode| proof_payload_is_supported(mode, config.target_payload.as_deref())) } /// Recompute `remaining_percent` from `used_percent` (matching the canonical @@ -739,8 +706,3 @@ mod tests { assert!(parse_seed_usage_snapshots(&json, None).is_err()); } } - -fn is_valid_proof_config(config: &ProofConfig) -> bool { - SurfaceMode::parse(&config.target_surface) - .is_some_and(|mode| proof_payload_is_supported(mode, config.target_payload.as_deref())) -} diff --git a/apps/desktop-tauri/src-tauri/src/state.rs b/apps/desktop-tauri/src-tauri/src/state.rs index 10d626ec43..174c13d113 100644 --- a/apps/desktop-tauri/src-tauri/src/state.rs +++ b/apps/desktop-tauri/src-tauri/src/state.rs @@ -143,6 +143,8 @@ pub struct AppState { pub installer_path: Option, /// Proof-harness configuration (set when `CODEXBAR_PROOF_MODE` is active). pub proof_config: Option, + /// True only after a validated proof seed was installed at startup. + pub provider_cache_seeded: bool, /// Persistent notification manager — tracks which alerts have fired to prevent spam. pub notification_manager: codexbar::notifications::NotificationManager, /// Instant when the tray panel was last shown — used to suppress @@ -204,6 +206,7 @@ impl AppState { last_update_check_ms: None, installer_path: None, proof_config: None, + provider_cache_seeded: false, notification_manager: codexbar::notifications::NotificationManager::new(), last_shown_at: None, last_blur_dismissed_at: None, From 5d2d39efe9803efbba24b265b516335984f345bd Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:04:11 +0700 Subject: [PATCH 54/62] Scope lineage invalidation to active dependencies --- rust/src/cost_scanner/codex/logical_target.rs | 75 +++++++++++++------ rust/src/cost_scanner/codex/scan.rs | 8 +- rust/src/cost_scanner/tests/lineage_cache.rs | 68 +++++++++++++++++ 3 files changed, 128 insertions(+), 23 deletions(-) diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 567e1a44a9..b261f174fe 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -246,34 +246,65 @@ impl CodexLineageGraph { } } - fn apply_candidate_plan(&self, candidates: &mut Vec) -> Vec { - if candidates.is_empty() { - return Vec::new(); - } - for (candidate_index, candidate) in candidates.iter_mut().enumerate() { - let node_index = self.candidate_node_indices[candidate_index]; - candidate.lineage_gate = self.gates[node_index]; - candidate.parent_owner_expected = self.parent_indices[node_index].is_some(); + fn apply_candidate_plan( + &self, + candidates: &mut Vec, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, + ) -> Vec { + if !candidates.is_empty() { + for (candidate_index, candidate) in candidates.iter_mut().enumerate() { + let node_index = self.candidate_node_indices[candidate_index]; + candidate.lineage_gate = self.gates[node_index]; + candidate.parent_owner_expected = self.parent_indices[node_index].is_some(); + } + + let mut remaining = candidates.drain(..).map(Some).collect::>(); + for candidate_index in &self.ordered_candidate_indices { + candidates.push( + remaining[*candidate_index] + .take() + .expect("candidate is ordered once"), + ); + } } - let mut remaining = candidates.drain(..).map(Some).collect::>(); - for candidate_index in &self.ordered_candidate_indices { - candidates.push( - remaining[*candidate_index] - .take() - .expect("candidate is ordered once"), - ); + // Keep the complete graph for parent resolution, but invalidate only + // unsafe cached nodes in the active range or in the ancestor closure + // required to resolve an active node. + let mut relevant = vec![false; self.nodes.len()]; + let mut pending = VecDeque::new(); + for (index, node) in self.nodes.iter().enumerate() { + if super::is_codex_path_in_scan_window(Path::new(&node.path), sessions_dirs, range) { + relevant[index] = true; + pending.push_back(index); + } + } + while let Some(index) = pending.pop_front() { + let Some(parent_id) = self.nodes[index].parent_id.as_deref() else { + continue; + }; + if let Some(owners) = self.session_owners.get(parent_id) { + for &owner in owners { + if !relevant[owner] { + relevant[owner] = true; + pending.push_back(owner); + } + } + } } self.nodes .iter() .zip(&self.gates) - .filter(|(node, gate)| { - node.candidate_index.is_none() + .enumerate() + .filter(|(index, (node, gate))| { + relevant[*index] + && node.candidate_index.is_none() && **gate == CodexLineageGate::Unsafe && !node.initially_unsafe }) - .map(|(node, _)| node.path.clone()) + .map(|(_, (node, _))| node.path.clone()) .collect() } } @@ -292,12 +323,14 @@ impl CodexLineagePlanner { pub(super) fn plan_candidates_by_lineage( cache: &CostUsageCache, candidates: &mut Vec, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, ) -> (Self, Vec) { let graph = Self::needs_graph(cache, Some(candidates)) .then(|| CodexLineageGraph::new(cache, Some(candidates))); - let unsafe_paths = graph - .as_ref() - .map_or_else(Vec::new, |graph| graph.apply_candidate_plan(candidates)); + let unsafe_paths = graph.as_ref().map_or_else(Vec::new, |graph| { + graph.apply_candidate_plan(candidates, sessions_dirs, range) + }); (Self { graph }, unsafe_paths) } diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index 758f887e44..94cb91f701 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -248,8 +248,12 @@ pub(super) fn scan_codex_detailed_with_cache( unprocessed.extend(cancelled_during_preparation); cached_lineage } else { - let (planner, unsafe_cached_paths) = - CodexLineagePlanner::plan_candidates_by_lineage(&cache, &mut work_queue); + let (planner, unsafe_cached_paths) = CodexLineagePlanner::plan_candidates_by_lineage( + &cache, + &mut work_queue, + &sessions_dirs, + scan_range, + ); invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); if invalidated_unsafe_lineage { cache.previous_report = None; diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs index 3fbca1dbd1..19ac317153 100644 --- a/rust/src/cost_scanner/tests/lineage_cache.rs +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -437,3 +437,71 @@ fn bounded_refresh_rejects_dependent_of_locally_inferred_parent() { assert_locally_inferred(&cache, &parent); assert_unresolved(&cache, &dependent); } + +#[test] +fn current_refresh_scopes_unsafe_cache_invalidation_to_range_and_dependencies() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let active_time = Utc::now() - Duration::hours(1); + let old_date = Local::now().date_naive() - Duration::days(30); + let old_day = old_date.format("%Y-%m-%d").to_string(); + let old_dir = sessions + .join(old_date.format("%Y").to_string()) + .join(old_date.format("%m").to_string()) + .join(old_date.format("%d").to_string()); + let mut cache = CostUsageCache::default(); + + { + let mut add_cached = |name: &str, session_id: &str, parent_id: Option<&str>| { + let path = old_dir.join(name).to_string_lossy().to_string(); + let mut usage = cached_usage_with_packed(&old_day, "gpt-5.6-sol", vec![100, 0, 5, 0]); + usage.codex_session_id = Some(session_id.to_string()); + usage.codex_forked_from_id = parent_id.map(str::to_string); + cache.files.insert(path, usage); + }; + add_cached("unrelated-a.jsonl", "unrelated-a", Some("unrelated-b")); + add_cached("unrelated-b.jsonl", "unrelated-b", Some("unrelated-a")); + add_cached("required-a.jsonl", "required-parent", None); + add_cached("required-b.jsonl", "required-parent", None); + } + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + + let active_child = write_codex_fork_session_fixture( + &sessions, + "active-child.jsonl", + "active-child", + Some("required-parent"), + active_time, + active_time + Duration::seconds(1), + &[1_000_000, 1_000_140], + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, refreshed) = scanner.scan_codex_detailed_with_cache(None); + let cached_path = |name: &str| old_dir.join(name).to_string_lossy().to_string(); + + for path in [ + cached_path("unrelated-a.jsonl"), + cached_path("unrelated-b.jsonl"), + ] { + let usage = refreshed + .files + .get(&path) + .expect("unrelated history retained"); + assert_eq!(usage.days[&old_day]["gpt-5.6-sol"], vec![100, 0, 5, 0]); + assert!(!usage.codex_unresolved_fork_parent); + } + for path in [ + cached_path("required-a.jsonl"), + cached_path("required-b.jsonl"), + ] { + assert_unresolved(&refreshed, Path::new(&path)); + } + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&refreshed, &active_child); +} From 6eec2e2f1f012fcf850c7285d56d742f92fc2891 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:46:46 +0700 Subject: [PATCH 55/62] Fall back to API key dashboard URLs --- .../src-tauri/src/commands/mod.rs | 16 +++++++++++++++- .../src-tauri/src/commands/tests.rs | 18 ++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index 06721a6e8e..6d54bcf9e3 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -116,13 +116,27 @@ fn provider_dashboard_url(id: ProviderId, settings: &Settings) -> Option .console_url() .to_string(), ), - _ => instantiate_provider(id) + ProviderId::OpenRouter => instantiate_provider(id) .metadata() .dashboard_url .map(str::to_string), + _ => provider_dashboard_url_from_sources( + instantiate_provider(id).metadata().dashboard_url, + codexbar::settings::api_keys::get_api_key_providers() + .into_iter() + .find(|provider| provider.id == id) + .and_then(|provider| provider.dashboard_url), + ), } } +fn provider_dashboard_url_from_sources( + metadata_url: Option<&'static str>, + api_key_catalog_url: Option<&'static str>, +) -> Option { + metadata_url.or(api_key_catalog_url).map(str::to_string) +} + fn validate_single_line_secret(value: &str, field: &str, max_len: usize) -> Result<(), String> { let trimmed = value.trim(); if trimmed.is_empty() { diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 685a2b2cb7..94ac348f11 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -720,6 +720,24 @@ fn provider_dashboard_url_uses_selected_regional_console() { ); } +#[test] +fn provider_dashboard_url_falls_back_to_api_key_catalog() { + assert_eq!( + super::provider_dashboard_url_from_sources(None, Some("https://catalog.example/dashboard")) + .as_deref(), + Some("https://catalog.example/dashboard") + ); + assert_eq!( + super::provider_dashboard_url_from_sources( + Some("https://metadata.example/dashboard"), + Some("https://catalog.example/dashboard"), + ) + .as_deref(), + Some("https://metadata.example/dashboard") + ); + assert_eq!(super::provider_dashboard_url_from_sources(None, None), None); +} + #[test] fn fetch_context_token_account_uses_web_cookie_header() { let settings = Settings::default(); From da5ce398e15f39a2db143840ff4a7dc64e7ba8ae Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:46:39 +0700 Subject: [PATCH 56/62] Compile Antigravity offline count helper only for tests --- rust/src/providers/antigravity/local_history.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs index 055f9b8d97..00519d93ca 100644 --- a/rust/src/providers/antigravity/local_history.rs +++ b/rust/src/providers/antigravity/local_history.rs @@ -63,6 +63,7 @@ pub fn offline_conversation_count() -> usize { offline_conversation_count_with_roots(&roots, &tokscale_sessions) } +#[cfg(test)] fn offline_conversation_count_in(home: &Path) -> usize { let roots = local_sqlite::database_roots(&home.join(".gemini")); let tokscale_sessions = local_sessions::tokscale_sessions_from_values(home, None); From e354f3f51d51a7e5fd6f85ae5f6cf7e5bf1002d1 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:52:37 +0700 Subject: [PATCH 57/62] Reinfer grown subagent history from start --- rust/src/cost_scanner/codex/logical_target.rs | 8 ++---- rust/src/cost_scanner/tests/lineage_cache.rs | 28 +++++++++++++++++++ 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index b261f174fe..dd42b96a98 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -526,18 +526,14 @@ impl CodexLineageDecision { } } Self::ParentAbsent => { - if let Some(state) = matching_cached_state + if let Some(state) = matching_cached_state.filter(|state| !state.locally_resolved) && let Some(baseline) = state.inherited_totals.clone() { return CodexAccountingMode::Baseline { baseline, paginated_continuation, remaining_inherited_totals: state.remaining_inherited_totals.clone(), - provenance: if state.locally_resolved { - CodexBaselineProvenance::CachedLocalInference - } else { - CodexBaselineProvenance::CachedValidatedParent - }, + provenance: CodexBaselineProvenance::CachedValidatedParent, }; } if metadata.is_subagent { diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs index 19ac317153..124733b098 100644 --- a/rust/src/cost_scanner/tests/lineage_cache.rs +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -137,6 +137,34 @@ fn assert_unresolved(cache: &CostUsageCache, path: &Path) { assert!(usage.codex_fork_accounting_state.is_none()); } +#[test] +fn appended_owned_token_row_reinfers_locally_resolved_subagent_from_start() { + use std::io::Write as _; + + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent(&sessions, "child.jsonl", "child-id", "missing-parent", base); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (initial, _, initial_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(initial.input_tokens, 50); + assert_locally_inferred(&initial_cache, &child); + + let appended_owned_row = lineage_token_row(base + Duration::seconds(2), 22, 1_100, 50); + std::fs::OpenOptions::new() + .append(true) + .open(&child) + .unwrap() + .write_all(format!("{appended_owned_row}\n").as_bytes()) + .unwrap(); + + let (grown, _, grown_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(grown.input_tokens, 100); + assert_locally_inferred(&grown_cache, &child); +} + #[test] fn replaced_parent_with_same_path_size_and_mtime_cannot_author_lineage() { let root = tempfile::tempdir().unwrap(); From aca5988accdd121f8a9065f540c9830dcb5e353f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:56:40 +0700 Subject: [PATCH 58/62] Preserve unknown Claude daily costs --- rust/src/cost_scanner.rs | 48 +++++++++++++++++-------- rust/src/cost_scanner/tests.rs | 66 ++++++++++++++++++++++++++++++++-- 2 files changed, 97 insertions(+), 17 deletions(-) diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 3c7fe4a06d..55444c6a74 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -673,6 +673,7 @@ impl CostScanner { }; let mut daily_cost = HashMap::new(); let mut daily_tokens = HashMap::new(); + let mut unknown_cost_dates = HashSet::new(); for days_ago in 0..self.days { let date = today - Duration::days(days_ago as i64); let key = date.format("%Y-%m-%d").to_string(); @@ -700,8 +701,11 @@ impl CostScanner { aggregation_complete &= record.timestamp.is_some(); aggregation_complete &= add_claude_record_to_summary(&mut summary, record); - aggregation_complete &= - add_claude_record_to_daily_costs(&mut daily_cost, record); + aggregation_complete &= add_claude_record_to_daily_costs( + &mut daily_cost, + &mut unknown_cost_dates, + record, + ); aggregation_complete &= add_claude_record_to_daily_tokens(&mut daily_tokens, record); if let Some(quota_record) = quota_history_record_from_usage(record) { @@ -739,11 +743,7 @@ impl CostScanner { is_cancelled(cancel), ); if complete { - for value in daily_cost.values_mut() { - if value.is_none() { - *value = Some(0.0); - } - } + zero_fill_uninitialized_claude_daily_costs(&mut daily_cost, &unknown_cost_dates); } let mut daily_cost = daily_cost.into_iter().collect::>(); @@ -1151,6 +1151,7 @@ fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option>, + unknown_cost_dates: &mut HashSet, record: &ClaudeUsageRecord, ) -> bool { let Some(timestamp) = record.timestamp else { @@ -1162,13 +1163,18 @@ fn add_claude_record_to_daily_costs( .format("%Y-%m-%d") .to_string(); if let Some(cost) = daily_costs.get_mut(&date_str) { + if unknown_cost_dates.contains(&date_str) { + return false; + } let Some(record_cost) = record.cost else { *cost = None; + unknown_cost_dates.insert(date_str); return false; }; let sum = cost.unwrap_or(0.0) + record_cost; if !sum.is_finite() { *cost = None; + unknown_cost_dates.insert(date_str); return false; } *cost = Some(sum); @@ -1176,6 +1182,17 @@ fn add_claude_record_to_daily_costs( true } +fn zero_fill_uninitialized_claude_daily_costs( + daily_costs: &mut HashMap>, + unknown_cost_dates: &HashSet, +) { + for (day, cost) in daily_costs { + if cost.is_none() && !unknown_cost_dates.contains(day) { + *cost = Some(0.0); + } + } +} + /// Check if any cost usage sources are available #[allow( dead_code, @@ -1256,6 +1273,7 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let mut claude_scan = ClaudeFileScanResult::default(); + let mut unknown_cost_dates = HashSet::new(); let traversal_read_failures = { let mut handle_file = |path: &Path| { let mut aggregation_complete = true; @@ -1267,8 +1285,11 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< &mut pricing, |record| { aggregation_complete &= record.timestamp.is_some(); - aggregation_complete &= - add_claude_record_to_daily_costs(&mut daily_costs, record); + aggregation_complete &= add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + record, + ); }, ); if !aggregation_complete { @@ -1283,11 +1304,10 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< .read_failures .saturating_add(traversal_read_failures); if claude_scan.is_complete() { - for slot in daily_costs.values_mut() { - if slot.is_none() { - *slot = Some(0.0); - } - } + zero_fill_uninitialized_claude_daily_costs( + &mut daily_costs, + &unknown_cost_dates, + ); } } } diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 4c2617f275..69f0c8dffa 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -865,14 +865,15 @@ fn daily_history_dedups_across_files_and_buckets_by_local_day() { .to_string() }; let mut daily_costs = HashMap::new(); - daily_costs.insert(day_key(&day_one), Some(0.0)); - daily_costs.insert(day_key(&day_two), Some(0.0)); + daily_costs.insert(day_key(&day_one), None); + daily_costs.insert(day_key(&day_two), None); + let mut unknown_cost_dates = HashSet::new(); let cutoff = Utc::now() - Duration::days(30); let mut seen = HashSet::new(); for path in [&file_a, &file_b] { for_each_claude_usage_record(path, &cutoff, &mut seen, None, |record| { - add_claude_record_to_daily_costs(&mut daily_costs, record); + add_claude_record_to_daily_costs(&mut daily_costs, &mut unknown_cost_dates, record); }); } @@ -891,6 +892,65 @@ fn daily_history_dedups_across_files_and_buckets_by_local_day() { let _removed_b = std::fs::remove_file(&file_b); } +fn claude_daily_cost_record(timestamp: DateTime, cost: Option) -> ClaudeUsageRecord { + ClaudeUsageRecord { + model: "claude-test".to_string(), + pricing_known: cost.is_some(), + timestamp: Some(timestamp), + dedup_key: None, + input: 1, + output: 1, + cache_create: 0, + cache_read: 0, + cost, + } +} + +#[test] +fn unknown_claude_cost_date_cannot_be_restored_by_later_priced_record() { + let timestamp = Utc::now(); + let day = timestamp + .with_timezone(&Local) + .date_naive() + .format("%Y-%m-%d") + .to_string(); + let mut daily_costs = HashMap::from([(day.clone(), None)]); + let mut unknown_cost_dates = HashSet::new(); + + assert!(add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, Some(0.75)), + )); + assert!(!add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, None), + )); + assert!(!add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, Some(1.25)), + )); + + assert_eq!(daily_costs[&day], None); + assert!(unknown_cost_dates.contains(&day)); +} + +#[test] +fn claude_daily_zero_fill_preserves_unknown_dates_and_fills_untouched_dates() { + let unknown_day = "2026-09-22".to_string(); + let untouched_day = "2026-09-23".to_string(); + let mut daily_costs = + HashMap::from([(unknown_day.clone(), None), (untouched_day.clone(), None)]); + let unknown_cost_dates = HashSet::from([unknown_day.clone()]); + + zero_fill_uninitialized_claude_daily_costs(&mut daily_costs, &unknown_cost_dates); + + assert_eq!(daily_costs[&unknown_day], None); + assert_eq!(daily_costs[&untouched_day], Some(0.0)); +} + #[test] fn claude_scan_counts_final_incomplete_jsonl_line() { let path = From a9af9ff4e683d273b0ca97aa8f6cb7162a60b90e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:15:22 +0700 Subject: [PATCH 59/62] Use public API key catalog export --- apps/desktop-tauri/src-tauri/src/commands/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index 6d54bcf9e3..039dc0ca10 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -122,7 +122,7 @@ fn provider_dashboard_url(id: ProviderId, settings: &Settings) -> Option .map(str::to_string), _ => provider_dashboard_url_from_sources( instantiate_provider(id).metadata().dashboard_url, - codexbar::settings::api_keys::get_api_key_providers() + codexbar::settings::get_api_key_providers() .into_iter() .find(|provider| provider.id == id) .and_then(|provider| provider.dashboard_url), From 63e243d0c0abc891acecf4f6d86964ae1d555d14 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:15:32 +0700 Subject: [PATCH 60/62] Remove obsolete cached local provenance --- rust/src/cost_scanner/codex.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 23730eb28f..6c0ebc5efa 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -32,7 +32,6 @@ enum CodexAccountingMode { enum CodexBaselineProvenance { ValidatedParent { replaces_cached_state: bool }, CachedValidatedParent, - CachedLocalInference, } impl CodexAccountingMode { @@ -44,16 +43,6 @@ impl CodexAccountingMode { matches!(self, Self::InferSubagent { .. }) } - fn locally_resolved(&self) -> bool { - matches!( - self, - Self::Baseline { - provenance: CodexBaselineProvenance::CachedLocalInference, - .. - } - ) - } - fn requires_cached_reparse(&self) -> bool { matches!( self, @@ -747,8 +736,7 @@ impl CostScanner { bytes_read: parse_result.bytes_read, is_complete: parse_result.is_complete, }; - let locally_resolved = - accounting_mode.locally_resolved() || parse_result.fork_baseline_locally_resolved; + let locally_resolved = parse_result.fork_baseline_locally_resolved; let codex_fork_accounting_state = if is_fork && (parse_result.fork_baseline.is_some() || parse_result.fork_baseline_locally_resolved) { From 08ad1985ceea7b0bd44bf0f61a023fe18832089d Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:15:42 +0700 Subject: [PATCH 61/62] Clarify Kimi manual cookie error --- rust/src/providers/kimi/web.rs | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 020854b6f4..c5ed5b8305 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -32,6 +32,15 @@ fn browser_import_allowed(cookie_source: &str) -> bool { cookie_source.eq_ignore_ascii_case("auto") || cookie_source.eq_ignore_ascii_case("browser") } +fn browser_import_error(cookie_source: &str) -> ProviderError { + let message = if cookie_source.eq_ignore_ascii_case("manual") { + "Kimi cookie source is Manual; provide a valid manual cookie header." + } else { + "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + }; + ProviderError::Other(message.into()) +} + /// Web auth token chain for both the web fetch and the Code-API enrichment /// (upstream `KimiWebEnrichmentTokenResolver.resolve`): /// 1. Manual cookie header (its `kimi-auth`/auth cookie), source-independent. @@ -135,10 +144,7 @@ pub(crate) async fn fetch_via_web( } if !browser_import_allowed(&source) { - return Err(ProviderError::Other( - "Kimi cookie source is Off; provide a manual cookie header or enable browser import." - .into(), - )); + return Err(browser_import_error(&source)); } let client = client()?; @@ -442,6 +448,25 @@ mod tests { assert!(!browser_import_allowed("manual")); } + #[test] + fn browser_import_error_matches_rejected_cookie_source() { + assert!(matches!( + browser_import_error("manual"), + ProviderError::Other(message) + if message == "Kimi cookie source is Manual; provide a valid manual cookie header." + )); + assert!(matches!( + browser_import_error("off"), + ProviderError::Other(message) + if message == "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + )); + assert!(matches!( + browser_import_error("unexpected"), + ProviderError::Other(message) + if message == "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + )); + } + #[test] fn subscription_stats_do_not_invent_a_membership_label() { let usage: KimiWebUsageResponse = serde_json::from_value(serde_json::json!({ From 2aa7ff59ec15ef608f0d30b9063e9721d37bd337 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Fri, 25 Sep 2026 01:50:42 +0700 Subject: [PATCH 62/62] Add isolated no-focus Antigravity proof mode --- .../src-tauri/src/commands/bridge.rs | 44 +- .../src-tauri/src/commands/locale_cmd.rs | 34 +- .../src-tauri/src/commands/providers.rs | 37 +- .../src-tauri/src/commands/settings.rs | 8 + .../src-tauri/src/commands/spend_contract.rs | 14 + .../src-tauri/src/commands/updater.rs | 37 ++ .../src-tauri/src/commands/usage_spend.rs | 341 ++++++++++- apps/desktop-tauri/src-tauri/src/main.rs | 140 +++-- .../src-tauri/src/proof_harness.rs | 37 ++ .../src-tauri/src/proof_runtime.rs | 569 ++++++++++++++++++ apps/desktop-tauri/src-tauri/src/shell/mod.rs | 4 +- .../src-tauri/src/shell/transition.rs | 92 ++- .../src-tauri/src/shell/window.rs | 29 +- apps/desktop-tauri/src-tauri/src/state.rs | 22 + .../src-tauri/src/test_support.rs | 77 +++ .../src-tauri/src/tray_bridge.rs | 10 +- rust/src/logging.rs | 21 +- .../providers/antigravity/local_history.rs | 15 +- .../providers/antigravity/local_sessions.rs | 66 ++ 19 files changed, 1497 insertions(+), 100 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/proof_runtime.rs create mode 100644 apps/desktop-tauri/src-tauri/src/test_support.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 96955cf740..a69be60f1c 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -721,9 +721,7 @@ pub struct SettingsSnapshot { provider_accent_colors: std::collections::HashMap, } -#[tauri::command] -pub fn get_bootstrap_state() -> BootstrapState { - let settings = Settings::load(); +fn bootstrap_state_from_settings(settings: Settings) -> BootstrapState { BootstrapState { contract_version: "v1", providers: provider_catalog_for(&settings), @@ -731,14 +729,48 @@ pub fn get_bootstrap_state() -> BootstrapState { } } +#[cfg(not(test))] #[tauri::command] -pub fn get_provider_catalog() -> Vec { +pub fn get_bootstrap_state( + state: tauri::State<'_, Mutex>, +) -> Result { + Ok(bootstrap_state_from_settings(settings_for_command(&state)?)) +} + +#[cfg(test)] +#[tauri::command] +pub fn get_bootstrap_state() -> BootstrapState { + bootstrap_state_from_settings(Settings::default()) +} + +#[tauri::command] +pub fn get_provider_catalog( + state: tauri::State<'_, Mutex>, +) -> Result, String> { + Ok(provider_catalog_for(&settings_for_command(&state)?)) +} + +pub fn get_provider_catalog_for_current_settings() -> Vec { provider_catalog_for(&Settings::load()) } #[tauri::command] -pub fn get_settings_snapshot() -> SettingsSnapshot { - SettingsSnapshot::from(Settings::load()) +pub fn get_settings_snapshot( + state: tauri::State<'_, Mutex>, +) -> Result { + Ok(SettingsSnapshot::from(settings_for_command(&state)?)) +} + +fn settings_for_command(state: &tauri::State<'_, Mutex>) -> Result { + let guard = state.lock().map_err(|error| error.to_string())?; + if guard.is_containment_proof() { + guard + .proof_settings() + .cloned() + .ok_or_else(|| "containment proof settings are unavailable".to_string()) + } else { + Ok(Settings::load()) + } } impl From for SettingsSnapshot { diff --git a/apps/desktop-tauri/src-tauri/src/commands/locale_cmd.rs b/apps/desktop-tauri/src-tauri/src/commands/locale_cmd.rs index a71224de8f..b80e0501ca 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/locale_cmd.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/locale_cmd.rs @@ -54,11 +54,22 @@ fn locale_strings_for(lang: Language) -> LocaleStrings { /// the persisted label (`"english"`, `"chinese"`), or the full name /// (`"English"`, `"Chinese"`, `"中文"`). #[tauri::command] -pub fn get_locale_strings(language: Option) -> Result { - let lang = match language.as_deref() { - None => locale::current_language(), - Some(raw) => { - parse_locale_language(raw).ok_or_else(|| format!("unknown language code: {raw}"))? +pub fn get_locale_strings( + state: tauri::State<'_, Mutex>, + language: Option, +) -> Result { + let proof_mode = state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true); + let lang = if proof_mode { + Language::English + } else { + match language.as_deref() { + None => locale::current_language(), + Some(raw) => { + parse_locale_language(raw).ok_or_else(|| format!("unknown language code: {raw}"))? + } } }; Ok(locale_strings_for(lang)) @@ -71,7 +82,18 @@ fn parse_locale_language(raw: &str) -> Option { /// Persist the UI language and emit a `locale-changed` event so the /// frontend can refetch its locale table without a restart. #[tauri::command] -pub fn set_ui_language(app: tauri::AppHandle, language: String) -> Result<(), String> { +pub fn set_ui_language( + app: tauri::AppHandle, + state: tauri::State<'_, Mutex>, + language: String, +) -> Result<(), String> { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return Err("settings mutations disabled in containment proof mode".to_string()); + } let lang = parse_locale_language(&language).ok_or_else(|| format!("unknown language: {language}"))?; let mut settings = Settings::load(); diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..779f9c801c 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -10,6 +10,7 @@ use serde::Serialize; use std::sync::Arc; const MAX_CONCURRENT_PROVIDER_FETCHES: usize = 8; +const PROOF_REFRESH_DISABLED: &str = "provider refresh disabled in containment proof mode"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RefreshScope { @@ -379,10 +380,13 @@ async fn do_refresh_providers_with_policy( scope: RefreshScope, ) -> Result { let state = app.state::>(); - let expected_generation = state - .lock() - .map_err(|e| e.to_string())? - .provider_refresh_generation; + let expected_generation = { + let guard = state.lock().map_err(|e| e.to_string())?; + if guard.is_containment_proof() { + return Err(PROOF_REFRESH_DISABLED.to_string()); + } + guard.provider_refresh_generation + }; let settings = Settings::load(); let enabled_ids = settings.get_enabled_provider_ids(); let refresh_ids = scope.provider_ids(&settings, &enabled_ids); @@ -1159,6 +1163,13 @@ pub struct DeepSeekPricingStatus { pub fn get_deepseek_pricing_status( state: tauri::State<'_, Mutex>, ) -> Option { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return None; + } let settings = Settings::load(); if !settings.enabled_providers.contains("deepseek") { return None; @@ -1203,11 +1214,21 @@ pub async fn refresh_providers_if_stale(app: tauri::AppHandle) -> Result<(), Str pub fn get_cached_providers( state: tauri::State<'_, Mutex>, ) -> Vec { - let snapshots = state + let (snapshots, proof_mode, proof_settings) = state .lock() - .map(|guard| guard.provider_cache.clone()) - .unwrap_or_default(); - let settings = Settings::load(); + .map(|guard| { + ( + guard.provider_cache.clone(), + guard.is_containment_proof(), + guard.proof_settings().cloned(), + ) + }) + .unwrap_or((Vec::new(), true, None)); + let settings = if proof_mode { + proof_settings.unwrap_or_default() + } else { + Settings::load() + }; snapshots .into_iter() diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index b687c9fd48..036004d771 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -517,8 +517,16 @@ fn parse_language(s: &str) -> Option { #[tauri::command] pub async fn update_settings( app: tauri::AppHandle, + state: tauri::State<'_, Mutex>, patch: SettingsUpdate, ) -> Result { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return Err("settings mutations disabled in containment proof mode".to_string()); + } let mut settings = Settings::load(); let notify_float_bar = patch.notifies_float_bar(); let refresh_provider_data = patch.refreshes_provider_data(); diff --git a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs index 4403ba7ad0..71644b8a99 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs @@ -1,15 +1,29 @@ //! Upstream 0.53 Usage & Spend accounting bridge. +use std::sync::Mutex; + use codexbar::cost_scanner::CostScanner; use codexbar::settings::Settings; use codexbar::spend_contract::{SpendContract, build_local_spend_contract_from_summary}; +use tauri::State; + +use crate::state::AppState; #[tauri::command] pub async fn get_spend_contract( + state: State<'_, Mutex>, provider_id: String, history_days: Option, include_open_codex: Option, ) -> Result { + let containment_proof = state + .lock() + .map_err(|_| "app state lock is poisoned".to_string())? + .is_containment_proof(); + if containment_proof { + return Err("spend contract is unavailable during containment proof".to_string()); + } + let provider = provider_id.trim().to_ascii_lowercase(); if !matches!(provider.as_str(), "codex" | "claude" | "pi" | "opencodego") { return Err(format!( diff --git a/apps/desktop-tauri/src-tauri/src/commands/updater.rs b/apps/desktop-tauri/src-tauri/src/commands/updater.rs index 3eeb7a1367..0bcaffd39b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/updater.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/updater.rs @@ -14,6 +14,13 @@ use codexbar::updater::UpdateInfo; #[tauri::command] pub fn get_update_state(state: tauri::State<'_, Mutex>) -> UpdateStatePayload { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return UpdateState::Idle.to_payload(); + } state .lock() .map(|guard| guard.update_payload()) @@ -28,6 +35,9 @@ pub async fn check_for_updates( // Guard: skip if already checking or downloading. { let mut guard = state.lock().map_err(|e| e.to_string())?; + if guard.is_containment_proof() { + return Ok(UpdateState::Idle.to_payload()); + } match guard.update_state { UpdateState::Checking | UpdateState::Downloading(_) => { return Ok(guard.update_payload()); @@ -79,6 +89,13 @@ pub async fn download_update( app: tauri::AppHandle, state: tauri::State<'_, Mutex>, ) -> Result { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return Ok(UpdateState::Idle.to_payload()); + } let info = match update_info_for_download(&state)? { DownloadStart::Ready(info) => info, DownloadStart::AlreadyDownloading(payload) => return Ok(payload), @@ -201,6 +218,12 @@ pub fn apply_update(state: tauri::State<'_, Mutex>) -> Result<(), Stri } pub(crate) fn apply_ready_update(state: &Mutex) -> Result<(), String> { + { + let guard = state.lock().map_err(|e| e.to_string())?; + if guard.is_containment_proof() { + return Err("updates disabled in containment proof mode".to_string()); + } + } let (path, expected_sha256) = { let guard = state.lock().map_err(|e| e.to_string())?; let path = guard @@ -223,6 +246,13 @@ pub fn dismiss_update( app: tauri::AppHandle, state: tauri::State<'_, Mutex>, ) -> Result { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return Ok(UpdateState::Idle.to_payload()); + } let payload = { let mut guard = state.lock().map_err(|e| e.to_string())?; guard.update_state = UpdateState::Idle; @@ -236,6 +266,13 @@ pub fn dismiss_update( #[tauri::command] pub fn open_release_page(state: tauri::State<'_, Mutex>) -> Result<(), String> { + if state + .lock() + .map(|guard| guard.is_containment_proof()) + .unwrap_or(true) + { + return Err("updates disabled in containment proof mode".to_string()); + } let url = { let guard = state.lock().map_err(|e| e.to_string())?; guard diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index ee7ae528c8..4114cccc53 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -8,8 +8,10 @@ use serde::Serialize; use tauri::State; use super::ProviderUsageSnapshot; +use crate::proof_runtime::ContainmentProof; use crate::state::AppState; use std::collections::{BTreeSet, HashMap}; +use std::path::PathBuf; use std::sync::{Mutex, OnceLock}; #[derive(Debug, Clone, Serialize)] @@ -188,26 +190,44 @@ pub async fn get_usage_spend_summary( history_days: Option, force_refresh: Option, ) -> Result { - let cached = { + let (cached, containment_proof) = { let guard = state.lock().map_err(|e| e.to_string())?; - guard.provider_cache.clone() + let containment_proof = guard.containment_proof.clone(); + if guard.is_containment_proof() && containment_proof.is_none() { + return Err("containment proof state is missing its manifest".to_string()); + } + (guard.provider_cache.clone(), containment_proof) }; let selected_days = history_days.unwrap_or(30); let force_refresh = force_refresh.unwrap_or(false); let built = tauri::async_runtime::spawn_blocking(move || { - build_usage_spend_summary_cached(&cached, selected_days, force_refresh) + build_usage_spend_summary_cached( + &cached, + selected_days, + force_refresh, + containment_proof.as_ref(), + ) }) .await .map_err(|e| format!("usage spend worker failed: {e}"))??; - let current_cached = state - .lock() - .map_err(|e| e.to_string()) - .map(|guard| guard.provider_cache.clone())?; - let current_key = usage_spend_cache_key( + let (current_cached, current_containment_proof) = { + let guard = state.lock().map_err(|e| e.to_string())?; + let containment_proof = guard.containment_proof.clone(); + if guard.is_containment_proof() && containment_proof.is_none() { + return Err("containment proof state is missing its manifest".to_string()); + } + (guard.provider_cache.clone(), containment_proof) + }; + let current_settings = current_containment_proof + .as_ref() + .map(ContainmentProof::proof_settings) + .unwrap_or_else(codexbar::settings::Settings::load); + let current_key = usage_spend_cache_key_for_runtime( ¤t_cached, selected_days, - &codexbar::settings::Settings::load(), + ¤t_settings, + current_containment_proof.as_ref(), ); if current_key != built.key { if let Some(owner) = built.refresh_owner.as_ref() { @@ -237,9 +257,13 @@ fn build_usage_spend_summary_cached( cached: &[ProviderUsageSnapshot], selected_days: u32, force_refresh: bool, + containment_proof: Option<&ContainmentProof>, ) -> Result { - let settings = codexbar::settings::Settings::load(); - let key = usage_spend_cache_key(cached, selected_days, &settings); + let settings = containment_proof + .map(ContainmentProof::proof_settings) + .unwrap_or_else(codexbar::settings::Settings::load); + let key = + usage_spend_cache_key_for_runtime(cached, selected_days, &settings, containment_proof); { let guard = usage_spend_coordinator() .lock() @@ -261,11 +285,16 @@ fn build_usage_spend_summary_cached( .map_err(|error| error.to_string())?; coordinator.begin(key.clone()) }; - let summary = build_usage_spend_summary(cached, selected_days, &settings, force_refresh); + let summary = containment_proof.map_or_else( + || build_usage_spend_summary(cached, selected_days, &settings, force_refresh), + |proof| build_usage_spend_summary_in_containment_proof(cached, selected_days, proof), + ); let refreshing = summary_is_refreshing(&summary); - let codex_scan_pause_reason = + let codex_scan_pause_reason = containment_proof.is_none().then(|| { codexbar::core::JsonlScanner::load_cache_status(codexbar::core::ProviderId::Codex, None) - .codex_scan_pause_reason; + .codex_scan_pause_reason + }); + let codex_scan_pause_reason = codex_scan_pause_reason.flatten(); let mut coordinator = usage_spend_coordinator() .lock() @@ -306,21 +335,61 @@ fn usage_spend_cache_key( selected_days: u32, settings: &codexbar::settings::Settings, ) -> String { - usage_spend_cache_key_with_privacy( + usage_spend_cache_key_with_identity( cached, selected_days, settings.open_codex_usage_logs_enabled, settings.hide_native_codex_cost_when_open_codex_present, settings.hide_personal_info, + None, ) } +fn usage_spend_cache_key_for_runtime( + cached: &[ProviderUsageSnapshot], + selected_days: u32, + settings: &codexbar::settings::Settings, + containment_proof: Option<&ContainmentProof>, +) -> String { + if containment_proof.is_none() { + return usage_spend_cache_key(cached, selected_days, settings); + } + let proof_identity = containment_proof.map(containment_proof_cache_identity); + usage_spend_cache_key_with_identity( + cached, + selected_days, + settings.open_codex_usage_logs_enabled, + settings.hide_native_codex_cost_when_open_codex_present, + settings.hide_personal_info, + proof_identity.as_deref(), + ) +} + +#[cfg(test)] fn usage_spend_cache_key_with_privacy( cached: &[ProviderUsageSnapshot], selected_days: u32, include_opencodex: bool, hide_native: bool, hide_personal_info: bool, +) -> String { + usage_spend_cache_key_with_identity( + cached, + selected_days, + include_opencodex, + hide_native, + hide_personal_info, + None, + ) +} + +fn usage_spend_cache_key_with_identity( + cached: &[ProviderUsageSnapshot], + selected_days: u32, + include_opencodex: bool, + hide_native: bool, + hide_personal_info: bool, + proof_identity: Option<&str>, ) -> String { let mut revisions: Vec = cached .iter() @@ -348,7 +417,7 @@ fn usage_spend_cache_key_with_privacy( }) .collect(); revisions.sort(); - format!( + let key = format!( "{}|{}|{}|{}|{}|{}", chrono::Local::now().date_naive(), selected_days, @@ -356,9 +425,148 @@ fn usage_spend_cache_key_with_privacy( hide_native, hide_personal_info, revisions.join(";") + ); + match proof_identity { + Some(identity) => format!("{key}|proof:{identity}"), + None => key, + } +} + +fn containment_proof_cache_identity(proof: &ContainmentProof) -> String { + format!( + "{}:{}:{}:{}:{}:{}:{}", + proof.manifest.schema, + proof.manifest.version, + proof.manifest.kind, + proof.manifest.provider, + proof.manifest.fixture_roots.gemini_cli_home.display(), + proof.manifest.fixture_roots.tokscale_config_dir.display(), + proof.manifest.now_utc.to_rfc3339(), ) } +fn containment_proof_database_roots(proof: &ContainmentProof) -> [PathBuf; 3] { + let gemini_cli_home = &proof.manifest.fixture_roots.gemini_cli_home; + [ + gemini_cli_home + .join("antigravity-cli") + .join("conversations"), + gemini_cli_home.join("antigravity"), + gemini_cli_home.join("antigravity").join("conversations"), + ] +} + +fn containment_proof_jsonl_sessions_root(proof: &ContainmentProof) -> PathBuf { + proof + .manifest + .fixture_roots + .tokscale_config_dir + .join("antigravity-cache") + .join("sessions") +} + +fn unavailable_codex_spend_contract(history_days: u32) -> SpendContract { + SpendContract { + provider_id: "codex".to_string(), + history_days: if history_days == 0 { + 365 + } else { + history_days.clamp(1, 365) + }, + known_cost_usd: None, + known_zero: false, + provenance: codexbar::spend_contract::CostProvenance::Unknown, + price_coverage: Default::default(), + price_coverage_ratio: None, + history_coverage_established: false, + token_mix: Default::default(), + conversation_count: 0, + models: Vec::new(), + projects: Vec::new(), + conversations: Vec::new(), + daily: Vec::new(), + hourly_activity: Vec::new(), + project_source_status: None, + custom_pricing_active: false, + imports: Vec::new(), + } +} + +fn build_usage_spend_summary_in_containment_proof( + cached: &[ProviderUsageSnapshot], + selected_days: u32, + proof: &ContainmentProof, +) -> UsageSpendSummary { + let database_roots = containment_proof_database_roots(proof); + let jsonl_sessions_root = containment_proof_jsonl_sessions_root(proof); + let seven = codexbar::providers::antigravity::local_sessions::summarize_from_roots( + &database_roots, + &jsonl_sessions_root, + proof.manifest.now_utc, + 7, + ); + let thirty = codexbar::providers::antigravity::local_sessions::summarize_from_roots( + &database_roots, + &jsonl_sessions_root, + proof.manifest.now_utc, + 30, + ); + let cached_snapshot = cached + .iter() + .find(|snapshot| snapshot.provider_id == "antigravity"); + let spend = antigravity_spend_values(cached_spend(cached_snapshot), &seven, &thirty); + let display_name = cached_snapshot + .map(|snapshot| snapshot.display_name.trim()) + .filter(|name| !name.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| "Antigravity".to_string()); + let currency = cached_snapshot + .and_then(|snapshot| snapshot.cost.as_ref()) + .map(|cost| cost.currency_code.clone()) + .unwrap_or_else(|| "USD".to_string()); + let daily = cached_snapshot + .and_then(|snapshot| snapshot.cost.as_ref()) + .map(|cost| { + cost.daily + .iter() + .map(|point| UsageSpendDailyPoint { + day: point.day.clone(), + amount: point.amount, + }) + .collect() + }) + .unwrap_or_default(); + let seven_day_estimate = Some(seven.cost_estimate.clone()); + let thirty_day_estimate = Some(thirty.cost_estimate.clone()); + let row = UsageSpendRow { + provider_id: "antigravity".to_string(), + display_name, + seven_day: spend.seven_day, + thirty_day: spend.thirty_day, + seven_day_estimate, + thirty_day_estimate, + seven_day_tokens: spend.seven_day_tokens, + thirty_day_tokens: spend.thirty_day_tokens, + currency, + source: spend.source, + included_in_overview: true, + daily, + refreshing: false, + stale_updated_at: None, + }; + UsageSpendSummary { + rows: vec![row], + contract: unavailable_codex_spend_contract(selected_days), + reporting_day: proof + .manifest + .now_utc + .date_naive() + .format("%Y-%m-%d") + .to_string(), + dashboard_timezone: codexbar::core::local_timezone_name(), + } +} + fn build_usage_spend_summary( cached: &[ProviderUsageSnapshot], selected_days: u32, @@ -804,6 +1012,7 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { #[cfg(test)] mod cache_key_tests { use super::*; + use std::fs; fn local_history( total_tokens: u64, @@ -912,4 +1121,104 @@ mod cache_key_tests { assert_eq!(spend.thirty_day_tokens, Some(0)); assert!(spend.source.contains("API list-price estimate")); } + + fn test_containment_proof() -> (crate::test_support::TempDir, ContainmentProof) { + let root = crate::test_support::TempDir::new(); + let gemini_cli_home = root.path().join("gemini"); + let tokscale_config_dir = root.path().join("tokscale"); + let scratch_root = root.path().join("scratch"); + fs::create_dir_all(&gemini_cli_home).unwrap(); + fs::create_dir_all(&tokscale_config_dir).unwrap(); + fs::create_dir_all(&scratch_root).unwrap(); + let manifest = serde_json::json!({ + "schema": "codexbar.containment-proof", + "version": 1, + "kind": "antigravityUsageSpend", + "provider": "antigravity", + "fixtureRoots": { + "geminiCliHome": gemini_cli_home, + "tokscaleConfigDir": tokscale_config_dir + }, + "scratchRoot": scratch_root, + "nowUtc": "2026-09-24T12:34:56Z" + }); + let proof = ContainmentProof::from_manifest_json(&manifest.to_string()).unwrap(); + (root, proof) + } + + fn write_proof_history(proof: &ContainmentProof) { + let sessions = containment_proof_jsonl_sessions_root(proof); + fs::create_dir_all(&sessions).unwrap(); + let now = proof.manifest.now_utc; + let recent = (now - chrono::Duration::days(1)).timestamp_millis(); + let older = (now - chrono::Duration::days(10)).timestamp_millis(); + let lines = [ + serde_json::json!({ + "type": "session_meta", + "modelId": "claude-sonnet-4-6" + }) + .to_string(), + serde_json::json!({ + "type": "usage", + "responseId": "recent", + "timestamp": recent, + "input": 100, + "output": 20 + }) + .to_string(), + serde_json::json!({ + "type": "usage", + "responseId": "older", + "timestamp": older, + "input": 200, + "output": 30 + }) + .to_string(), + ] + .join("\n"); + fs::write(sessions.join("fixture.jsonl"), lines).unwrap(); + } + + #[test] + fn containment_proof_routes_to_exact_antigravity_singleton_with_fixed_totals() { + let (_root, proof) = test_containment_proof(); + write_proof_history(&proof); + + let summary = build_usage_spend_summary_in_containment_proof(&[], 30, &proof); + + assert_eq!(summary.rows.len(), 1); + let row = &summary.rows[0]; + assert_eq!(row.provider_id, "antigravity"); + assert_eq!(row.seven_day_tokens, Some(120)); + assert_eq!(row.thirty_day_tokens, Some(350)); + assert_eq!(summary.reporting_day, "2026-09-24"); + } + + #[test] + fn containment_proof_returns_empty_unavailable_codex_contract() { + let (_root, proof) = test_containment_proof(); + let summary = build_usage_spend_summary_in_containment_proof(&[], 30, &proof); + + assert_eq!(summary.contract.provider_id, "codex"); + assert_eq!(summary.contract.known_cost_usd, None); + assert!(!summary.contract.history_coverage_established); + assert_eq!( + summary.contract.provenance, + codexbar::spend_contract::CostProvenance::Unknown + ); + assert!(!summary.contract.custom_pricing_active); + assert!(summary.contract.models.is_empty()); + assert!(summary.contract.imports.is_empty()); + } + + #[test] + fn containment_proof_identity_is_part_of_cache_key() { + let (_root, proof) = test_containment_proof(); + let settings = proof.proof_settings(); + let proof_key = usage_spend_cache_key_for_runtime(&[], 30, &settings, Some(&proof)); + let production_key = usage_spend_cache_key_for_runtime(&[], 30, &settings, None); + + assert_ne!(proof_key, production_key); + assert!(proof_key.contains("proof:")); + } } diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index f7bbbef11f..347160e42b 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -11,11 +11,14 @@ mod floatbar; mod geometry_store; mod powertoys; mod proof_harness; +mod proof_runtime; mod shell; mod shortcut_bridge; mod state; mod surface; mod surface_target; +#[cfg(test)] +mod test_support; mod tray_accounts; mod tray_bridge; mod tray_menu; @@ -113,6 +116,20 @@ fn should_suppress_blur_dismiss(launch: LaunchBehavior, proof_mode: bool) -> boo } fn main() { + let containment_proof = match proof_runtime::ContainmentProof::from_env() { + Ok(proof) => proof, + Err(error) => { + eprintln!("CodexBar containment proof rejected: {error}"); + std::process::exit(2); + } + }; + if let Some(proof) = containment_proof.as_ref() + && let Err(error) = proof.install() + { + eprintln!("CodexBar containment proof setup failed: {error}"); + std::process::exit(2); + } + // Per-process log file names: the shell writes codexbar-desktop.log so // its cached handle never blocks the CLI's rotation on Windows. // SAFETY: runs before any thread spawns; no concurrent env access exists. @@ -120,52 +137,81 @@ fn main() { codexbar::logging::install_panic_hook(); codexbar::logging::init(false, false).expect("failed to initialize logging"); - let proof_config = proof_harness::ProofConfig::from_env(); + let containment_active = containment_proof.is_some(); + let proof_config = if containment_active { + Some(proof_harness::ProofConfig { + target_surface: "settings".to_string(), + settings_tab: Some("usageSpend".to_string()), + target_payload: Some("usageSpend".to_string()), + }) + } else { + proof_harness::ProofConfig::from_env() + }; let is_proof_mode = proof_config.is_some(); - let force_start_visible = std::env::var_os("CODEXBAR_START_VISIBLE").is_some(); - let settings = codexbar::settings::Settings::load(); + let force_start_visible = + !containment_active && std::env::var_os("CODEXBAR_START_VISIBLE").is_some(); + let settings = containment_proof + .as_ref() + .map(proof_runtime::ContainmentProof::proof_settings) + .unwrap_or_else(codexbar::settings::Settings::load); let launch = launch_behavior( force_start_visible, settings.start_minimized, std::env::args().skip(1), ); - let mut initial_state = AppState::new(); + let mut initial_state = containment_proof + .as_ref() + .map(|proof| AppState::new_for_containment_proof(proof.clone())) + .unwrap_or_else(AppState::new); initial_state.proof_config = proof_config; - // Validate the complete proof seed before installing any snapshots, so an - // invalid multi-provider fixture cannot leave a partial cache behind. - if let Some(snapshots) = - proof_harness::seed_usage_snapshots_from_env(initial_state.proof_config.as_ref()) - { - let seeded_at = std::time::Instant::now(); - for snapshot in &snapshots { - tracing::info!( - "proof-harness: seeded provider snapshot for '{}'", - snapshot.provider_id - ); - if let Some(provider) = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id) - { - initial_state - .provider_cache_updated_at_by_provider - .insert(provider, seeded_at); + if !containment_active { + // Validate the complete proof seed before installing any snapshots, so an + // invalid multi-provider fixture cannot leave a partial cache behind. + if let Some(snapshots) = + proof_harness::seed_usage_snapshots_from_env(initial_state.proof_config.as_ref()) + { + let seeded_at = std::time::Instant::now(); + for snapshot in &snapshots { + tracing::info!( + "proof-harness: seeded provider snapshot for '{}'", + snapshot.provider_id + ); + if let Some(provider) = + codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id) + { + initial_state + .provider_cache_updated_at_by_provider + .insert(provider, seeded_at); + } } + initial_state.provider_cache.extend(snapshots); + initial_state.provider_cache_seeded = true; + initial_state.provider_cache_updated_at = Some(seeded_at); } - initial_state.provider_cache.extend(snapshots); - initial_state.provider_cache_seeded = true; - initial_state.provider_cache_updated_at = Some(seeded_at); } - tauri::Builder::default() + let builder = tauri::Builder::default() .manage(Mutex::new(initial_state)) - .plugin(shortcut_bridge::plugin()) - .plugin(tauri_plugin_dialog::init()) - .plugin(tauri_plugin_single_instance::init(|app, args, _cwd| { - if should_reopen_primary_window_from_instance_args(args.iter().skip(1)) { - let request = primary_window_request(); - let _ = - shell::reopen_to_target(app, request.mode, request.target, request.position); - } - })) + .plugin(tauri_plugin_dialog::init()); + let builder = if containment_active { + builder + } else { + builder + .plugin(shortcut_bridge::plugin()) + .plugin(tauri_plugin_single_instance::init(|app, args, _cwd| { + if should_reopen_primary_window_from_instance_args(args.iter().skip(1)) { + let request = primary_window_request(); + let _ = shell::reopen_to_target( + app, + request.mode, + request.target, + request.position, + ); + } + })) + }; + builder .invoke_handler(tauri::generate_handler![ commands::get_bootstrap_state, commands::get_provider_catalog, @@ -288,19 +334,24 @@ fn main() { floatbar::set_float_bar_orientation, ]) .setup(move |app| { - if let Err(error) = codexbar::providers::claude::accounts::cleanup_abandoned_logins() { + if !containment_active + && let Err(error) = + codexbar::providers::claude::accounts::cleanup_abandoned_logins() + { tracing::warn!("failed to clean abandoned Claude sign-in directories: {error}"); } if let Some(window) = app.get_webview_window("main") { shell::dwm::force_dark_caption(&window); window.hide()?; } - tray_bridge::setup(app)?; - shortcut_bridge::register(app.handle()); - floatbar::install(app.handle()); - auto_refresh::install(app.handle().clone()); - if settings.powertoys_status_pipe_enabled { - powertoys::install(app.handle().clone()); + if !containment_active { + tray_bridge::setup(app)?; + shortcut_bridge::register(app.handle()); + floatbar::install(app.handle()); + auto_refresh::install(app.handle().clone()); + if settings.powertoys_status_pipe_enabled { + powertoys::install(app.handle().clone()); + } } // Give the WebView/event loop one turn to finish startup before @@ -310,7 +361,14 @@ fn main() { let app_handle = app.handle().clone(); tauri::async_runtime::spawn(async move { tokio::time::sleep(PROOF_ACTIVATION_DELAY).await; - proof_harness::activate(&app_handle); + if containment_active { + if let Err(error) = proof_harness::activate_without_focus(&app_handle) { + tracing::error!("containment proof reveal failed: {error}"); + app_handle.exit(2); + } + } else { + proof_harness::activate(&app_handle); + } }); } else if launch.open_primary_window_at_start { let app = app.handle().clone(); diff --git a/apps/desktop-tauri/src-tauri/src/proof_harness.rs b/apps/desktop-tauri/src-tauri/src/proof_harness.rs index 4d05b39b78..c69bee013e 100644 --- a/apps/desktop-tauri/src-tauri/src/proof_harness.rs +++ b/apps/desktop-tauri/src-tauri/src/proof_harness.rs @@ -138,6 +138,43 @@ pub fn activate(app: &AppHandle) { } } +/// Show the containment proof surface behind the user's other windows and +/// abort if Windows ever makes it the foreground window. +pub fn activate_without_focus(app: &AppHandle) -> Result<(), String> { + let config = { + let st = app.state::>(); + st.lock() + .map_err(|_| "containment proof state lock is poisoned".to_string())? + .proof_config + .clone() + .ok_or_else(|| "containment proof surface configuration is missing".to_string())? + }; + let window = app + .get_webview_window("main") + .ok_or_else(|| "containment proof main window is unavailable".to_string())?; + let hwnd = crate::proof_runtime::native_window_handle(&window)?; + crate::proof_runtime::start_foreground_guard(app.clone(), hwnd); + + let target = config.surface_mode(); + let position = match target { + SurfaceMode::Settings | SurfaceMode::PopOut => None, + _ => proof_window_position(app), + }; + tracing::info!( + "containment-proof: showing surface={} tab={:?} without activation", + config.target_surface, + config.settings_tab, + ); + let mode = shell::transition_to_target_without_activation( + app, + target, + config.surface_target(), + position, + )?; + tracing::info!("containment-proof: nonactivating transition succeeded → {mode:?}"); + Ok(()) +} + /// Bottom inset (physical px) kept between the proof panel's bottom edge and /// the monitor work-area bottom (#265). const PROOF_BOTTOM_INSET_PX: i32 = 8; diff --git a/apps/desktop-tauri/src-tauri/src/proof_runtime.rs b/apps/desktop-tauri/src-tauri/src/proof_runtime.rs new file mode 100644 index 0000000000..6e54ef0a1f --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/proof_runtime.rs @@ -0,0 +1,569 @@ +//! Fail-closed runtime containment for deterministic provider proof runs. + +use std::collections::HashSet; +use std::fs; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use tauri::WebviewWindow; + +pub const CONTAINMENT_PROOF_ENV: &str = "CODEXBAR_CONTAINMENT_PROOF"; +pub const LEGACY_PROOF_ENV: &str = "CODEXBAR_PROOF_MODE"; +pub const LEGACY_SEED_ENV: &str = "CODEXBAR_SEED_USAGE_JSON"; +pub const PROOF_KIND: &str = "antigravityUsageSpend"; +pub const PROOF_PROVIDER: &str = "antigravity"; +const PROOF_SCHEMA: &str = "codexbar.containment-proof"; +const PROOF_VERSION: u32 = 1; + +/// Return the native HWND for the proof window. The containment proof is a +/// Windows-only UI path; other platforms fail closed instead of falling back +/// to Tauri's activating `show()` behavior. +pub fn native_window_handle(window: &WebviewWindow) -> Result { + #[cfg(windows)] + { + use raw_window_handle::{HasWindowHandle, RawWindowHandle}; + + let handle = window + .window_handle() + .map_err(|error| format!("cannot inspect proof window handle: {error}"))?; + match handle.as_raw() { + RawWindowHandle::Win32(handle) => Ok(handle.hwnd.get()), + _ => Err("containment proof requires a Win32 window handle".to_string()), + } + } + + #[cfg(not(windows))] + { + let _ = window; + Err("containment proof UI is supported only on Windows".to_string()) + } +} + +/// Reveal the proof window behind other windows without activating it. +pub fn show_window_without_activation(window: &WebviewWindow) -> Result<(), String> { + let hwnd = native_window_handle(window)?; + #[cfg(windows)] + { + const HWND_BOTTOM: isize = 1; + const HWND_NOTOPMOST: isize = -2; + const SW_SHOWNOACTIVATE: i32 = 4; + const SWP_NOMOVE: u32 = 0x0002; + const SWP_NOSIZE: u32 = 0x0001; + const SWP_NOACTIVATE: u32 = 0x0010; + const SWP_SHOWWINDOW: u32 = 0x0040; + + let native = hwnd; + // SAFETY: `native` comes from the live Tauri WebviewWindow above; all + // calls preserve activation and only adjust this window's Z-order. + unsafe { + if set_window_pos( + native, + HWND_NOTOPMOST, + 0, + 0, + 0, + 0, + SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE, + ) == 0 + { + return Err("could not remove proof window from topmost Z-order".to_string()); + } + show_window(native, SW_SHOWNOACTIVATE); + if set_window_pos( + native, + HWND_BOTTOM, + 0, + 0, + 0, + 0, + SWP_NOMOVE | SWP_NOSIZE | SWP_NOACTIVATE | SWP_SHOWWINDOW, + ) == 0 + { + return Err("could not show proof window without activation".to_string()); + } + if get_foreground_window() == native { + show_window(native, 0); // SW_HIDE + return Err("proof window unexpectedly became foreground".to_string()); + } + } + Ok(()) + } + + #[cfg(not(windows))] + { + let _ = hwnd; + Err("containment proof UI is supported only on Windows".to_string()) + } +} + +/// Stop the proof run immediately if Windows ever reports its HWND as the +/// foreground window. This is a guard for regressions in the native reveal +/// path; normal proof flow never activates the window. +pub fn start_foreground_guard(app: tauri::AppHandle, hwnd: isize) { + #[cfg(windows)] + tauri::async_runtime::spawn(async move { + use std::time::Duration; + + let native_handle = hwnd; + loop { + tokio::time::sleep(Duration::from_millis(10)).await; + // SAFETY: `native_handle` is the HWND obtained from the live proof + // WebviewWindow, and GetForegroundWindow is a read-only query. + if unsafe { get_foreground_window() } == native_handle { + tracing::error!("containment proof aborted: proof window became foreground"); + // SAFETY: hide only the proof HWND before exiting so the run + // cannot continue interacting with the user's desktop. + unsafe { show_window(native_handle, 0) }; // SW_HIDE + app.exit(2); + break; + } + } + }); + + #[cfg(not(windows))] + { + let _ = (app, hwnd); + } +} + +#[cfg(windows)] +#[link(name = "user32")] +unsafe extern "system" { + #[link_name = "ShowWindow"] + fn show_window(hwnd: isize, command: i32) -> i32; + #[link_name = "SetWindowPos"] + fn set_window_pos( + hwnd: isize, + insert_after: isize, + x: i32, + y: i32, + width: i32, + height: i32, + flags: u32, + ) -> i32; + #[link_name = "GetForegroundWindow"] + fn get_foreground_window() -> isize; +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawFixtureRoots { + #[serde(rename = "geminiCliHome")] + gemini_cli_home: PathBuf, + #[serde(rename = "tokscaleConfigDir")] + tokscale_config_dir: PathBuf, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawManifest { + schema: String, + #[serde(alias = "schemaVersion")] + version: u32, + kind: String, + provider: String, + #[serde(rename = "fixtureRoots")] + fixture_roots: RawFixtureRoots, + #[serde(rename = "scratchRoot")] + scratch_root: PathBuf, + #[serde(rename = "nowUtc")] + now_utc: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FixtureRoots { + pub gemini_cli_home: PathBuf, + pub tokscale_config_dir: PathBuf, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ContainmentManifest { + pub schema: String, + pub version: u32, + pub kind: String, + pub provider: String, + pub fixture_roots: FixtureRoots, + pub scratch_root: PathBuf, + pub now_utc: DateTime, +} + +#[derive(Debug, Clone)] +pub struct ContainmentProof { + pub manifest: ContainmentManifest, + pub config_root: PathBuf, + pub log_root: PathBuf, + pub webview_user_data_folder: PathBuf, +} + +impl ContainmentProof { + pub fn from_env() -> Result, String> { + let Some(raw_path) = std::env::var_os(CONTAINMENT_PROOF_ENV) else { + return Ok(None); + }; + + for variable in [LEGACY_PROOF_ENV, LEGACY_SEED_ENV] { + if std::env::var_os(variable).is_some() { + return Err(format!( + "{CONTAINMENT_PROOF_ENV} cannot be combined with {variable}" + )); + } + } + + let path = PathBuf::from(raw_path); + if !path.is_absolute() { + return Err(format!( + "{CONTAINMENT_PROOF_ENV} must name an absolute manifest path" + )); + } + let raw = fs::read_to_string(&path).map_err(|error| { + format!( + "cannot read {CONTAINMENT_PROOF_ENV} manifest {}: {error}", + path.display() + ) + })?; + Self::from_manifest_json(&raw).map(Some) + } + + pub fn from_manifest_json(raw: &str) -> Result { + let raw_manifest: RawManifest = serde_json::from_str(raw) + .map_err(|error| format!("invalid containment proof manifest: {error}"))?; + if raw_manifest.schema != PROOF_SCHEMA { + return Err(format!( + "unsupported containment proof schema {:?}", + raw_manifest.schema + )); + } + if raw_manifest.version != PROOF_VERSION { + return Err(format!( + "unsupported containment proof version {}", + raw_manifest.version + )); + } + if raw_manifest.kind != PROOF_KIND { + return Err(format!("containment proof kind must be {PROOF_KIND:?}")); + } + if raw_manifest.provider != PROOF_PROVIDER { + return Err(format!( + "containment proof provider must be {PROOF_PROVIDER:?}" + )); + } + let now_utc = DateTime::parse_from_rfc3339(&raw_manifest.now_utc) + .map_err(|error| format!("nowUtc must be RFC3339: {error}"))? + .with_timezone(&Utc); + + let fixture_roots = FixtureRoots { + gemini_cli_home: canonical_directory( + "fixtureRoots.geminiCliHome", + &raw_manifest.fixture_roots.gemini_cli_home, + )?, + tokscale_config_dir: canonical_directory( + "fixtureRoots.tokscaleConfigDir", + &raw_manifest.fixture_roots.tokscale_config_dir, + )?, + }; + let scratch_root = canonical_empty_directory("scratchRoot", &raw_manifest.scratch_root)?; + let temp_root = fs::canonicalize(std::env::temp_dir()) + .map_err(|error| format!("system temporary root cannot be canonicalized: {error}"))?; + let roots = [ + fixture_roots.gemini_cli_home.as_path(), + fixture_roots.tokscale_config_dir.as_path(), + scratch_root.as_path(), + ]; + for (label, root) in [ + ( + "fixtureRoots.geminiCliHome", + fixture_roots.gemini_cli_home.as_path(), + ), + ( + "fixtureRoots.tokscaleConfigDir", + fixture_roots.tokscale_config_dir.as_path(), + ), + ("scratchRoot", scratch_root.as_path()), + ] { + if !is_dedicated_temp_child(root, &temp_root) { + return Err(format!( + "{label} must be a child of the system temporary root {}", + temp_root.display() + )); + } + } + for (index, left) in roots.iter().enumerate() { + for right in roots.iter().skip(index + 1) { + if paths_overlap(left, right) { + return Err(format!( + "containment proof roots overlap: {} and {}", + left.display(), + right.display() + )); + } + } + } + + Ok(Self { + manifest: ContainmentManifest { + schema: raw_manifest.schema, + version: raw_manifest.version, + kind: raw_manifest.kind, + provider: raw_manifest.provider, + fixture_roots, + scratch_root: scratch_root.clone(), + now_utc, + }, + config_root: scratch_root.clone(), + log_root: scratch_root.join("logs"), + webview_user_data_folder: scratch_root.join("webview2-user-data"), + }) + } + + pub fn proof_settings(&self) -> codexbar::settings::Settings { + codexbar::settings::Settings { + enabled_providers: HashSet::from([PROOF_PROVIDER.to_string()]), + provider_order: vec![PROOF_PROVIDER.to_string()], + refresh_interval_secs: 0, + adaptive_refresh: false, + refresh_all_providers_on_menu_open: false, + start_minimized: false, + powertoys_status_pipe_enabled: false, + float_bar_enabled: false, + auto_download_updates: false, + install_updates_on_quit: false, + ..Default::default() + } + } + + pub fn install(&self) -> Result<(), String> { + fs::create_dir_all(&self.log_root) + .map_err(|error| format!("cannot create proof log root: {error}"))?; + fs::create_dir_all(&self.webview_user_data_folder) + .map_err(|error| format!("cannot create proof WebView root: {error}"))?; + codexbar::logging::install_config_root_override(self.config_root.clone())?; + set_process_env("WEBVIEW2_USER_DATA_FOLDER", &self.webview_user_data_folder); + set_process_env( + "GEMINI_CLI_HOME", + &self.manifest.fixture_roots.gemini_cli_home, + ); + set_process_env( + "TOKSCALE_CONFIG_DIR", + &self.manifest.fixture_roots.tokscale_config_dir, + ); + set_process_env( + "CODEXBAR_CONTAINMENT_NOW_UTC", + Path::new(&self.manifest.now_utc.to_rfc3339()), + ); + Ok(()) + } +} + +fn set_process_env(name: &str, value: &Path) { + // SAFETY: startup runs before Tauri creates worker threads. + unsafe { std::env::set_var(name, value.as_os_str()) }; +} + +fn canonical_directory(label: &str, path: &Path) -> Result { + ensure_absolute(label, path)?; + reject_reparse_chain(label, path)?; + let canonical = fs::canonicalize(path) + .map_err(|error| format!("{label} cannot be canonicalized: {error}"))?; + let metadata = fs::metadata(&canonical) + .map_err(|error| format!("{label} metadata unavailable: {error}"))?; + if !metadata.is_dir() { + return Err(format!("{label} must be a directory")); + } + reject_reparse_tree(label, &canonical)?; + Ok(canonical) +} + +fn canonical_empty_directory(label: &str, path: &Path) -> Result { + let canonical = canonical_directory(label, path)?; + if fs::read_dir(&canonical) + .map_err(|error| format!("{label} cannot be read: {error}"))? + .next() + .is_some() + { + return Err(format!("{label} must be empty")); + } + Ok(canonical) +} + +fn ensure_absolute(label: &str, path: &Path) -> Result<(), String> { + if path.as_os_str().is_empty() || !path.is_absolute() { + return Err(format!("{label} must be a non-empty absolute path")); + } + Ok(()) +} + +fn is_dedicated_temp_child(root: &Path, temp_root: &Path) -> bool { + root != temp_root && root.starts_with(temp_root) +} + +fn reject_reparse_tree(label: &str, root: &Path) -> Result<(), String> { + reject_reparse_point(label, root)?; + let mut pending = vec![root.to_path_buf()]; + while let Some(directory) = pending.pop() { + for entry in + fs::read_dir(&directory).map_err(|error| format!("{label} cannot be read: {error}"))? + { + let entry = entry.map_err(|error| format!("{label} entry cannot be read: {error}"))?; + let path = entry.path(); + reject_reparse_point(label, &path)?; + if entry + .file_type() + .map_err(|error| format!("{label} entry type unavailable: {error}"))? + .is_dir() + { + pending.push(path); + } + } + } + Ok(()) +} + +fn reject_reparse_point(label: &str, path: &Path) -> Result<(), String> { + let metadata = fs::symlink_metadata(path).map_err(|error| { + format!( + "{label} metadata unavailable for {}: {error}", + path.display() + ) + })?; + if metadata.file_type().is_symlink() || is_windows_reparse_point(&metadata) { + return Err(format!( + "{label} contains a reparse or symlink path: {}", + path.display() + )); + } + Ok(()) +} + +fn reject_reparse_chain(label: &str, path: &Path) -> Result<(), String> { + let mut current = PathBuf::new(); + for component in path.components() { + current.push(component); + if current.exists() { + reject_reparse_point(label, ¤t)?; + } + } + Ok(()) +} + +#[cfg(target_os = "windows")] +fn is_windows_reparse_point(metadata: &fs::Metadata) -> bool { + use std::os::windows::fs::MetadataExt; + + metadata.file_attributes() & 0x400 != 0 +} + +#[cfg(not(target_os = "windows"))] +fn is_windows_reparse_point(_metadata: &fs::Metadata) -> bool { + false +} + +fn paths_overlap(left: &Path, right: &Path) -> bool { + #[cfg(target_os = "windows")] + { + let left = left.to_string_lossy().to_ascii_lowercase(); + let right = right.to_string_lossy().to_ascii_lowercase(); + left == right + || Path::new(&left).starts_with(Path::new(&right)) + || Path::new(&right).starts_with(Path::new(&left)) + } + #[cfg(not(target_os = "windows"))] + { + left == right || left.starts_with(right) || right.starts_with(left) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn manifest_json(fixture_a: &Path, fixture_b: &Path, scratch: &Path) -> String { + serde_json::json!({ + "schema": PROOF_SCHEMA, + "version": PROOF_VERSION, + "kind": PROOF_KIND, + "provider": PROOF_PROVIDER, + "fixtureRoots": { + "geminiCliHome": fixture_a, + "tokscaleConfigDir": fixture_b + }, + "scratchRoot": scratch, + "nowUtc": "2026-09-24T12:34:56Z" + }) + .to_string() + } + + #[test] + fn strict_manifest_accepts_absolute_disjoint_directories() { + let root = crate::test_support::TempDir::new(); + let fixture_a = root.path().join("gemini"); + let fixture_b = root.path().join("tokscale"); + let scratch = root.path().join("scratch"); + fs::create_dir_all(&fixture_a).unwrap(); + fs::create_dir_all(&fixture_b).unwrap(); + fs::create_dir_all(&scratch).unwrap(); + + let proof = + ContainmentProof::from_manifest_json(&manifest_json(&fixture_a, &fixture_b, &scratch)) + .unwrap(); + assert_eq!(proof.manifest.provider, PROOF_PROVIDER); + assert_eq!( + proof.manifest.now_utc.to_rfc3339(), + "2026-09-24T12:34:56+00:00" + ); + } + + #[test] + fn strict_manifest_rejects_unknown_fields_and_wrong_kind() { + let root = crate::test_support::TempDir::new(); + let fixture_a = root.path().join("gemini"); + let fixture_b = root.path().join("tokscale"); + let scratch = root.path().join("scratch"); + fs::create_dir_all(&fixture_a).unwrap(); + fs::create_dir_all(&fixture_b).unwrap(); + fs::create_dir_all(&scratch).unwrap(); + + let mut value: serde_json::Value = + serde_json::from_str(&manifest_json(&fixture_a, &fixture_b, &scratch)).unwrap(); + value["unexpected"] = serde_json::json!(true); + assert!(ContainmentProof::from_manifest_json(&value.to_string()).is_err()); + value.as_object_mut().unwrap().remove("unexpected"); + value["kind"] = serde_json::json!("settings"); + assert!(ContainmentProof::from_manifest_json(&value.to_string()).is_err()); + } + + #[test] + fn path_safety_rejects_nonempty_scratch_and_overlapping_roots() { + let root = crate::test_support::TempDir::new(); + let fixture_a = root.path().join("fixture"); + let fixture_b = root.path().join("fixture").join("nested"); + let scratch = root.path().join("scratch"); + fs::create_dir_all(&fixture_b).unwrap(); + fs::create_dir_all(&scratch).unwrap(); + fs::write(scratch.join("not-empty"), b"x").unwrap(); + assert!( + ContainmentProof::from_manifest_json(&manifest_json(&fixture_a, &fixture_b, &scratch,)) + .is_err() + ); + + fs::remove_file(scratch.join("not-empty")).unwrap(); + assert!( + ContainmentProof::from_manifest_json(&manifest_json(&fixture_a, &fixture_b, &scratch,)) + .is_err() + ); + } + + #[test] + fn fixture_roots_must_be_dedicated_children_of_system_temp() { + let temp_root = std::env::temp_dir(); + assert!(is_dedicated_temp_child( + &temp_root.join("codexbar-proof-fixtures/run-1/gemini"), + &temp_root + )); + assert!(!is_dedicated_temp_child(&temp_root, &temp_root)); + assert!(!is_dedicated_temp_child( + &temp_root.parent().unwrap().join("user-data"), + &temp_root + )); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/shell/mod.rs b/apps/desktop-tauri/src-tauri/src/shell/mod.rs index 6e77233d57..dc8d60f84c 100644 --- a/apps/desktop-tauri/src-tauri/src/shell/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/shell/mod.rs @@ -19,7 +19,9 @@ mod tests; pub(crate) use position::inferred_tray_panel_position_for_monitor_size; pub use position::{remember_current_geometry_if_eligible, tray_panel_position}; -pub use transition::{reopen_to_target, transition_to_target}; +pub use transition::{ + reopen_to_target, transition_to_target, transition_to_target_without_activation, +}; pub use window::hide_to_tray_if_current; #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/apps/desktop-tauri/src-tauri/src/shell/transition.rs b/apps/desktop-tauri/src-tauri/src/shell/transition.rs index fa57b8380c..484f4b9762 100644 --- a/apps/desktop-tauri/src-tauri/src/shell/transition.rs +++ b/apps/desktop-tauri/src-tauri/src/shell/transition.rs @@ -13,7 +13,10 @@ use crate::window_positioner::{self, PanelSize, Rect}; use super::geometry::surface_panel_size; use super::position::default_surface_position; -use super::window::{apply_window_layout, apply_window_properties, show_window}; +use super::window::{ + apply_window_layout, apply_window_properties, apply_window_properties_without_activation, + show_window, +}; use super::{SHELL_TRANSITION_SERIAL, ShellTransitionRequest}; /// Positions from the positioner pipeline are in Tauri's physical coordinate @@ -75,6 +78,12 @@ pub(super) enum TransitionResolution { }, } +#[derive(Clone, Copy)] +pub(super) enum RevealStrategy { + Activate, + NoActivate, +} + pub fn transition_to_target( app: &AppHandle, mode: SurfaceMode, @@ -89,6 +98,27 @@ pub fn transition_to_target( position, }, false, + RevealStrategy::Activate, + ) +} + +/// Transition to a surface using the native no-activate path reserved for +/// isolated containment proof runs. +pub fn transition_to_target_without_activation( + app: &AppHandle, + mode: SurfaceMode, + target: SurfaceTarget, + position: Option<(i32, i32)>, +) -> Result { + apply_transition_request_with_strategy( + app, + ShellTransitionRequest { + mode, + target, + position, + }, + false, + RevealStrategy::NoActivate, ) } @@ -106,6 +136,7 @@ pub fn reopen_to_target( position, }, true, + RevealStrategy::Activate, ) } @@ -113,14 +144,16 @@ fn apply_transition_request_with_strategy( app: &AppHandle, request: ShellTransitionRequest, force_same_mode_apply: bool, + reveal: RevealStrategy, ) -> Result { - apply_transition_request(app, request, force_same_mode_apply) + apply_transition_request(app, request, force_same_mode_apply, reveal) } fn apply_transition_request( app: &AppHandle, request: ShellTransitionRequest, force_same_mode_apply: bool, + reveal: RevealStrategy, ) -> Result { let _transition_guard = SHELL_TRANSITION_SERIAL.lock().unwrap(); let window = app @@ -145,11 +178,17 @@ fn apply_transition_request( .or_else(|| preserved_visible_mode_change_position(&window, &resolution)); match resolution { - TransitionResolution::ModeChange { transition, target } => { - apply_transition(app, &window, &transition, &previous, target, position) - } + TransitionResolution::ModeChange { transition, target } => apply_transition( + app, + &window, + &transition, + &previous, + target, + position, + reveal, + ), TransitionResolution::SameModeRetarget { mode, target } => { - apply_same_mode_target_update(app, &window, mode, target, position) + apply_same_mode_target_update(app, &window, mode, target, position, reveal) } TransitionResolution::SameModeReopen { mode, target } => { let transition = SurfaceTransition { @@ -157,7 +196,15 @@ fn apply_transition_request( to: mode, properties: mode.window_properties(), }; - apply_transition(app, &window, &transition, &previous, target, position) + apply_transition( + app, + &window, + &transition, + &previous, + target, + position, + reveal, + ) } TransitionResolution::Noop { mode } => Ok(mode), } @@ -430,6 +477,7 @@ fn apply_same_mode_target_update( mode: SurfaceMode, target: SurfaceTarget, position: Option<(i32, i32)>, + reveal: RevealStrategy, ) -> Result { if let Some((x, y)) = position { let _ = window.set_position(os_position(window, x, y)); @@ -445,7 +493,14 @@ fn apply_same_mode_target_update( }, )?; events::emit_surface_mode_changed(app, mode, mode, target); - if show_window(window).is_ok() && mode == SurfaceMode::TrayPanel { + let shown = match reveal { + RevealStrategy::Activate => show_window(window).is_ok(), + RevealStrategy::NoActivate => { + crate::proof_runtime::show_window_without_activation(window)?; + true + } + }; + if shown && mode == SurfaceMode::TrayPanel { mark_tray_panel_shown(app); } Ok(mode) @@ -458,6 +513,7 @@ pub(super) fn apply_transition( previous: &SurfaceSnapshot, current_target: SurfaceTarget, position: Option<(i32, i32)>, + reveal: RevealStrategy, ) -> Result { if let Some((x, y)) = position { let _ = window.set_position(os_position(window, x, y)); @@ -485,7 +541,14 @@ pub(super) fn apply_transition( // ever target Hidden/PopOut/Settings, none of which defer their // own reveal.) if needs_show { - let _ = show_window(window); + match reveal { + RevealStrategy::Activate => { + let _ = show_window(window); + } + RevealStrategy::NoActivate => { + crate::proof_runtime::show_window_without_activation(window)?; + } + } } clamp_current_window_to_work_area(window); @@ -494,9 +557,14 @@ pub(super) fn apply_transition( Err(err) => { let recovery = recovery_snapshot_for_failed_transition(transition, previous, ¤t_target); - if let Err(recovery_err) = restore_recovery_surface(&recovery, |mode, properties| { - apply_window_properties(window, mode, properties) - }) { + if let Err(recovery_err) = + restore_recovery_surface(&recovery, |mode, properties| match reveal { + RevealStrategy::Activate => apply_window_properties(window, mode, properties), + RevealStrategy::NoActivate => { + apply_window_properties_without_activation(window, mode, properties) + } + }) + { let hidden = hidden_surface_snapshot(); if let Err(hide_err) = window.hide().map_err(|e| e.to_string()) { tracing::warn!( diff --git a/apps/desktop-tauri/src-tauri/src/shell/window.rs b/apps/desktop-tauri/src-tauri/src/shell/window.rs index 256926ccc9..24864314d6 100644 --- a/apps/desktop-tauri/src-tauri/src/shell/window.rs +++ b/apps/desktop-tauri/src-tauri/src/shell/window.rs @@ -9,7 +9,9 @@ use crate::surface::{SurfaceMode, SurfaceTransition, WindowProperties}; use crate::surface_target::SurfaceTarget; use super::SHELL_TRANSITION_SERIAL; -use super::transition::{SurfaceSnapshot, apply_transition, current_surface_snapshot}; +use super::transition::{ + RevealStrategy, SurfaceSnapshot, apply_transition, current_surface_snapshot, +}; pub(super) struct HideToTrayPlan { pub previous: SurfaceSnapshot, @@ -30,6 +32,20 @@ pub fn apply_window_properties( Ok(()) } +/// Apply the requested surface without giving the proof window input focus. +/// This path is used only by the isolated containment proof harness. +pub fn apply_window_properties_without_activation( + window: &WebviewWindow, + mode: SurfaceMode, + props: &WindowProperties, +) -> Result<(), String> { + let needs_show = apply_window_layout(window, mode, props)?; + if needs_show { + crate::proof_runtime::show_window_without_activation(window)?; + } + Ok(()) +} + /// Apply layout properties (decorations, size, always-on-top) WITHOUT making /// the window visible. Returns `true` when the caller should subsequently /// call [`show_window`] to make it visible, or `false` when the mode hides @@ -181,7 +197,16 @@ where }; if let Some(transition) = plan.transition { - apply_transition(app, &window, &transition, &plan.previous, plan.target, None).map(Some) + apply_transition( + app, + &window, + &transition, + &plan.previous, + plan.target, + None, + RevealStrategy::Activate, + ) + .map(Some) } else { let _ = window.hide(); Ok(Some(SurfaceMode::Hidden)) diff --git a/apps/desktop-tauri/src-tauri/src/state.rs b/apps/desktop-tauri/src-tauri/src/state.rs index 174c13d113..d5f5637d4f 100644 --- a/apps/desktop-tauri/src-tauri/src/state.rs +++ b/apps/desktop-tauri/src-tauri/src/state.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use crate::commands::ProviderUsageSnapshot; use crate::proof_harness::ProofConfig; +use crate::proof_runtime::ContainmentProof; use crate::surface::{SurfaceMode, SurfaceStateMachine, SurfaceTransition}; use crate::surface_target::SurfaceTarget; @@ -143,6 +144,10 @@ pub struct AppState { pub installer_path: Option, /// Proof-harness configuration (set when `CODEXBAR_PROOF_MODE` is active). pub proof_config: Option, + /// Strict, isolated containment proof runtime, when active. + pub containment_proof: Option, + /// In-memory settings route used by containment proof consumers. + pub proof_settings: Option, /// True only after a validated proof seed was installed at startup. pub provider_cache_seeded: bool, /// Persistent notification manager — tracks which alerts have fired to prevent spam. @@ -206,6 +211,8 @@ impl AppState { last_update_check_ms: None, installer_path: None, proof_config: None, + containment_proof: None, + proof_settings: None, provider_cache_seeded: false, notification_manager: codexbar::notifications::NotificationManager::new(), last_shown_at: None, @@ -217,6 +224,21 @@ impl AppState { } } + pub fn new_for_containment_proof(proof: ContainmentProof) -> Self { + let mut state = Self::new(); + state.proof_settings = Some(proof.proof_settings()); + state.containment_proof = Some(proof); + state + } + + pub fn is_containment_proof(&self) -> bool { + self.containment_proof.is_some() + } + + pub fn proof_settings(&self) -> Option<&codexbar::settings::Settings> { + self.proof_settings.as_ref() + } + pub fn mark_blur_dismissed(&mut self, dismissed_at: std::time::Instant) { self.last_blur_dismissed_at = Some(dismissed_at); } diff --git a/apps/desktop-tauri/src-tauri/src/test_support.rs b/apps/desktop-tauri/src-tauri/src/test_support.rs new file mode 100644 index 0000000000..dce065fd64 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/test_support.rs @@ -0,0 +1,77 @@ +//! Small test-only helpers that avoid adding runtime or dev dependencies. + +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +const TEMP_PREFIX: &str = "win-codexbar-test-"; + +pub struct TempDir { + path: PathBuf, +} + +impl TempDir { + pub fn new() -> Self { + let root = fs::canonicalize(std::env::temp_dir()) + .expect("system temporary directory must be available"); + let nonce = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock must be after Unix epoch") + .as_nanos(); + for attempt in 0..32 { + let path = root.join(format!( + "{TEMP_PREFIX}{}-{nonce}-{attempt}", + std::process::id() + )); + match fs::create_dir(&path) { + Ok(()) => return Self { path }, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => panic!("cannot create test directory {}: {error}", path.display()), + } + } + panic!( + "could not allocate a unique test directory under {}", + root.display() + ); + } + + pub fn path(&self) -> &Path { + &self.path + } +} + +impl Drop for TempDir { + fn drop(&mut self) { + let Ok(temp_root) = fs::canonicalize(std::env::temp_dir()) else { + return; + }; + let Ok(canonical) = fs::canonicalize(&self.path) else { + return; + }; + if canonical.parent() != Some(temp_root.as_path()) + || !canonical + .file_name() + .is_some_and(|name| name.to_string_lossy().starts_with(TEMP_PREFIX)) + { + return; + } + let Ok(metadata) = fs::symlink_metadata(&self.path) else { + return; + }; + if metadata.file_type().is_symlink() || is_reparse_point(&metadata) { + return; + } + let _ = fs::remove_dir_all(canonical); + } +} + +#[cfg(windows)] +fn is_reparse_point(metadata: &fs::Metadata) -> bool { + use std::os::windows::fs::MetadataExt; + metadata.file_attributes() & 0x400 != 0 +} + +#[cfg(not(windows))] +fn is_reparse_point(_metadata: &fs::Metadata) -> bool { + false +} diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index eb277cf2b0..b656562b17 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -250,7 +250,11 @@ fn store_anchor(app: &AppHandle, rect: &tauri::Rect, click_position: tauri::Phys /// - **Left-click** toggles the custom tray panel via the surface state machine. /// - **Right-click** opens the native context menu with shell actions. pub fn setup(app: &mut tauri::App) -> Result<(), Box> { - let menu = build_native_tray_menu(app.handle(), &crate::commands::get_provider_catalog(), &[])?; + let menu = build_native_tray_menu( + app.handle(), + &crate::commands::get_provider_catalog_for_current_settings(), + &[], + )?; // Embed the icon at compile time so it works regardless of working directory. let icon_bytes = include_bytes!("../../../../rust/icons/icon.png"); @@ -400,7 +404,7 @@ fn handle_menu_event(app: &AppHandle, id: &str) { /// Rebuild the native tray menu from current provider + settings state. pub(crate) fn rebuild_tray_menu(app: &AppHandle) { - let catalog = crate::commands::get_provider_catalog(); + let catalog = crate::commands::get_provider_catalog_for_current_settings(); let settings = Settings::load(); let status_labels = if let Some(st) = app.try_state::>() { let guard = st.lock().unwrap(); @@ -421,7 +425,7 @@ pub fn update_tray_status_items( app: &AppHandle, snapshots: &[crate::commands::ProviderUsageSnapshot], ) { - let catalog = crate::commands::get_provider_catalog(); + let catalog = crate::commands::get_provider_catalog_for_current_settings(); let settings = Settings::load(); let status_labels = TrayPresentationPlan::resolve(&settings, snapshots).status_labels(settings.ui_language); diff --git a/rust/src/logging.rs b/rust/src/logging.rs index 96efbfa5b1..159e850198 100755 --- a/rust/src/logging.rs +++ b/rust/src/logging.rs @@ -11,7 +11,7 @@ use std::io::Write as _; use std::path::PathBuf; -use std::sync::LazyLock; +use std::sync::{LazyLock, OnceLock}; use tracing_subscriber::{EnvFilter, fmt, prelude::*}; /// Convert a displayable error into a frontend/log-safe message. @@ -21,7 +21,10 @@ pub fn safe_error_message(err: impl std::fmt::Display) -> String { /// Canonical application config root that hosts the settings file and logs. pub fn config_root() -> Option { - dirs::config_dir().map(|p| p.join("CodexBar")) + CONFIG_ROOT_OVERRIDE + .get() + .cloned() + .or_else(|| dirs::config_dir().map(|p| p.join("CodexBar"))) } /// Settings directory that hosts the app settings file (also the log root). @@ -37,6 +40,20 @@ pub const LOG_MAX_BYTES: u64 = 1024 * 1024; pub const LOG_FILE_STEM_CLI: &str = "codexbar-cli"; pub const LOG_FILE_STEM_DESKTOP: &str = "codexbar-desktop"; +static CONFIG_ROOT_OVERRIDE: OnceLock = OnceLock::new(); + +/// Install a process-local config root before logging or settings are first +/// touched. The desktop containment proof uses this to keep settings and logs +/// out of the user's normal profile. +pub fn install_config_root_override(root: PathBuf) -> Result<(), String> { + if !root.is_absolute() { + return Err(format!("config root must be absolute: {}", root.display())); + } + CONFIG_ROOT_OVERRIDE + .set(root) + .map_err(|_| "config root override was already installed".to_string()) +} + static LOG_FILE_STEM: LazyLock<&'static str> = LazyLock::new(|| { // The Tauri shell sets CODEXBAR_PROCESS=desktop before logging::init; // everything else (the `codexbar` binary, tests) gets the CLI name. diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs index 00519d93ca..464201feab 100644 --- a/rust/src/providers/antigravity/local_history.rs +++ b/rust/src/providers/antigravity/local_history.rs @@ -40,6 +40,17 @@ fn summarize_local_usage_from( } } +pub(super) fn summarize_local_usage_from_explicit_roots( + database_roots: &[PathBuf], + jsonl_sessions_root: &Path, + now: DateTime, + days: u32, +) -> LocalTokenHistorySummary { + summarize_local_usage_from(database_roots, now, days, || { + local_sessions::summarize_jsonl_at(jsonl_sessions_root, now, days) + }) +} + pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { let now = Utc::now(); let Some(home) = dirs::home_dir() else { @@ -47,9 +58,7 @@ pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { }; let roots = configured_database_roots(&home); let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); - summarize_local_usage_from(&roots, now, days, || { - local_sessions::summarize_jsonl_at(&tokscale_sessions, now, days) - }) + summarize_local_usage_from_explicit_roots(&roots, &tokscale_sessions, now, days) } /// Count local Antigravity conversation artifacts for the quota provider's diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 4efb9af063..d2d1be7b74 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -2,3 +2,69 @@ pub use super::local_history::{offline_conversation_count, summarize_local_usage pub use crate::spend_contract::{ LocalHistoryCoverage, LocalTokenHistorySummary as LocalSessionSummary, }; + +use chrono::{DateTime, Utc}; +use std::path::{Path, PathBuf}; + +pub fn summarize_from_roots( + database_roots: &[PathBuf], + jsonl_sessions_root: &Path, + now: DateTime, + days: u32, +) -> LocalSessionSummary { + super::local_history::summarize_local_usage_from_explicit_roots( + database_roots, + jsonl_sessions_root, + now, + days, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + use std::fs; + + #[test] + fn explicit_roots_use_only_caller_paths_and_fixed_window() { + let dir = tempfile::tempdir().unwrap(); + let caller_database_root = dir.path().join("caller-databases"); + let caller_root = dir.path().join("caller-sessions"); + let unrelated_root = dir.path().join("unrelated-sessions"); + fs::create_dir_all(&caller_database_root).unwrap(); + fs::create_dir_all(&caller_root).unwrap(); + fs::create_dir_all(&unrelated_root).unwrap(); + let connection = + rusqlite::Connection::open(caller_database_root.join("foreign.db")).unwrap(); + connection + .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) + .unwrap(); + + fs::write( + caller_root.join("caller.jsonl"), + concat!( + "{\"type\":\"usage\",\"responseId\":\"in-window\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", + "{\"type\":\"usage\",\"responseId\":\"old\",\"timestamp\":1784894400000,\"input\":900,\"output\":90}\n" + ), + ) + .unwrap(); + fs::write( + unrelated_root.join("unrelated.jsonl"), + b"{\"type\":\"usage\",\"responseId\":\"unrelated\",\"timestamp\":1787572800000,\"input\":9000,\"output\":900}\n", + ) + .unwrap(); + + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_from_roots( + std::slice::from_ref(&caller_database_root), + &caller_root, + now, + 7, + ); + + assert_eq!(summary.total_tokens, 120); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } +}