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 001/259] 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 002/259] 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 003/259] 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 004/259] 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 005/259] 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 006/259] 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 007/259] 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 008/259] 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 009/259] 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 010/259] 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 011/259] 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 012/259] 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 013/259] 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 014/259] 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 015/259] 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 016/259] 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 017/259] 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 018/259] 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 019/259] 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 020/259] 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 021/259] 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 022/259] 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 023/259] 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 024/259] 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 025/259] 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 026/259] 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 027/259] 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 028/259] 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 029/259] 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 030/259] 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 031/259] 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 032/259] 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 033/259] 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 034/259] 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 035/259] 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 036/259] 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 037/259] 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 038/259] 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 039/259] 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 a0de08cc0cb6d8f90c68417d45ed87004e194fb7 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:57:13 +0700 Subject: [PATCH 040/259] Animate budget disclosure chevron --- apps/desktop-tauri/src/surfaces/menu-surface-layout.css | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/apps/desktop-tauri/src/surfaces/menu-surface-layout.css b/apps/desktop-tauri/src/surfaces/menu-surface-layout.css index e79ac3293a..60961fe029 100644 --- a/apps/desktop-tauri/src/surfaces/menu-surface-layout.css +++ b/apps/desktop-tauri/src/surfaces/menu-surface-layout.css @@ -78,10 +78,16 @@ .menu-metric__budget-header::before { content: "›"; margin-right: 4px; + transition: transform 120ms ease; } .menu-metric__budget-header[aria-expanded="true"]::before { transform: rotate(90deg); } +@media (prefers-reduced-motion: reduce) { + .menu-metric__budget-header::before { + transition: none; + } +} .codex-menu-accounts { margin: 0; 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 041/259] 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 32734c321b445905750423bc345783e52388512c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:19:31 +0700 Subject: [PATCH 042/259] Add preferred display currency support --- .../src-tauri/src/commands/bridge.rs | 20 ++ .../src-tauri/src/commands/currency.rs | 177 ++++++++++++++++++ .../src-tauri/src/commands/mod.rs | 2 + .../src-tauri/src/commands/settings.rs | 30 +++ apps/desktop-tauri/src-tauri/src/main.rs | 2 + .../src-tauri/src/tray_bridge.rs | 6 +- apps/desktop-tauri/src/App.tsx | 5 +- .../src/components/MenuCardDetails.tsx | 67 ++----- apps/desktop-tauri/src/floatbar/FloatBar.tsx | 11 +- .../src/hooks/CurrencyProvider.tsx | 72 +++++++ apps/desktop-tauri/src/i18n/keys.ts | 2 + apps/desktop-tauri/src/lib/currency.test.ts | 50 +++++ apps/desktop-tauri/src/lib/currency.ts | 117 ++++++++++++ apps/desktop-tauri/src/lib/tauri.ts | 5 + .../src/lib/usageSpendSharing.test.ts | 39 +++- .../src/lib/usageSpendSharing.ts | 22 ++- apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 23 ++- .../providers/sections/CostSection.tsx | 8 +- .../ClaudeSwapAccountsSection.test.tsx | 23 ++- .../credentials/ClaudeSwapAccountsSection.tsx | 15 +- .../settings/tabs/GeneralTab.test.tsx | 17 ++ .../src/surfaces/settings/tabs/GeneralTab.tsx | 16 ++ .../surfaces/settings/tabs/UsageSpendTab.tsx | 48 ++++- apps/desktop-tauri/src/types/bridge.ts | 6 + rust/src/currency.rs | 175 +++++++++++++++++ rust/src/lib.rs | 1 + rust/src/locale.rs | 2 + rust/src/locale/en-US.ftl | 2 + rust/src/settings.rs | 9 + rust/src/settings/raw.rs | 10 + rust/src/settings/tests.rs | 22 +++ 31 files changed, 909 insertions(+), 95 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/currency.rs create mode 100644 apps/desktop-tauri/src/hooks/CurrencyProvider.tsx create mode 100644 apps/desktop-tauri/src/lib/currency.test.ts create mode 100644 apps/desktop-tauri/src/lib/currency.ts create mode 100644 rust/src/currency.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ab5efbef8a..06ceb78be0 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -126,6 +126,24 @@ fn default_cost_period() -> String { /// otherwise falling back to the currency-code prefix. Used by tray surfaces /// that render a spend amount without a rate-window percent (MonthlyPlan). pub(crate) fn format_cost_amount(cost: &CostSnapshotBridge) -> String { + if let Some((amount, currency)) = + crate::commands::convert_preferred_amount(cost.used, &cost.currency_code) + { + let symbol = match currency.as_str() { + "USD" => Some("$"), + "EUR" => Some("€"), + "GBP" => Some("£"), + "TRY" => Some("₺"), + _ => None, + }; + return symbol.map_or_else( + || format!("{amount:.2} {currency}"), + |symbol| format!("{symbol}{amount:.2}"), + ); + } + if !cost.formatted_used.is_empty() { + return cost.formatted_used.clone(); + } if let Some(ref symbol) = cost.currency_symbol { format!("{}{:.2}", symbol, cost.used) } else { @@ -641,6 +659,7 @@ pub struct ProviderCatalogEntry { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct SettingsSnapshot { + preferred_currency_code: String, enabled_providers: Vec, provider_order: Vec, refresh_interval_secs: u64, @@ -763,6 +782,7 @@ impl From for SettingsSnapshot { .collect(); Self { + preferred_currency_code: settings.preferred_currency_code, enabled_providers, provider_order, refresh_interval_secs: settings.refresh_interval_secs, diff --git a/apps/desktop-tauri/src-tauri/src/commands/currency.rs b/apps/desktop-tauri/src-tauri/src/commands/currency.rs new file mode 100644 index 0000000000..2c46394b2d --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/currency.rs @@ -0,0 +1,177 @@ +use std::{collections::HashMap, path::PathBuf}; + +use codexbar::currency::{ + SUPPORTED_CURRENCY_CODES, convert_amount, fallback_rates, fetch_exchange_rates, + normalize_preferred_currency, +}; +use serde::{Deserialize, Serialize}; +use tauri::State; +use tokio::sync::Mutex; + +const CACHE_MAX_AGE_SECS: i64 = 24 * 60 * 60; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct PersistedRates { + fetched_at_unix: i64, + rates: HashMap, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct CurrencyRatesSnapshot { + pub rates: HashMap, +} + +#[derive(Default)] +pub struct CurrencyRateCache { + inner: Mutex, +} + +#[derive(Default)] +struct CacheState { + loaded: bool, + persisted: Option, +} + +#[tauri::command] +pub async fn get_currency_rates( + app: tauri::AppHandle, + cache: State<'_, CurrencyRateCache>, + preferred_currency_code: String, +) -> CurrencyRatesSnapshot { + let preferred = normalize_preferred_currency(&preferred_currency_code); + let mut state = cache.inner.lock().await; + if !state.loaded { + state.persisted = read_persisted_rates(); + state.loaded = true; + } + + if preferred != "AUTO" { + let now = unix_now(); + let fresh = state.persisted.as_ref().is_some_and(|cached| { + let age = now.saturating_sub(cached.fetched_at_unix); + (0..CACHE_MAX_AGE_SECS).contains(&age) + }); + if !fresh { + match fetch_exchange_rates().await { + Ok(rates) => { + let entry = PersistedRates { + fetched_at_unix: now, + rates, + }; + persist_rates(&entry); + state.persisted = Some(entry); + crate::tray_bridge::refresh_tray_presentation(&app); + } + Err(error) => tracing::debug!(%error, "currency rates unavailable; using cached or offline rates"), + } + } + } + + CurrencyRatesSnapshot { + rates: merged_rates(state.persisted.as_ref()), + } +} + +pub(crate) fn convert_preferred_amount(amount: f64, source_code: &str) -> Option<(f64, String)> { + let preferred = normalize_preferred_currency( + &codexbar::settings::Settings::load().preferred_currency_code, + ); + if preferred == "AUTO" { + return None; + } + let cached = read_persisted_rates(); + let rates = merged_rates(cached.as_ref()); + convert_amount(amount, source_code, &preferred, &rates).map(|converted| (converted, preferred)) +} + +fn unix_now() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_secs().min(i64::MAX as u64) as i64) + .unwrap_or_default() +} + +fn cache_path() -> Option { + codexbar::settings::Settings::settings_path()? + .parent() + .map(|parent| parent.join("currency-rates.json")) +} + +fn clean_rates(rates: HashMap) -> HashMap { + let mut clean = HashMap::new(); + for code in SUPPORTED_CURRENCY_CODES { + if let Some(rate) = rates.get(*code).copied() + && rate.is_finite() + && rate > 0.0 + { + clean.insert((*code).to_string(), rate); + } + } + if clean.get("USD").is_none_or(|rate| (*rate - 1.0).abs() > f64::EPSILON) { + return HashMap::new(); + } + clean +} + +fn read_persisted_rates() -> Option { + let path = cache_path()?; + let bytes = std::fs::read(path).ok()?; + let mut cached: PersistedRates = serde_json::from_slice(&bytes).ok()?; + cached.rates = clean_rates(cached.rates); + (!cached.rates.is_empty()).then_some(cached) +} + +fn merged_rates(cached: Option<&PersistedRates>) -> HashMap { + let mut rates = fallback_rates(); + if let Some(cached) = cached { + for (code, rate) in clean_rates(cached.rates.clone()) { + rates.insert(code, rate); + } + } + rates +} + +fn persist_rates(cached: &PersistedRates) { + let Some(path) = cache_path() else { return }; + let Some(parent) = path.parent() else { return }; + let write_result = (|| -> Result<(), Box> { + std::fs::create_dir_all(parent)?; + let bytes = serde_json::to_vec(cached)?; + codexbar::atomic_file::write_atomic(&path, &bytes)?; + Ok(()) + })(); + if let Err(error) = write_result { + tracing::debug!(%error, "could not persist currency rate cache"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codexbar::currency::FALLBACK_RATES; + + #[test] + fn fallback_table_covers_every_preferred_currency() { + let rates = fallback_rates(); + assert_eq!(rates.len(), FALLBACK_RATES.len()); + for code in SUPPORTED_CURRENCY_CODES { + assert!(rates.get(*code).is_some_and(|rate| rate.is_finite() && *rate > 0.0)); + } + } + + #[test] + fn persisted_rates_are_sanitized_before_merging() { + let rates = HashMap::from([ + ("USD".to_string(), 1.0), + ("TRY".to_string(), 48.0), + ("EUR".to_string(), f64::NAN), + ("BTC".to_string(), 100.0), + ]); + let clean = clean_rates(rates); + assert_eq!(clean.get("TRY"), Some(&48.0)); + assert!(!clean.contains_key("EUR")); + assert!(!clean.contains_key("BTC")); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index fdfaeda967..fda0557740 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -24,6 +24,7 @@ use crate::surface::SurfaceMode; use crate::surface_target::SurfaceTarget; mod chart; +mod currency; mod spend_contract; mod tokens; mod updater; @@ -79,6 +80,7 @@ pub(crate) use usage_items::*; mod tests; pub use chart::*; +pub use currency::*; pub use spend_contract::*; pub use tokens::*; pub use updater::*; diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..f85da91d88 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -7,6 +7,7 @@ use super::*; #[derive(Debug, Clone, Default, Deserialize)] #[serde(default, rename_all = "camelCase")] pub struct SettingsUpdate { + pub preferred_currency_code: Option, pub enabled_providers: Option>, pub refresh_interval_secs: Option, pub adaptive_refresh: Option, @@ -130,6 +131,7 @@ impl SettingsUpdate { || self.overview_layout.is_some() || self.provider_metrics.is_some() || self.provider_hidden_usage_item_ids.is_some() + || self.preferred_currency_code.is_some() || self.codex_spark_usage_visible.is_some() || self.copilot_seat_credit_entitlement.is_some() || self.enabled_providers.is_some() @@ -200,6 +202,15 @@ impl SettingsUpdate { } fn apply_general_settings(self, settings: &mut Settings) -> Result { + if let Some(value) = self.preferred_currency_code.as_deref() { + let normalized = codexbar::currency::normalize_preferred_currency(value); + if !value.trim().eq_ignore_ascii_case("AUTO") + && normalized == "AUTO" + { + return Err(format!("Unsupported preferred currency: {value}")); + } + settings.preferred_currency_code = normalized; + } if let Some(v) = self.start_at_login { settings.set_start_at_login(v).map_err(|e| e.to_string())?; } @@ -572,6 +583,25 @@ pub async fn update_settings( mod tests { use super::*; + #[test] + fn preferred_currency_patch_accepts_supported_codes_and_rejects_unknown_codes() { + let mut settings = Settings::default(); + SettingsUpdate { + preferred_currency_code: Some("try".to_string()), + ..SettingsUpdate::default() + } + .apply_to(&mut settings) + .expect("TRY is supported"); + assert_eq!(settings.preferred_currency_code, "TRY"); + + let result = SettingsUpdate { + preferred_currency_code: Some("BTC".to_string()), + ..SettingsUpdate::default() + } + .apply_to(&mut settings); + assert!(matches!(result, Err(error) if error.contains("Unsupported preferred currency"))); + } + #[test] fn only_data_affecting_settings_refresh_providers() { assert!( diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 02f48672b6..28c6179b87 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -146,6 +146,7 @@ fn main() { tauri::Builder::default() .manage(Mutex::new(initial_state)) + .manage(commands::CurrencyRateCache::default()) .plugin(shortcut_bridge::plugin()) .plugin(tauri_plugin_dialog::init()) .plugin(tauri_plugin_single_instance::init(|app, args, _cwd| { @@ -159,6 +160,7 @@ fn main() { commands::get_bootstrap_state, commands::get_provider_catalog, commands::get_settings_snapshot, + commands::get_currency_rates, commands::list_agent_sessions, commands::focus_agent_session, commands::update_settings, diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index aebb7c2115..5ddda446b1 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -564,11 +564,7 @@ fn provider_status_label( 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) - }; + let amount = crate::commands::format_cost_amount(cost); return ( snapshot.provider_id.clone(), format!("{} {}", snapshot.display_name, amount), diff --git a/apps/desktop-tauri/src/App.tsx b/apps/desktop-tauri/src/App.tsx index aa775058ad..7550c8136e 100644 --- a/apps/desktop-tauri/src/App.tsx +++ b/apps/desktop-tauri/src/App.tsx @@ -17,6 +17,7 @@ import { LocaleProvider } from "./i18n/LocaleProvider"; import type { BootstrapState, ThemePreference } from "./types/bridge"; import type { SurfaceSnapshot } from "./hooks/useSurfaceSnapshot"; import { useDeepSeekPricingStatus } from "./hooks/useDeepSeekPricingStatus"; +import { CurrencyProvider } from "./hooks/CurrencyProvider"; const Settings = lazy(() => import("./surfaces/Settings")); const PopOutPanel = lazy(() => import("./surfaces/PopOutPanel")); @@ -50,7 +51,9 @@ function initialSettingsTab(): string { export default function App() { return ( - + + + ); } diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 31634bd910..0b674e2f97 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -12,6 +12,7 @@ import type { SessionEquivalentForecastSnapshot, } from "../types/bridge"; import { useLocale } from "../hooks/useLocale"; +import { useCurrency } from "../hooks/CurrencyProvider"; import { providerAllowsPace } from "../lib/providerPace"; import { useFormattedResetTime, @@ -63,7 +64,6 @@ function formatSessionEquivalentEstimate( return `Estimated: ${display} ${unit} left`; } -const currencyFormatters = new Map(); const compactCountFormat0 = new Intl.NumberFormat("en-US", { notation: "compact", maximumFractionDigits: 0, @@ -73,22 +73,6 @@ const compactCountFormat1 = new Intl.NumberFormat("en-US", { maximumFractionDigits: 1, }); -function formatCurrency(amount: number, code: string): string { - try { - let formatter = currencyFormatters.get(code); - if (!formatter) { - formatter = new Intl.NumberFormat("en-US", { - style: "currency", - currency: code, - }); - currencyFormatters.set(code, formatter); - } - return formatter.format(amount); - } catch { - return `${code} ${amount.toFixed(2)}`; - } -} - function formatCompactCount(value: number | null): string { if (value == null || value <= 0) return "—"; return (value >= 1_000_000 ? compactCountFormat1 : compactCountFormat0).format( @@ -112,6 +96,7 @@ function LocalUsageBlock({ costHistory: DailyCostPoint[]; }) { const { t } = useLocale(); + const { format } = useCurrency(); const isCodex = providerId === "codex"; const isMuse = providerId === "muse"; const visibleHistory = costHistory.slice(-30); @@ -131,7 +116,7 @@ function LocalUsageBlock({ ? formatCompactCount(summary.latestTokens) : "—") : summary.todayCost != null - ? formatCurrency(summary.todayCost, "USD") + ? format(summary.todayCost, "USD") : "—"} @@ -140,7 +125,7 @@ function LocalUsageBlock({ {t("PanelThirtyDayCost")} {summary.thirtyDayCost != null - ? formatCurrency(summary.thirtyDayCost, "USD") + ? format(summary.thirtyDayCost, "USD") : "—"} @@ -166,7 +151,7 @@ function LocalUsageBlock({ height: `${point.value == null || maxCost <= 0 ? 1 : Math.max(4, Math.round((point.value / maxCost) * 64))}px`, opacity: point.value == null ? 0 : undefined, }} - title={point.value == null ? point.date : `${point.date}: ${formatCurrency(point.value, "USD")}`} + title={point.value == null ? point.date : `${point.date}: ${format(point.value, "USD")}`} /> ))} @@ -528,6 +513,15 @@ export default function MenuCardDetails({ onLayoutChange, }: MenuCardDetailsProps) { const { t } = useLocale(); + const { format, preferredCode } = useCurrency(); + const formatProviderCost = ( + amount: number | null | undefined, + currencyCode: string, + sourceSymbol?: string | null, + sourceFormatted?: string | null, + ) => preferredCode === "AUTO" && sourceFormatted + ? sourceFormatted + : format(amount, currencyCode, sourceSymbol); const paceEnabled = display.showPace !== false && providerAllowsPace(provider.providerId, provider.sourceLabel); @@ -626,49 +620,30 @@ export default function MenuCardDetails({ {provider.cost.balance != null && provider.cost.limit == null ? (
- {provider.cost.formattedBalance || - formatCurrency( - provider.cost.balance, - provider.cost.currencyCode, - )} + {formatProviderCost(provider.cost.balance, provider.cost.currencyCode, provider.cost.currencySymbol, provider.cost.formattedBalance)}
) : ( <>
{t("DetailCostUsed")}:{" "} - {provider.cost.formattedUsed || - formatCurrency( - provider.cost.used, - provider.cost.currencyCode, - )} + {formatProviderCost(provider.cost.used, provider.cost.currencyCode, provider.cost.currencySymbol, provider.cost.formattedUsed)} {provider.cost.limit != null && ( <> {" / "} - {provider.cost.formattedLimit || - formatCurrency( - provider.cost.limit, - provider.cost.currencyCode, - )} + {formatProviderCost(provider.cost.limit, provider.cost.currencyCode, provider.cost.currencySymbol, provider.cost.formattedLimit)} )}
{costStyle === "detailed" && provider.cost.balance != null && (
{t("DetailCostBalance")}:{" "} - {provider.cost.formattedBalance || - formatCurrency( - provider.cost.balance, - provider.cost.currencyCode, - )} + {formatProviderCost(provider.cost.balance, provider.cost.currencyCode, provider.cost.currencySymbol, provider.cost.formattedBalance)}
)} {costStyle === "detailed" && provider.cost.remaining != null && (
{t("DetailCostRemaining")}:{" "} - {formatCurrency( - provider.cost.remaining, - provider.cost.currencyCode, - )} + {format(provider.cost.remaining, provider.cost.currencyCode, provider.cost.currencySymbol)}
)} {costStyle === "detailed" && formattedCostReset && ( @@ -681,9 +656,7 @@ export default function MenuCardDetails({ {provider.providerId === "mistral" && provider.cost && (
{t("MistralMonthlySpend")}:{" "} - {provider.cost.currencySymbol - ? `${provider.cost.currencySymbol}${provider.cost.used.toFixed(2)}` - : provider.cost.formattedUsed} + {formatProviderCost(provider.cost.used, provider.cost.currencyCode, provider.cost.currencySymbol, provider.cost.formattedUsed)}
)} diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.tsx index e8d71af8fe..4a55f2faa5 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.tsx @@ -10,6 +10,7 @@ import { import { listen } from "@tauri-apps/api/event"; import { getCurrentWindow } from "@tauri-apps/api/window"; import { useFormattedResetTime } from "../hooks/useFormattedResetTime"; +import { useCurrency } from "../hooks/CurrencyProvider"; import { useLocale } from "../hooks/useLocale"; import { useProviders } from "../hooks/useProviders"; import { @@ -109,11 +110,6 @@ function hasLocalCost(summary: ProviderLocalUsageSummary | null): summary is Pro return summary?.todayCost != null || summary?.thirtyDayCost != null; } -function formatUsd(value: number | null): string | null { - if (value == null || !Number.isFinite(value)) return null; - return `$${value.toFixed(2)}`; -} - function CostPill({ summary, scale, @@ -127,8 +123,9 @@ function CostPill({ thirtyDayLabel: string; estimateLabel: string; }) { - const today = formatUsd(summary.todayCost); - const thirtyDay = formatUsd(summary.thirtyDayCost); + const { format } = useCurrency(); + const today = summary.todayCost == null ? null : format(summary.todayCost, "USD"); + const thirtyDay = summary.thirtyDayCost == null ? null : format(summary.thirtyDayCost, "USD"); const iconSize = Math.round(10 * scale); const brand = getProviderIcon(summary.providerId).brandColor; const title = [ diff --git a/apps/desktop-tauri/src/hooks/CurrencyProvider.tsx b/apps/desktop-tauri/src/hooks/CurrencyProvider.tsx new file mode 100644 index 0000000000..c31cf90029 --- /dev/null +++ b/apps/desktop-tauri/src/hooks/CurrencyProvider.tsx @@ -0,0 +1,72 @@ +import { createContext, useCallback, useContext, useEffect, useRef, useState, type ReactNode } from "react"; +import { listen } from "@tauri-apps/api/event"; +import { getCurrencyRates, getSettingsSnapshot } from "../lib/tauri"; +import { FALLBACK_CURRENCY_RATES, formatDisplayCurrency, mergeValidCurrencyRates, normalizePreferredCurrency } from "../lib/currency"; +import type { SettingsSnapshot } from "../types/bridge"; + +interface CurrencyContextValue { + preferredCode: string; + rates: Record; + format: (amount: number | null | undefined, sourceCode: string, sourceSymbol?: string | null) => string; +} + +const CurrencyContext = createContext({ + preferredCode: "AUTO", + rates: FALLBACK_CURRENCY_RATES, + format: (amount, sourceCode, sourceSymbol) => formatDisplayCurrency(amount, sourceCode, "AUTO", FALLBACK_CURRENCY_RATES, sourceSymbol), +}); + +export function CurrencyProvider({ children }: { children: ReactNode }) { + const [preferredCode, setPreferredCode] = useState("AUTO"); + const [rates, setRates] = useState(FALLBACK_CURRENCY_RATES); + const requestId = useRef(0); + + const applySettings = useCallback((settings: SettingsSnapshot) => { + const selected = normalizePreferredCurrency(settings.preferredCurrencyCode); + setPreferredCode(selected); + if (selected === "AUTO") { + requestId.current += 1; + return; + } + const id = ++requestId.current; + void getCurrencyRates(selected) + .then((snapshot) => { + if (requestId.current === id) setRates(mergeValidCurrencyRates(snapshot.rates)); + }) + .catch(() => { + // Keep the offline fallback; exchange-rate availability never blocks app surfaces. + }); + }, []); + + useEffect(() => { + let active = true; + void getSettingsSnapshot().then((settings) => { if (active) applySettings(settings); }).catch(() => {}); + const onUpdated = (event: Event) => { + const settings = (event as CustomEvent).detail; + if (settings) applySettings(settings); + }; + window.addEventListener("codexbar:settings-updated", onUpdated); + let unlisten: (() => void) | undefined; + void listen("settings-changed", () => { + void getSettingsSnapshot().then((settings) => { if (active) applySettings(settings); }).catch(() => {}); + }).then((stop) => { if (active) unlisten = stop; else stop(); }).catch(() => {}); + return () => { + active = false; + requestId.current += 1; + window.removeEventListener("codexbar:settings-updated", onUpdated); + unlisten?.(); + }; + }, [applySettings]); + + const format = useCallback( + (amount: number | null | undefined, sourceCode: string, sourceSymbol?: string | null) => + formatDisplayCurrency(amount, sourceCode, preferredCode, rates, sourceSymbol), + [preferredCode, rates], + ); + + return {children}; +} + +export function useCurrency(): CurrencyContextValue { + return useContext(CurrencyContext); +} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..8a10817b45 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -508,6 +508,8 @@ export const ALL_LOCALE_KEYS = [ "SectionUsageRendering", "SectionTime", "SectionLanguage", + "PreferredCurrencyLabel", + "PreferredCurrencyHelper", "SectionCredentialsSecurity", "SectionDebug", "SectionApiKeys", diff --git a/apps/desktop-tauri/src/lib/currency.test.ts b/apps/desktop-tauri/src/lib/currency.test.ts new file mode 100644 index 0000000000..96a05d5879 --- /dev/null +++ b/apps/desktop-tauri/src/lib/currency.test.ts @@ -0,0 +1,50 @@ +import { describe, expect, it } from "vitest"; +import { + FALLBACK_CURRENCY_RATES, + convertCurrencyAmount, + formatDisplayCurrency, + mergeValidCurrencyRates, + normalizePreferredCurrency, + sumDisplayCurrencyAmounts, +} from "./currency"; + +describe("preferred currency display", () => { + it("normalizes supported preferences and falls back safely for unknown codes", () => { + expect(normalizePreferredCurrency(undefined)).toBe("AUTO"); + expect(normalizePreferredCurrency(" try ")).toBe("TRY"); + expect(normalizePreferredCurrency("BTC")).toBe("AUTO"); + }); + + it("converts both currencies through the USD pivot and rounds for display", () => { + expect(convertCurrencyAmount(10, "USD", "TRY", FALLBACK_CURRENCY_RATES)).toBe(485); + expect(convertCurrencyAmount(10, "GBP", "TRY", FALLBACK_CURRENCY_RATES)).toBeCloseTo(613.92405, 4); + const display = formatDisplayCurrency(10, "USD", "TRY", FALLBACK_CURRENCY_RATES); + expect(display).not.toContain("10.00"); + expect(display).toMatch(/485/); + }); + + it("keeps AUTO, credits, unknown units, and missing-rate values in source units", () => { + expect(formatDisplayCurrency(4.25, "USD", "AUTO", FALLBACK_CURRENCY_RATES)).toMatch(/4\.25/); + expect(formatDisplayCurrency(4.25, "Credits", "TRY", FALLBACK_CURRENCY_RATES)).toBe("4.25 Credits"); + expect(formatDisplayCurrency(10, "USD", "TRY", {} , "$" )).toBe("$10.00"); + expect(convertCurrencyAmount(8, "Quota", "TRY", FALLBACK_CURRENCY_RATES)).toBeNull(); + }); + + it("rejects malformed exchange rates and preserves offline fallbacks", () => { + const rates = mergeValidCurrencyRates({ USD: 1, TRY: Number.NaN, EUR: -2, BTC: 90 }); + expect(rates.TRY).toBe(48.5); + expect(rates.EUR).toBe(0.92); + expect(rates.BTC).toBeUndefined(); + }); + + it("sums only converted overview rows and reports incomplete coverage", () => { + const result = sumDisplayCurrencyAmounts([ + { amount: 10, currency: "USD" }, + { amount: 10, currency: "EUR" }, + { amount: 5, currency: "Credits" }, + ], "TRY", FALLBACK_CURRENCY_RATES); + expect(result.included).toBe(2); + expect(result.considered).toBe(3); + expect(result.total).toBeCloseTo(10 * 48.5 + (10 / 0.92) * 48.5, 8); + }); +}); diff --git a/apps/desktop-tauri/src/lib/currency.ts b/apps/desktop-tauri/src/lib/currency.ts new file mode 100644 index 0000000000..4155b58987 --- /dev/null +++ b/apps/desktop-tauri/src/lib/currency.ts @@ -0,0 +1,117 @@ +export const SUPPORTED_CURRENCIES = [ + "USD", "GBP", "EUR", "CZK", "CNY", "JPY", "KRW", "CAD", "AUD", "HKD", "TWD", "SGD", + "INR", "CHF", "AED", "TRY", +] as const; + +export const FALLBACK_CURRENCY_RATES: Record = { + USD: 1, + GBP: 0.79, + EUR: 0.92, + CZK: 21, + CNY: 7.27, + JPY: 154, + KRW: 1428.9, + CAD: 1.38, + AUD: 1.55, + HKD: 7.8, + TWD: 32.3, + SGD: 1.34, + INR: 84.5, + CHF: 0.8, + AED: 3.6725, + TRY: 48.5, +}; + +export function normalizePreferredCurrency(value: string | null | undefined): string { + const code = value?.trim().toUpperCase() || "AUTO"; + return code === "AUTO" || SUPPORTED_CURRENCIES.includes(code as (typeof SUPPORTED_CURRENCIES)[number]) + ? code + : "AUTO"; +} + +export function convertCurrencyAmount( + amount: number, + sourceCode: string, + targetCode: string, + rates: Record, +): number | null { + if (!Number.isFinite(amount)) return null; + const source = sourceCode.trim().toUpperCase(); + const target = targetCode.trim().toUpperCase(); + if (!SUPPORTED_CURRENCIES.includes(source as (typeof SUPPORTED_CURRENCIES)[number]) || + !SUPPORTED_CURRENCIES.includes(target as (typeof SUPPORTED_CURRENCIES)[number])) return null; + if (source === target) return amount; + const sourceRate = source === "USD" ? 1 : rates[source]; + const targetRate = target === "USD" ? 1 : rates[target]; + if (!Number.isFinite(sourceRate) || sourceRate <= 0 || !Number.isFinite(targetRate) || targetRate <= 0) return null; + const result = (amount / sourceRate) * targetRate; + return Number.isFinite(result) ? result : null; +} + +function formatOriginal(amount: number, code: string, symbol?: string | null): string { + if (symbol) return `${symbol}${amount.toFixed(2)}`; + if (!/^[A-Z]{3}$/.test(code)) return `${amount.toFixed(2)} ${code}`; + try { + return new Intl.NumberFormat("en-US", { style: "currency", currency: code }).format(amount); + } catch { + return `${amount.toFixed(2)} ${code}`; + } +} + +export function formatDisplayCurrency( + amount: number | null | undefined, + sourceCode: string, + preferredCode: string, + rates: Record, + sourceSymbol?: string | null, +): string { + if (amount == null || !Number.isFinite(amount)) return "—"; + const source = sourceCode.trim().toUpperCase(); + const preferred = normalizePreferredCurrency(preferredCode); + if (preferred === "AUTO") return formatOriginal(amount, source, sourceSymbol); + const converted = convertCurrencyAmount(amount, source, preferred, rates); + if (converted == null) return formatOriginal(amount, source, sourceSymbol); + try { + return new Intl.NumberFormat(undefined, { + style: "currency", + currency: preferred, + maximumFractionDigits: 2, + }).format(converted); + } catch { + return `${converted.toFixed(2)} ${preferred}`; + } +} + +export function mergeValidCurrencyRates(input: Record): Record { + const rates = { ...FALLBACK_CURRENCY_RATES }; + for (const code of SUPPORTED_CURRENCIES) { + const value = input[code]; + if (Number.isFinite(value) && value > 0 && (code !== "USD" || Math.abs(value - 1) <= Number.EPSILON)) { + rates[code] = value; + } + } + return rates; +} + +export function sumDisplayCurrencyAmounts( + rows: Array<{ amount: number | null | undefined; currency: string }>, + preferredCode: string, + rates: Record, +): { total: number | null; included: number; considered: number } { + const target = normalizePreferredCurrency(preferredCode); + let total = 0; + let included = 0; + for (const row of rows) { + if (row.amount == null || !Number.isFinite(row.amount)) continue; + let amount: number | null; + if (target === "AUTO") { + amount = row.currency.trim().toUpperCase() === "USD" ? row.amount : null; + } else { + amount = convertCurrencyAmount(row.amount, row.currency || "USD", target, rates); + } + if (amount == null) continue; + total += amount; + included += 1; + } + return { total: included > 0 && Number.isFinite(total) ? total : null, included, considered: rows.length }; +} diff --git a/apps/desktop-tauri/src/lib/tauri.ts b/apps/desktop-tauri/src/lib/tauri.ts index 6ccd63c0de..18b7ad4302 100644 --- a/apps/desktop-tauri/src/lib/tauri.ts +++ b/apps/desktop-tauri/src/lib/tauri.ts @@ -45,6 +45,7 @@ import type { CodexAccountsStateBridge, CodexSwitchResult, DeepSeekPricingStatus, + CurrencyRatesSnapshot, } from "../types/bridge"; export const claudeAccountsList = () => invoke("claude_accounts_list"); @@ -87,6 +88,10 @@ export function getSettingsSnapshot(): Promise { return invoke("get_settings_snapshot"); } +export function getCurrencyRates(preferredCurrencyCode: string): Promise { + return invoke("get_currency_rates", { preferredCurrencyCode }); +} + export function updateSettings( patch: SettingsUpdate, ): Promise { diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index bf693ce7b8..148cf95b34 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from "vitest"; +import { describe, expect, it, vi } from "vitest"; import { formatUsageSpendReportingDay, @@ -104,4 +104,41 @@ describe("usage spend sharing", () => { expect(Object.keys(row).filter((key) => /email|org|token|account/i.test(key))).toEqual([]); expect(unsafeKeys).toEqual(["providerId", "includedInOverview"]); }); + + it("uses the supplied display formatter and currency resolver for PNG cells", () => { + const canvas = document.createElement("canvas"); + const context = { + scale: vi.fn(), fillRect: vi.fn(), strokeRect: vi.fn(), fillText: vi.fn(), + beginPath: vi.fn(), moveTo: vi.fn(), lineTo: vi.fn(), stroke: vi.fn(), + measureText: vi.fn(() => ({ width: 1 })), + } as unknown as CanvasRenderingContext2D; + vi.spyOn(canvas, "getContext").mockReturnValue(context); + vi.spyOn(canvas, "toDataURL").mockReturnValue("data:image/png;base64,test"); + const createElement = document.createElement.bind(document); + vi.spyOn(document, "createElement").mockImplementation((tagName, options) => + tagName === "canvas" ? canvas : createElement(tagName, options), + ); + const formatMetric = vi.fn(() => "₺32.00 · 10 tokens"); + const displayCurrency = vi.fn(() => "TRY"); + const summary: UsageSpendSummary = { + contract: {} as SpendContract, + reportingDay: "2026-09-19", + dashboardTimezone: "UTC", + rows: [{ + providerId: "codex", displayName: "Codex", sevenDay: 1, thirtyDay: 2, + currency: "USD", source: "local", includedInOverview: true, + }], + }; + + try { + renderUsageSpendSharePng(summary, "Usage & Spend", { formatMetric, displayCurrency }); + expect(formatMetric).toHaveBeenCalledWith(1, undefined, "USD", "tokens"); + expect(formatMetric).toHaveBeenCalledWith(2, undefined, "USD", "tokens"); + expect(displayCurrency).toHaveBeenCalledWith("USD"); + expect(context.fillText).toHaveBeenCalledWith("₺32.00 · 10 tokens", expect.any(Number), expect.any(Number)); + expect(context.fillText).toHaveBeenCalledWith("TRY", expect.any(Number), expect.any(Number)); + } finally { + vi.restoreAllMocks(); + } + }); }); diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 2246e98505..3da6473f4a 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -139,6 +139,11 @@ export function formatSpendMetric( return parts.length > 0 ? parts.join(" · ") : "—"; } +export interface UsageSpendSharePresentation { + formatMetric?: (cost: number | null | undefined, tokens: number | null | undefined, sourceCurrency: string, tokenLabel: string) => string; + displayCurrency?: (sourceCurrency: string) => string; +} + /** * Render the sanitized share-card PNG. * @@ -148,7 +153,11 @@ export function formatSpendMetric( * states the guarantee. Keep it that way — do not add account fields to the * drawn cells or the footer. */ -export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: string): string { +export function renderUsageSpendSharePng( + summary: UsageSpendSummary, + title: string, + presentation: UsageSpendSharePresentation = {}, +): string { const rows = summary.rows; const pad = 24; const rowH = 28; @@ -202,9 +211,11 @@ 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"), - row.currency || "USD", + presentation.formatMetric?.(row.sevenDay, row.sevenDayTokens, row.currency || "USD", "tokens") + ?? formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, "tokens"), + presentation.formatMetric?.(row.thirtyDay, row.thirtyDayTokens, row.currency || "USD", "tokens") + ?? formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, "tokens"), + presentation.displayCurrency?.(row.currency || "USD") ?? (row.currency || "USD"), row.source, ]; let cellX = pad; @@ -250,10 +261,11 @@ export function shareUsageSpendPng( summary: UsageSpendSummary | null, title: string, filename: string, + presentation: UsageSpendSharePresentation = {}, ): string | null { if (!summary) return "UsageSpendShareEmpty"; try { - const dataUrl = renderUsageSpendSharePng(summary, title); + const dataUrl = renderUsageSpendSharePng(summary, title, presentation); if (!dataUrl) return "UsageSpendShareFailed"; downloadPng(dataUrl, filename); return null; diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..b068920c2b 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -2,6 +2,8 @@ import { Fragment, useEffect, useState, type CSSProperties } from "react"; import { getCurrentWindow } from "@tauri-apps/api/window"; import type { BootstrapState, ProviderUsageSnapshot, UsageSpendSummary } from "../types/bridge"; import type { LocaleKey } from "../i18n/keys"; +import { useCurrency } from "../hooks/CurrencyProvider"; +import { normalizePreferredCurrency, sumDisplayCurrencyAmounts } from "../lib/currency"; import { beginFlyoutGesture, getUsageSpendSummary, @@ -337,6 +339,7 @@ function TrayResizeHandles() { } function OverviewSpendSummary({ providerIds, t }: { providerIds: string[]; t: (key: LocaleKey) => string }) { + const { preferredCode, rates, format } = useCurrency(); const [summary, setSummary] = useState(null); const [shareError, setShareError] = useState(null); @@ -364,21 +367,25 @@ function OverviewSpendSummary({ providerIds, t }: { providerIds: string[]; t: (k }; const rows = overviewSummary.rows; - const summable = rows.filter((row) => (row.currency || "USD") === "USD"); - const known = summable.filter((row) => row.thirtyDay != null && Number.isFinite(row.thirtyDay)); - if (known.length === 0) return null; - const total = known.reduce((sum, row) => sum + (row.thirtyDay ?? 0), 0); - const partial = known.length < rows.length; - const formatter = new Intl.NumberFormat(undefined, { style: "currency", currency: "USD", maximumFractionDigits: 2 }); + const target = normalizePreferredCurrency(preferredCode); + const aggregate = sumDisplayCurrencyAmounts( + rows.map((row) => ({ amount: row.thirtyDay, currency: row.currency || "USD" })), + target, + rates, + ); + const partial = aggregate.included < aggregate.considered; + const displayedTotal = aggregate.total == null + ? "—" + : `${partial ? "~" : ""}${format(aggregate.total, target === "AUTO" ? "USD" : target)}`; return (
{t("OverviewSpendTitle")} - {partial ? "~" : ""}{formatter.format(total)} + {displayedTotal}
- {known.length} of {rows.length} {t("OverviewSpendProviderCoverage")} · {t("OverviewSpendEstimate")} + {aggregate.included} of {aggregate.considered} {t("OverviewSpendProviderCoverage")} · {t("OverviewSpendEstimate")}
+ {!compactOverview && detailText && ( +
{detailText}
+ )} {!compactOverview && !isInformational && snap.isExhausted && (
{exhaustedLabel}
)} diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index 88860e087d..f24e1fda7d 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -47,6 +47,7 @@ type RateWindowOptions = { informational?: boolean; resetsAt?: string | null; resetDescription?: string | null; + descriptionIsDetail?: boolean; }; function rateWindow( @@ -61,6 +62,7 @@ function rateWindow( resetDescription: opts.resetDescription ?? null, isExhausted: opts.exhausted ?? false, isInformational: opts.informational, + descriptionIsDetail: opts.descriptionIsDetail, reservePercent: null, reserveDescription: null, }; @@ -76,6 +78,7 @@ function snapshot( errorState?: ProviderUsageSnapshot["errorState"]; resetsAt?: string | null; resetDescription?: string | null; + descriptionIsDetail?: boolean; informational?: boolean; secondary?: { used: number; @@ -325,6 +328,25 @@ describe("FloatBar", () => { }); }); + it("never shows a detail-backed description as reset text", async () => { + tauriMocks.getCachedProviders.mockResolvedValue([ + snapshot("claude", "Claude", 13, { + resetDescription: "34.07 EUR / 255.00 EUR · 220.93 EUR remaining", + descriptionIsDetail: true, + }), + ]); + tauriMocks.getSettingsSnapshot.mockResolvedValue( + settings({ floatBarShowResetInline: true }), + ); + + const { container } = renderFloatBar(bootstrap({ floatBarShowResetInline: true })); + await waitFor(() => { + const pill = container.querySelector(".floatbar__pill"); + expect(pill?.getAttribute("title")).toBe("Claude: 13% used"); + expect(container.querySelector(".floatbar__reset")).toBeNull(); + }); + }); + it("keeps an informational primary window when no secondary window is available", async () => { tauriMocks.getCachedProviders.mockResolvedValue([ snapshot("claude", "Claude", 10, { informational: true }), diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.tsx index e8d71af8fe..c01f3501c3 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.tsx @@ -20,6 +20,7 @@ import { import { ProviderIcon } from "../components/providers/ProviderIcon"; import { getProviderIcon } from "../components/providers/providerIcons"; import { describeProviderState } from "../lib/providerState"; +import { resetDescriptionFallback } from "../lib/usageWindows"; import type { BootstrapState, ProviderLocalUsageSummary, @@ -221,7 +222,7 @@ function ProviderPill({ const label = state.isProblem ? stateLabel : `${Math.round(displayPercent)}%`; const resetText = useFormattedResetTime( rateWindow.resetsAt, - rateWindow.resetDescription, + resetDescriptionFallback(rateWindow), resetRelative, ); const resetSuffix = resetText ? `\n${resetText}` : ""; diff --git a/apps/desktop-tauri/src/lib/usageWindows.test.ts b/apps/desktop-tauri/src/lib/usageWindows.test.ts index 2876b33810..e5183d7153 100644 --- a/apps/desktop-tauri/src/lib/usageWindows.test.ts +++ b/apps/desktop-tauri/src/lib/usageWindows.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it } from "vitest"; import type { RateWindowSnapshot } from "../types/bridge"; -import { selectSingleMetricUsageWindow } from "./usageWindows"; +import { + resetDescriptionFallback, + selectSingleMetricUsageWindow, + windowDetailText, +} from "./usageWindows"; function rateWindow( overrides: Partial = {}, @@ -46,3 +50,32 @@ describe("selectSingleMetricUsageWindow", () => { expect(selectSingleMetricUsageWindow({ primary, secondary })).toBe(primary); }); }); + +describe("detail-backed descriptions", () => { + const amounts = "34.07 EUR / 255.00 EUR · 220.93 EUR remaining"; + + it("keeps ordinary descriptions as the reset fallback and shows no detail line", () => { + const window = rateWindow({ resetDescription: "Resets in 2h" }); + + expect(resetDescriptionFallback(window)).toBe("Resets in 2h"); + expect(windowDetailText(window)).toBeNull(); + }); + + it("moves a detail-backed description out of the reset fallback", () => { + const window = rateWindow({ resetDescription: ` ${amounts} `, descriptionIsDetail: true }); + + expect(resetDescriptionFallback(window)).toBeNull(); + expect(windowDetailText(window)).toBe(amounts); + }); + + it("ignores blank and informational detail-backed descriptions", () => { + expect( + windowDetailText(rateWindow({ resetDescription: " ", descriptionIsDetail: true })), + ).toBeNull(); + expect( + windowDetailText( + rateWindow({ resetDescription: amounts, descriptionIsDetail: true, isInformational: true }), + ), + ).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/lib/usageWindows.ts b/apps/desktop-tauri/src/lib/usageWindows.ts index 9a55daf098..9179e1dce3 100644 --- a/apps/desktop-tauri/src/lib/usageWindows.ts +++ b/apps/desktop-tauri/src/lib/usageWindows.ts @@ -12,3 +12,22 @@ export function selectSingleMetricUsageWindow( ? secondary : primary; } + +/** + * Reset-wording fallback for a window. A detail-backed description (for + * example spend amounts) is never reset text, so it must not reach + * `normalizeResetDescription` through the reset formatter. + */ +export function resetDescriptionFallback( + window: Pick, +): string | null { + return window.descriptionIsDetail ? null : window.resetDescription; +} + +/** Secondary detail line for a detail-backed window, or null. */ +export function windowDetailText( + window: Pick, +): string | null { + if (!window.descriptionIsDetail || window.isInformational) return null; + return window.resetDescription?.trim() || null; +} diff --git a/apps/desktop-tauri/src/styles.css b/apps/desktop-tauri/src/styles.css index a41405732c..b01bee4339 100644 --- a/apps/desktop-tauri/src/styles.css +++ b/apps/desktop-tauri/src/styles.css @@ -2100,7 +2100,8 @@ body:has(.tray-panel-reveal) { transition: width 0.3s; } -.provider-usage-bar__reset { +.provider-usage-bar__reset, +.provider-usage-bar__detail { font-size: 0.68rem; color: var(--text-muted); } @@ -3779,7 +3780,8 @@ html:has(.menu-surface--tray) { max-width: 100%; } .menu-surface--tray .menu-metric__pct, -.menu-surface--tray .menu-metric__reset { +.menu-surface--tray .menu-metric__reset, +.menu-surface--tray .menu-metric__detail { font-size: 10px; } .menu-surface--tray .menu-metric__reserve { @@ -4305,6 +4307,15 @@ html:has(.menu-surface--tray) { white-space: nowrap; } +/* Detail-backed window description (for example spend amounts): its own + secondary line under the metric row, never reset wording. */ +.menu-metric__detail { + font-size: 11px; + color: var(--text-secondary); + min-width: 0; + overflow-wrap: anywhere; +} + .menu-metric__reset { font-size: 11px; color: var(--text-secondary); @@ -5051,7 +5062,8 @@ html:has(.menu-surface--tray) { } .menu-metric__pct, -.menu-metric__reset { +.menu-metric__reset, +.menu-metric__detail { font-size: var(--font-caption); } .menu-card__cost-line, diff --git a/apps/desktop-tauri/src/surfaces/menu-surface-layout.css b/apps/desktop-tauri/src/surfaces/menu-surface-layout.css index e79ac3293a..7e57ed3dea 100644 --- a/apps/desktop-tauri/src/surfaces/menu-surface-layout.css +++ b/apps/desktop-tauri/src/surfaces/menu-surface-layout.css @@ -61,7 +61,8 @@ } :is(.menu-surface--tray, .menu-surface--popout) .menu-metric__pct, -:is(.menu-surface--tray, .menu-surface--popout) .menu-metric__reset { +:is(.menu-surface--tray, .menu-surface--popout) .menu-metric__reset, +:is(.menu-surface--tray, .menu-surface--popout) .menu-metric__detail { font-size: 11px; line-height: 1.4; } diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx index cf7ca9b884..734d47c414 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx @@ -86,6 +86,51 @@ describe("UsageSection", () => { expect(screen.getByText("42%")).toBeInTheDocument(); }); + it.each([false, true])( + "renders detail-backed amounts as a detail line, not reset text (resetsAt: %s)", + async (hasReset) => { + const detail = provider(); + const resetsAt = hasReset + ? new Date(Date.now() + 3 * 60 * 60 * 1000 + 30_000).toISOString() + : null; + detail.session = { + ...rateWindow(75), + resetsAt, + resetDescription: "19.17 EUR / 25.50 EUR · 6.33 EUR remaining", + descriptionIsDetail: true, + }; + detail.extraRateWindows = [ + { + id: "mistral-monthly-plan", + title: "Monthly Plan", + window: { + ...rateWindow(13), + resetsAt, + resetDescription: "34.07 EUR / 255.00 EUR · 220.93 EUR remaining", + descriptionIsDetail: true, + }, + }, + ]; + + const { container } = render( + + key} /> + , + ); + + const detailLines = await screen.findAllByText(/EUR remaining/); + expect(detailLines.map((line) => line.textContent)).toEqual([ + "19.17 EUR / 25.50 EUR · 6.33 EUR remaining", + "34.07 EUR / 255.00 EUR · 220.93 EUR remaining", + ]); + detailLines.forEach((line) => expect(line).toHaveClass("provider-usage-bar__detail")); + expect(screen.queryByText(/Resets .*EUR/)).not.toBeInTheDocument(); + expect(container.querySelectorAll(".provider-usage-bar__reset")).toHaveLength( + hasReset ? 2 : 0, + ); + }, + ); + it("filters only hidden metric and extra rows", async () => { const detail = provider(); detail.weekly = rateWindow(30); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index b2d24dbdfc..b5babf826c 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -9,6 +9,7 @@ import { ProviderDisplayRow } from "../../../../components/ProviderDisplayRow"; import type { LocaleKey } from "../../../../i18n/keys"; import { useFormattedResetTime } from "../../../../hooks/useFormattedResetTime"; import { isUsageItemVisible } from "../../../../lib/usageItemVisibility"; +import { resetDescriptionFallback, windowDetailText } from "../../../../lib/usageWindows"; interface Props { provider: ProviderDetail; @@ -121,9 +122,10 @@ function UsageBar({ const usedPct = Number.isFinite(rate.usedPercent) ? Math.max(0, rate.usedPercent) : 0; const pct = Math.min(100, usedPct); const isInformational = rate.isInformational === true; + const detailText = windowDetailText(rate); const formattedReset = useFormattedResetTime( rate.resetsAt, - rate.resetDescription, + resetDescriptionFallback(rate), resetTimeRelative, ); const resetHint = formattedReset @@ -161,6 +163,7 @@ function UsageBar({ {!isInformational && resetHint && ( {resetHint} )} + {detailText && {detailText}} ); } \ No newline at end of file diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 88754baa7d..53d8c843ea 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -552,6 +552,8 @@ export interface RateWindowSnapshot { resetDescription: string | null; isExhausted: boolean; isInformational?: boolean; + /** `resetDescription` is a detail line (for example spend amounts), not reset wording. */ + descriptionIsDetail?: boolean; reservePercent: number | null; reserveDescription: string | null; reserveWillLastToReset?: boolean; diff --git a/rust/src/cli/usage/render.rs b/rust/src/cli/usage/render.rs index 62fd50310d..95f20b9d77 100644 --- a/rust/src/cli/usage/render.rs +++ b/rust/src/cli/usage/render.rs @@ -325,8 +325,31 @@ fn append_window_line(lines: &mut Vec, label: &str, window: &RateWindow, } let bar = render_progress_bar(window.used_percent, 20, use_color); - let reset = window - .format_countdown() + let countdown = window.format_countdown(); + if window.description_is_detail { + // The description is a detail line (for example spend amounts), so + // the reset stays on its own line and only exists with a real date. + lines.push(format!( + " {:<8} {} {} used", + format!("{}:", label), + bar, + format_percent(window.used_percent), + )); + if let Some(countdown) = countdown { + lines.push(format!(" resets in {countdown}")); + } + if let Some(detail) = window + .reset_description + .as_deref() + .map(str::trim) + .filter(|detail| !detail.is_empty()) + { + lines.push(format!(" {detail}")); + } + return; + } + + let reset = countdown .map(|c| format!(" (resets in {})", c)) .unwrap_or_default(); lines.push(format!( diff --git a/rust/src/cli/usage_tests.rs b/rust/src/cli/usage_tests.rs index d332131b81..dcba8fad01 100644 --- a/rust/src/cli/usage_tests.rs +++ b/rust/src/cli/usage_tests.rs @@ -2,8 +2,8 @@ use super::*; use crate::core::{ - CostSnapshot, ProviderAccountData, ProviderDisplayDetail, ProviderInventoryItem, RateWindow, - TokenAccount, TokenAccountSupport, UsageSnapshot, + CostSnapshot, NamedRateWindow, ProviderAccountData, ProviderDisplayDetail, + ProviderInventoryItem, RateWindow, TokenAccount, TokenAccountSupport, UsageSnapshot, }; use crate::providers::claude::claude_swap::ClaudeSwapAccount; use crate::status::{ProviderStatus as StatusInfo, StatusLevel}; @@ -324,3 +324,91 @@ fn json_inventory_is_additive_and_contains_no_redemption_token() { .contains("coupon-token-secret") ); } + +fn detail_window(used: f64, detail: &str, resets_at: Option>) -> RateWindow { + RateWindow::with_details(used, None, resets_at, Some(detail.to_string())) + .with_description_as_detail() +} + +fn detail_backed_result(resets_at: Option>) -> ProviderFetchResult { + let mut usage = UsageSnapshot::new(detail_window( + 75.0, + "19.17 EUR / 25.50 EUR · 6.33 EUR remaining", + resets_at, + )) + .with_primary_label("Included API"); + usage.extra_rate_windows.push(NamedRateWindow::new( + "mistral-monthly-plan", + "Monthly Plan", + detail_window( + 13.0, + "34.07 EUR / 255.00 EUR · 220.93 EUR remaining", + resets_at, + ), + )); + fetch_result(usage) +} + +#[test] +fn detail_backed_windows_print_reset_then_amounts_lines() { + let resets_at = Utc::now() + chrono::Duration::minutes(61); + let output = render_text( + ProviderId::Mistral, + &detail_backed_result(Some(resets_at)), + false, + ); + let lines: Vec<&str> = output.lines().collect(); + + let primary = lines + .iter() + .position(|line| line.starts_with(" Included API:")) + .expect("primary line"); + assert!(lines[primary].ends_with("75% used")); + assert!(lines[primary + 1].starts_with(" resets in ")); + assert_eq!( + lines[primary + 2], + " 19.17 EUR / 25.50 EUR · 6.33 EUR remaining" + ); + let plan = lines + .iter() + .position(|line| line.starts_with(" Monthly Plan:")) + .expect("plan line"); + assert!(lines[plan].ends_with("13% used")); + assert!(lines[plan + 1].starts_with(" resets in ")); + assert_eq!( + lines[plan + 2], + " 34.07 EUR / 255.00 EUR · 220.93 EUR remaining" + ); + assert!(!output.contains("(resets in")); +} + +#[test] +fn detail_backed_windows_omit_reset_line_without_reset_date() { + let output = render_text(ProviderId::Mistral, &detail_backed_result(None), false); + let lines: Vec<&str> = output.lines().collect(); + + assert!(!output.contains("resets in")); + let plan = lines + .iter() + .position(|line| line.starts_with(" Monthly Plan:")) + .expect("plan line"); + assert_eq!( + lines[plan + 1], + " 34.07 EUR / 255.00 EUR · 220.93 EUR remaining" + ); +} + +#[test] +fn detail_backed_flag_stays_out_of_json_output() { + let json = render_json_result(ProviderId::Mistral, detail_backed_result(None), None); + let windows = json["usage"]["extra_rate_windows"] + .as_array() + .expect("extra windows"); + assert_eq!(windows[0]["id"], "mistral-monthly-plan"); + assert_eq!( + windows[0]["window"]["reset_description"], + "34.07 EUR / 255.00 EUR · 220.93 EUR remaining" + ); + assert!(windows[0]["window"].get("descriptionIsDetail").is_none()); + assert!(windows[0]["window"].get("description_is_detail").is_none()); +} diff --git a/rust/src/core/rate_window.rs b/rust/src/core/rate_window.rs index 1e943df983..cf250ef73d 100755 --- a/rust/src/core/rate_window.rs +++ b/rust/src/core/rate_window.rs @@ -92,6 +92,13 @@ pub struct RateWindow { #[serde(default)] pub is_informational: bool, + /// Whether `reset_description` is a detail line (for example spend + /// amounts) rather than reset wording. Renderers show it as its own + /// secondary line and never as a reset time. Provider-owned; kept out of + /// serialized output so CLI JSON is unchanged. + #[serde(default, skip_serializing)] + pub description_is_detail: bool, + /// Whether the provider explicitly supplied the usage percentage. /// Internal-only metadata prevents a missing value normalized to zero from /// becoming an exported quota measurement. @@ -112,6 +119,7 @@ impl RateWindow { resets_at: None, reset_description: None, is_informational: false, + description_is_detail: false, usage_known: true, } } @@ -152,10 +160,17 @@ impl RateWindow { resets_at, reset_description, is_informational: false, + description_is_detail: false, usage_known: true, } } + /// Mark `reset_description` as a detail line instead of reset wording. + pub fn with_description_as_detail(mut self) -> Self { + self.description_is_detail = true; + self + } + pub(crate) fn with_usage_known(mut self, usage_known: bool) -> Self { self.usage_known = usage_known; self diff --git a/rust/src/providers/copilot/api.rs b/rust/src/providers/copilot/api.rs index c1f8793c21..9b97b0d7b5 100755 --- a/rust/src/providers/copilot/api.rs +++ b/rust/src/providers/copilot/api.rs @@ -496,6 +496,7 @@ impl UsableQuota { resets_at: reset, reset_description, is_informational: false, + description_is_detail: false, usage_known: true, } } diff --git a/rust/src/providers/mistral/mod.rs b/rust/src/providers/mistral/mod.rs index bcaa15e48c..33f82bc908 100644 --- a/rust/src/providers/mistral/mod.rs +++ b/rust/src/providers/mistral/mod.rs @@ -390,7 +390,9 @@ impl MistralProvider { "No usage this month".to_string() }; - let primary = RateWindow::with_details(0.0, None, reset_date, Some(cost_description)); + // The description carries spend amounts, not reset wording. + let primary = RateWindow::with_details(0.0, None, reset_date, Some(cost_description)) + .with_description_as_detail(); let mut usage = UsageSnapshot::new(primary); if summary.model_count > 0 { usage = usage.with_login_method(format!("{} model(s)", summary.model_count)); @@ -443,6 +445,7 @@ impl MistralProvider { budget.resets_at, Some(description), ) + .with_description_as_detail() } fn build_price_index(prices: Vec) -> HashMap { @@ -662,6 +665,19 @@ mod tests { "mistral-monthly-plan" ); assert_eq!(result.cost.as_ref().map(|cost| cost.used), Some(12.5)); + assert!(result.usage.primary.description_is_detail); + assert!( + result.usage.extra_rate_windows[0] + .window + .description_is_detail + ); + assert_eq!( + result.usage.extra_rate_windows[0] + .window + .reset_description + .as_deref(), + Some("10.00 EUR / 20.00 EUR · 10.00 EUR remaining") + ); } #[test] From a1980aeb40b24f3035754cec06a5551a526e54e6 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:48:49 +0700 Subject: [PATCH 090/259] Port upstream 0.67.0: Add Aixy provider (key budgets and 7-day usage) --- README.md | 1 + .../src/commands/provider_settings.rs | 40 +- .../src-tauri/src/commands/providers.rs | 7 +- .../providers/icons/ProviderIcon-aixy.svg | 1 + .../src/components/providers/providerIcons.ts | 3 + apps/desktop-tauri/src/i18n/keys.ts | 5 + apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 2 +- .../settings/providers/ProviderDetailPane.tsx | 14 +- .../sections/MenuBarMetricSection.test.tsx | 19 + .../sections/MenuBarMetricSection.tsx | 35 +- .../sections/WayfinderGatewaySection.test.tsx | 53 ++ .../sections/WayfinderGatewaySection.tsx | 43 +- .../providers/sections/usageSourcePolicy.ts | 6 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + docs/PROVIDERS.md | 24 + rust/src/core/provider.rs | 10 +- rust/src/core/provider_factory.rs | 34 +- rust/src/core/token_accounts.rs | 24 + rust/src/locale.rs | 5 + rust/src/locale/en-US.ftl | 5 + .../providers/aixy/fixtures/key-usage.json | 134 +++++ rust/src/providers/aixy/mod.rs | 215 ++++++++ rust/src/providers/aixy/model.rs | 411 ++++++++++++++ rust/src/providers/aixy/present.rs | 173 ++++++ rust/src/providers/aixy/tests.rs | 504 ++++++++++++++++++ rust/src/providers/bifrost/mod.rs | 16 +- rust/src/providers/mod.rs | 22 + rust/src/settings/api_keys.rs | 11 + 28 files changed, 1760 insertions(+), 58 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-aixy.svg create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.test.tsx create mode 100644 rust/src/providers/aixy/fixtures/key-usage.json create mode 100644 rust/src/providers/aixy/mod.rs create mode 100644 rust/src/providers/aixy/model.rs create mode 100644 rust/src/providers/aixy/present.rs create mode 100644 rust/src/providers/aixy/tests.rs diff --git a/README.md b/README.md index 2e052ad2f9..bd9e0de719 100755 --- a/README.md +++ b/README.md @@ -116,6 +116,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Grok | Cookies / auth.json | Billing | | Helmcode (also NaN Builders) | Browser cookies / manual Cookie header | Per-model token quotas, reset windows, Helmcode prepaid balance | | Replicate | Cookies / token accounts | Monthly spend, credit balance | +| Aixy | API Key / token accounts | Applicable budget balances, 7-day key usage | | ElevenLabs | API Key | Subscription Credits, Voice Slots | | Deepgram | API Key | Project Usage | | Groq | API Key | Enterprise Metrics | 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 3aa5518b59..314baa56d8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -400,15 +400,48 @@ mod tests { ); } + #[test] + fn fetch_context_carries_saved_gateway_urls_for_every_gateway_provider() { + use codexbar::settings::{ApiKeys, ManualCookies, Settings}; + use std::collections::HashMap; + + let mut settings = Settings::default(); + for (id, url) in [ + (ProviderId::Wayfinder, "http://localhost:8787"), + (ProviderId::Bifrost, "https://bifrost.example.com"), + (ProviderId::Aixy, "https://aixy.example.com/prefix"), + ] { + settings.set_gateway_url(id, url); + let ctx = super::super::providers::build_fetch_context( + id, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + assert_eq!(ctx.gateway_url.as_deref(), Some(url), "{id:?}"); + } + + let ctx = super::super::providers::build_fetch_context( + ProviderId::Codex, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + assert_eq!(ctx.gateway_url, None); + } + #[test] fn maps_gitkraken_organization_provider() { assert_eq!(workspace_provider("gitkraken"), Some(ProviderId::GitKraken)); } #[test] - fn gateway_provider_exposes_wayfinder_and_bifrost_only() { + fn gateway_provider_exposes_gateway_providers_only() { assert_eq!(gateway_provider("wayfinder"), Some(ProviderId::Wayfinder)); assert_eq!(gateway_provider("bifrost"), Some(ProviderId::Bifrost)); + assert_eq!(gateway_provider("aixy"), Some(ProviderId::Aixy)); assert_eq!(gateway_provider("codex"), None); } @@ -466,6 +499,7 @@ fn gateway_provider(provider_id: &str) -> Option { match provider_id { "wayfinder" => Some(codexbar::core::ProviderId::Wayfinder), "bifrost" => Some(codexbar::core::ProviderId::Bifrost), + "aixy" => Some(codexbar::core::ProviderId::Aixy), _ => None, } } @@ -491,6 +525,10 @@ pub fn set_provider_gateway_url(provider_id: String, gateway_url: String) -> Res codexbar::providers::bifrost::validate_gateway_url(gateway_url) .map_err(|error| error.to_string())?; } + codexbar::core::ProviderId::Aixy => { + codexbar::providers::aixy::validate_gateway_url(gateway_url) + .map_err(|error| error.to_string())?; + } _ => unreachable!("gateway_provider only returns gateway providers"), } diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 779f9c801c..6ad574d3ca 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -214,8 +214,11 @@ pub(crate) fn build_fetch_context( let workspace_id = settings.workspace_id(id).trim().to_string(); let api_region = settings.api_region(id).trim().to_string(); - let gateway_url = (id == ProviderId::Wayfinder && !settings.gateway_url(id).is_empty()) - .then(|| settings.gateway_url(id).to_string()); + // Every gateway-style provider (Wayfinder, Bifrost, Aixy) stores its base + // URL here; providers without one report an empty string. + let gateway_url = Some(settings.gateway_url(id)) + .filter(|url| !url.is_empty()) + .map(str::to_owned); // Local-first Auto providers (OpenCode Go) flip to web-first when a // token account or manual cookie source scopes the session to web creds. let auto_prefer_web = token_override.is_some() || cookie_source == "manual"; diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-aixy.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-aixy.svg new file mode 100644 index 0000000000..224bb877f1 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-aixy.svg @@ -0,0 +1 @@ +Aixy diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index ee24a2c1f1..99ef91ee3d 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -20,6 +20,7 @@ import deepgram from "./icons/ProviderIcon-deepgram.svg?raw"; import deepinfra from "./icons/ProviderIcon-deepinfra.svg?raw"; import fireworks from "./icons/ProviderIcon-fireworks.svg?raw"; import aiand from "./icons/ProviderIcon-aiand.svg?raw"; +import aixy from "./icons/ProviderIcon-aixy.svg?raw"; import clinepass from "./icons/ProviderIcon-clinepass.svg?raw"; import longcat from "./icons/ProviderIcon-longcat.svg?raw"; import neuralwatt from "./icons/ProviderIcon-neuralwatt.svg?raw"; @@ -107,6 +108,7 @@ const RAW: Record = { deepinfra: tint(deepinfra), fireworks: tint(fireworks), aiand: tint(aiand), + aixy: tint(aixy), clinepass: tint(clinepass), longcat: tint(longcat), neuralwatt: tint(neuralwatt), @@ -186,6 +188,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { grok: { id: "grok", brandColor: "#111827", fallbackLetter: "G", svgPath: RAW.grok }, groq: { id: "groq", brandColor: "#f55036", fallbackLetter: "G", svgPath: RAW.groq }, bifrost: { id: "bifrost", brandColor: "#5b7cfa", fallbackLetter: "B" }, + aixy: { id: "aixy", brandColor: "#123650", fallbackLetter: "A", svgPath: RAW.aixy }, gitkraken: { id: "gitkraken", brandColor: "#179287", fallbackLetter: "G" }, huggingface: { id: "huggingface", brandColor: "#ffd21e", fallbackLetter: "H", svgPath: RAW.huggingface }, hyper: { id: "hyper", brandColor: "#7c3aed", fallbackLetter: "H" }, diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 1c07040562..904dd541f8 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -96,6 +96,11 @@ export const ALL_LOCALE_KEYS = [ "WayfinderGatewayLabel", "WayfinderGatewayHelp", "WayfinderGatewayStatus", + "BifrostGatewayTitle", + "BifrostGatewayHelp", + "AixyGatewayTitle", + "AixyGatewayLabel", + "AixyGatewayHelp", "WayfinderModels", "WayfinderRequests", "WayfinderTokens", diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 09ead3c1d3..8e9c989a21 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", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", + "aiand", "aixy", "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/providers/ProviderDetailPane.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx index ae3e6766a5..5d8099a174 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx @@ -47,7 +47,10 @@ import { AccentColorSection } from "./sections/AccentColorSection"; import { ProviderIssueNotice } from "./sections/ProviderIssueNotice"; import { CredentialStorageSection } from "./sections/CredentialStorageSection"; import { CredentialsDispatcher } from "./sections/CredentialsDispatcher"; -import { WayfinderGatewaySection } from "./sections/WayfinderGatewaySection"; +import { + isGatewayProviderId, + WayfinderGatewaySection, +} from "./sections/WayfinderGatewaySection"; import { AzureApiVersionSection } from "./sections/AzureApiVersionSection"; interface Props { @@ -145,9 +148,8 @@ export function ProviderDetailPane({ } }, []); - const gatewayProviderId = providerId === "wayfinder" || providerId === "bifrost" - ? providerId - : null; + const gatewayProviderId = + providerId !== null && isGatewayProviderId(providerId) ? providerId : null; useEffect(() => { setGatewayLoadedProviderId(null); @@ -345,9 +347,10 @@ export function ProviderDetailPane({ t={t} onChanged={reload} /> - {(detail.id === "wayfinder" || detail.id === "bifrost") && + {isGatewayProviderId(detail.id) && gatewayLoadedProviderId === detail.id && ( void saveGateway()} t={t} - bifrost={detail.id === "bifrost"} /> )} { providerMetrics: { copilot: "extraUsage" }, }); }); + + it("offers only Automatic for Aixy, even with extra budget windows", () => { + const aixy = provider(); + aixy.id = "aixy"; + aixy.displayName = "Aixy"; + aixy.weekly = rateWindow(40); + render( + key} + onChange={vi.fn()} + />, + ); + + const options = screen.getAllByRole("option").map((option) => option.textContent); + expect(options).toEqual(["Automatic"]); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx index 7bc89e0d7d..f8ff1eb89f 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx @@ -20,6 +20,8 @@ interface MetricOption { label: string; } +const AUTOMATIC_ONLY_PROVIDERS: ReadonlySet = new Set(["aixy"]); + export function MenuBarMetricSection({ provider, providerMetrics, @@ -76,11 +78,15 @@ function metricOptions( selected: MetricPreference, t: (key: LocaleKey) => string, ): MetricOption[] { - const options: MetricOption[] = [ - { value: "automatic", label: t("Automatic") }, - { value: "session", label: t("ProviderSessionLabel") }, - ]; + const options: MetricOption[] = [{ value: "automatic", label: t("Automatic") }]; + + // Aixy's primary budget depends on which limits currently apply to the key, + // so a fixed session/weekly lane would be misleading. Offer Automatic only. + if (AUTOMATIC_ONLY_PROVIDERS.has(provider.id)) { + return withSelected(options, selected, t); + } + options.push({ value: "session", label: t("ProviderSessionLabel") }); if (provider.weekly) { options.push({ value: "weekly", label: t("ProviderWeeklyLabel") }); } @@ -104,12 +110,21 @@ function metricOptions( if (provider.id === "gemini" && provider.weekly) { options.push({ value: "average", label: t("Average") }); } - if (!options.some((option) => option.value === selected)) { - options.push({ + return withSelected(options, selected, t); +} + +/** Keep a previously saved preference visible even if it is no longer offered. */ +function withSelected( + options: MetricOption[], + selected: MetricPreference, + t: (key: LocaleKey) => string, +): MetricOption[] { + if (options.some((option) => option.value === selected)) return options; + return [ + ...options, + { value: selected, label: selected === "credits" ? t("CreditsLabel") : selected, - }); - } - - return options; + }, + ]; } diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.test.tsx new file mode 100644 index 0000000000..1b95267446 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.test.tsx @@ -0,0 +1,53 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import { + isGatewayProviderId, + WayfinderGatewaySection, + type GatewayProviderId, +} from "./WayfinderGatewaySection"; + +function renderSection(providerId: GatewayProviderId, onSave = vi.fn()) { + render( + key} + />, + ); + return onSave; +} + +describe("WayfinderGatewaySection", () => { + it.each([ + ["wayfinder", "WayfinderGatewayTitle", "WayfinderGatewayLabel", "WayfinderGatewayHelp"], + ["bifrost", "BifrostGatewayTitle", "WayfinderGatewayLabel", "BifrostGatewayHelp"], + ["aixy", "AixyGatewayTitle", "AixyGatewayLabel", "AixyGatewayHelp"], + ] as const)("uses localized %s copy", (providerId, title, label, help) => { + renderSection(providerId); + + expect(screen.getByRole("heading", { name: title })).toBeInTheDocument(); + expect(screen.getByLabelText(label)).toHaveValue("https://gateway.example.com"); + expect(screen.getByText(help)).toBeInTheDocument(); + }); + + it("saves through the shared button", () => { + const onSave = renderSection("aixy"); + + fireEvent.click(screen.getByRole("button", { name: "Save" })); + + expect(onSave).toHaveBeenCalledTimes(1); + }); + + it("recognizes only gateway providers", () => { + expect(isGatewayProviderId("aixy")).toBe(true); + expect(isGatewayProviderId("bifrost")).toBe(true); + expect(isGatewayProviderId("wayfinder")).toBe(true); + expect(isGatewayProviderId("codex")).toBe(false); + expect(isGatewayProviderId("toString")).toBe(false); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.tsx index 77b4067771..9c454f0ae1 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/WayfinderGatewaySection.tsx @@ -1,6 +1,37 @@ import type { LocaleKey } from "../../../../i18n/keys"; +export type GatewayProviderId = "wayfinder" | "bifrost" | "aixy"; + +interface GatewayCopy { + title: LocaleKey; + label: LocaleKey; + help: LocaleKey; +} + +const GATEWAY_COPY: Record = { + wayfinder: { + title: "WayfinderGatewayTitle", + label: "WayfinderGatewayLabel", + help: "WayfinderGatewayHelp", + }, + bifrost: { + title: "BifrostGatewayTitle", + label: "WayfinderGatewayLabel", + help: "BifrostGatewayHelp", + }, + aixy: { + title: "AixyGatewayTitle", + label: "AixyGatewayLabel", + help: "AixyGatewayHelp", + }, +}; + +export function isGatewayProviderId(id: string): id is GatewayProviderId { + return Object.prototype.hasOwnProperty.call(GATEWAY_COPY, id); +} + interface Props { + providerId: GatewayProviderId; draft: string; error: string | null; busy: boolean; @@ -8,10 +39,10 @@ interface Props { onDraftChange: (draft: string) => void; onSave: () => void; t: (key: LocaleKey) => string; - bifrost?: boolean; } export function WayfinderGatewaySection({ + providerId, draft, error, busy, @@ -19,13 +50,13 @@ export function WayfinderGatewaySection({ onDraftChange, onSave, t, - bifrost = false, }: Props) { + const copy = GATEWAY_COPY[providerId]; return (
-

{bifrost ? "Bifrost gateway" : t("WayfinderGatewayTitle")}

+

{t(copy.title)}

-

- {bifrost ? "Base URL of your Bifrost gateway." : t("WayfinderGatewayHelp")} -

+

{t(copy.help)}

{error &&

{error}

} + {period !== null && } +
- {([7, 30, 0] as const).map((days) => ( - - ))} -
)} + {!compactOverview && openAiApiUsage && ( +
+ {t("OpenAIChartTitle")} +
+ {/* Tray cards open often; skip the bar entrance animation there. */} + +
+
+ )} + {!compactOverview && (localUsage || hasPace || hasCharts) && (
{t("PanelUsageDetails")} diff --git a/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx b/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx new file mode 100644 index 0000000000..4dc02025fc --- /dev/null +++ b/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx @@ -0,0 +1,174 @@ +import { render, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const tauriMocks = vi.hoisted(() => ({ + getProviderChartData: vi.fn(), + getDeepSeekPricingStatus: vi.fn(), + getLocaleStrings: vi.fn(), + setUiLanguage: vi.fn(), + claudeAccountsList: vi.fn(), +})); + +const eventMocks = vi.hoisted(() => ({ + listen: vi.fn(), +})); + +vi.mock("../lib/tauri", async (importOriginal) => ({ + ...(await importOriginal()), + ...tauriMocks, +})); +vi.mock("@tauri-apps/api/event", () => eventMocks); + +import { LocaleProvider } from "../i18n/LocaleProvider"; +import { buildBundle } from "../test/localeHarness"; +import type { OpenAiApiUsageSnapshot, ProviderUsageSnapshot } from "../types/bridge"; +import MenuCard from "./MenuCard"; +import { describeCard } from "./MenuCardDetails"; + +const START = 1_788_220_800; + +function usage(days = 2): OpenAiApiUsageSnapshot { + return { + historyDays: 30, + projectId: null, + daily: Array.from({ length: days }, (_, i) => ({ + startTime: START + i * 86_400, + endTime: START + (i + 1) * 86_400, + costUsd: 1 + i, + requests: 10, + inputTokens: 1000, + cachedInputTokens: 100, + outputTokens: 400, + totalTokens: 1400, + lineItems: [{ name: "Text tokens", costUsd: 1 + i }], + models: [], + })), + }; +} + +function rateWindow() { + return { + usedPercent: 10, + remainingPercent: 90, + windowMinutes: null, + resetsAt: null, + resetDescription: null, + isExhausted: false, + reservePercent: null, + reserveDescription: null, + reserveWillLastToReset: false, + reserveEtaSeconds: null, + }; +} + +function snapshot( + providerId: string, + openAiApiUsage: OpenAiApiUsageSnapshot | null, + error: string | null = null, +): ProviderUsageSnapshot { + return { + providerId, + displayName: providerId, + primary: rateWindow(), + selectedMetric: rateWindow(), + primaryLabel: "Session", + secondary: null, + modelSpecific: null, + tertiary: null, + extraRateWindows: [], + cost: null, + planName: null, + accountEmail: null, + sourceLabel: "api", + updatedAt: "2026-05-24T00:00:00Z", + error, + errorState: "unknown", + pace: null, + accountOrganization: null, + trayStatusLabel: null, + fetchDurationMs: null, + openAiApiUsage, + }; +} + +function renderCard(snap: ProviderUsageSnapshot, compactOverview = false) { + return render( + + + , + ); +} + +describe("MenuCard OpenAI daily usage section", () => { + beforeEach(() => { + vi.clearAllMocks(); + tauriMocks.claudeAccountsList.mockResolvedValue([]); + tauriMocks.getLocaleStrings.mockResolvedValue(buildBundle({})); + tauriMocks.getDeepSeekPricingStatus.mockResolvedValue(null); + tauriMocks.getProviderChartData.mockResolvedValue({ + providerId: "openaiapi", + costHistory: [], + creditsHistory: [], + usageBreakdown: [], + localUsage: null, + }); + eventMocks.listen.mockResolvedValue(() => {}); + }); + + it("renders a collapsed Daily usage section for openaiapi", async () => { + const { container } = renderCard(snapshot("openaiapi", usage())); + await waitFor(() => { + expect(container.querySelector(".menu-card__daily-usage")).not.toBeNull(); + }); + const section = container.querySelector(".menu-card__daily-usage") as HTMLDetailsElement; + expect(section.tagName).toBe("DETAILS"); + expect(section.open).toBe(false); + expect(section.querySelector("summary")).toHaveTextContent("OpenAIChartTitle"); + expect(section.querySelectorAll('[role="option"]')).toHaveLength(2); + }); + + it("does not show another provider's OpenAI payload (provider siloing)", async () => { + const { container } = renderCard(snapshot("claude", usage())); + await waitFor(() => { + expect(tauriMocks.getProviderChartData).toHaveBeenCalled(); + }); + expect(container.querySelector(".menu-card__daily-usage")).toBeNull(); + expect(container.querySelector('[role="option"]')).toBeNull(); + }); + + it("hides the section on error, in compact overview and with an empty window", async () => { + const errored = renderCard(snapshot("openaiapi", usage(), "boom")); + await waitFor(() => expect(tauriMocks.getLocaleStrings).toHaveBeenCalled()); + expect(errored.container.querySelector(".menu-card__daily-usage")).toBeNull(); + errored.unmount(); + + const compact = renderCard(snapshot("openaiapi", usage()), true); + expect(compact.container.querySelector(".menu-card__daily-usage")).toBeNull(); + compact.unmount(); + + const empty = renderCard(snapshot("openaiapi", usage(0))); + expect(empty.container.querySelector(".menu-card__daily-usage")).toBeNull(); + }); +}); + +describe("describeCard OpenAI usage presence", () => { + it("is present only for openaiapi with days and no error", () => { + expect(describeCard(snapshot("openaiapi", usage()), null, []).openAiApiUsage).not.toBeNull(); + expect(describeCard(snapshot("claude", usage()), null, []).openAiApiUsage).toBeNull(); + expect(describeCard(snapshot("openaiapi", usage(0)), null, []).openAiApiUsage).toBeNull(); + expect(describeCard(snapshot("openaiapi", null), null, []).openAiApiUsage).toBeNull(); + expect(describeCard(snapshot("openaiapi", usage(), "boom"), null, []).openAiApiUsage).toBeNull(); + }); + + it("makes a card that only has daily usage count as having details", () => { + expect(describeCard(snapshot("openaiapi", usage()), null, []).hasDetails).toBe(true); + expect(describeCard(snapshot("openaiapi", usage(0)), null, []).hasDetails).toBe(false); + }); + + it("keeps compact overview header-only when daily usage is the only content", () => { + expect(describeCard(snapshot("openaiapi", usage()), null, [], "detailed", true, true).hasDetails).toBe(false); + }); +}); diff --git a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx new file mode 100644 index 0000000000..de38174925 --- /dev/null +++ b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx @@ -0,0 +1,226 @@ +import { fireEvent, render, screen, within } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import type { LocaleKey } from "../i18n/keys"; +import type { + OpenAiApiDailyUsageSnapshot, + OpenAiApiUsageSnapshot, +} from "../types/bridge"; +import { OpenAIApiUsageChart } from "./OpenAIApiUsageChart"; + +const DAY = 86_400; +// 2026-09-01T00:00:00Z +const START = 1_788_220_800; + +const strings: Partial> = { + OpenAIChartMore: "+{} more", + OpenAIChartWindowToday: "Today", + OpenAIChartWindowDays: "Last {} days", + UsageSpendRequests: "requests", +}; +const t = (key: LocaleKey) => strings[key] ?? key; + +function day( + index: number, + overrides: Partial = {}, +): OpenAiApiDailyUsageSnapshot { + return { + startTime: START + index * DAY, + endTime: START + (index + 1) * DAY, + costUsd: 1 + index, + requests: 10 * (index + 1), + inputTokens: 1_000 * (index + 1), + cachedInputTokens: 100 * (index + 1), + outputTokens: 500 * (index + 1), + totalTokens: 1_500 * (index + 1), + lineItems: [{ name: "Text tokens", costUsd: 1 + index }], + models: [ + { + name: `gpt-model-${index}`, + requests: 10 * (index + 1), + inputTokens: 1_000 * (index + 1), + cachedInputTokens: 100 * (index + 1), + outputTokens: 500 * (index + 1), + totalTokens: 1_500 * (index + 1), + }, + ], + ...overrides, + }; +} + +function usage(days: OpenAiApiDailyUsageSnapshot[], historyDays = 30): OpenAiApiUsageSnapshot { + return { historyDays, projectId: null, daily: days }; +} + +function renderChart(snapshot: OpenAiApiUsageSnapshot, onLayoutChange?: () => void) { + return render( + , + ); +} + +function bars() { + return screen.getAllByRole("option"); +} + +describe("OpenAIApiUsageChart", () => { + it("selects the latest day by default and shows its detail", () => { + renderChart(usage([day(0), day(1), day(2)])); + expect(bars()).toHaveLength(3); + expect(bars()[2]).toHaveAttribute("aria-selected", "true"); + expect(bars()[2]).toHaveAttribute("tabindex", "0"); + expect(bars()[0]).toHaveAttribute("tabindex", "-1"); + const detail = document.querySelector(".openai-usage__detail") as HTMLElement; + expect(within(detail).getByText("2026-09-03")).toBeInTheDocument(); + // Day cost in the stat grid plus the single line item carrying the same cost. + expect(within(detail).getAllByText("$3.00")).toHaveLength(2); + expect(within(detail).getByText("30")).toBeInTheDocument(); + expect(within(detail).getByText("3,000")).toBeInTheDocument(); + expect(within(detail).getByText("300")).toBeInTheDocument(); + expect(within(detail).getByText("1,500")).toBeInTheDocument(); + expect(within(detail).getByText("gpt-model-2")).toBeInTheDocument(); + }); + + it("shows the window total in the footer", () => { + renderChart(usage([day(0), day(1), day(2)])); + expect(screen.getByText("Last 30 days").nextElementSibling).toHaveTextContent("$6.00"); + }); + + it("labels a one-day window Today", () => { + renderChart(usage([day(0)], 1)); + expect(screen.getByText("Today")).toBeInTheDocument(); + }); + + it("selects a day on hover", () => { + renderChart(usage([day(0), day(1), day(2)])); + fireEvent.mouseEnter(bars()[0]); + expect(bars()[0]).toHaveAttribute("aria-selected", "true"); + expect(document.querySelector(".openai-usage__day")).toHaveTextContent("2026-09-01"); + expect(screen.getByText("gpt-model-0")).toBeInTheDocument(); + expect(screen.queryByText("gpt-model-2")).toBeNull(); + }); + + it("moves the selection with Left/Right/Home/End and keeps focus on the bar", () => { + renderChart(usage([day(0), day(1), day(2), day(3)])); + const list = screen.getByRole("listbox"); + fireEvent.keyDown(list, { key: "ArrowLeft" }); + expect(bars()[2]).toHaveAttribute("aria-selected", "true"); + expect(bars()[2]).toHaveFocus(); + fireEvent.keyDown(list, { key: "Home" }); + expect(bars()[0]).toHaveAttribute("aria-selected", "true"); + fireEvent.keyDown(list, { key: "ArrowLeft" }); + expect(bars()[0]).toHaveAttribute("aria-selected", "true"); + fireEvent.keyDown(list, { key: "End" }); + expect(bars()[3]).toHaveAttribute("aria-selected", "true"); + fireEvent.keyDown(list, { key: "ArrowRight" }); + expect(bars()[3]).toHaveAttribute("aria-selected", "true"); + expect(bars()[3]).toHaveAttribute("tabindex", "0"); + }); + + it("keeps the selected day when a refresh delivers new data", () => { + const { rerender } = renderChart(usage([day(0), day(1), day(2)])); + fireEvent.mouseEnter(bars()[1]); + rerender( + , + ); + expect(bars()).toHaveLength(4); + expect(bars()[1]).toHaveAttribute("aria-selected", "true"); + }); + + it("toggles between Cost and Tokens with the tab strip", () => { + renderChart(usage([day(0), day(1)])); + const cost = screen.getByRole("tab", { name: "OpenAIChartMetricCost" }); + const tokens = screen.getByRole("tab", { name: "OpenAIChartMetricTokens" }); + expect(cost).toHaveAttribute("aria-selected", "true"); + expect(tokens).toHaveAttribute("tabindex", "-1"); + expect(bars()[1]).toHaveAttribute("aria-label", "2026-09-02: $2.00"); + + fireEvent.click(tokens); + expect(tokens).toHaveAttribute("aria-selected", "true"); + expect(bars()[1]).toHaveAttribute("aria-label", "2026-09-02: 3K"); + expect(screen.getByText("Last 30 days").nextElementSibling).toHaveTextContent("4.5K"); + expect(screen.getByRole("tabpanel")).toHaveAttribute("aria-labelledby", tokens.id); + }); + + it("moves between the metric tabs with the arrow keys", () => { + renderChart(usage([day(0), day(1)])); + const cost = screen.getByRole("tab", { name: "OpenAIChartMetricCost" }); + const tokens = screen.getByRole("tab", { name: "OpenAIChartMetricTokens" }); + fireEvent.keyDown(cost, { key: "ArrowRight" }); + expect(tokens).toHaveAttribute("aria-selected", "true"); + expect(tokens).toHaveFocus(); + fireEvent.keyDown(tokens, { key: "ArrowLeft" }); + expect(cost).toHaveAttribute("aria-selected", "true"); + expect(cost).toHaveFocus(); + }); + + it("caps line items and models at five rows plus a more row", () => { + const many = day(0, { + lineItems: Array.from({ length: 7 }, (_, i) => ({ name: `item-${i}`, costUsd: 7 - i })), + models: Array.from({ length: 6 }, (_, i) => ({ + name: `model-${i}`, + requests: 1, + inputTokens: 1, + cachedInputTokens: 0, + outputTokens: 1, + totalTokens: 2, + })), + }); + renderChart(usage([many])); + expect(screen.getByText("item-4")).toBeInTheDocument(); + expect(screen.queryByText("item-5")).toBeNull(); + expect(screen.getByText("+2 more")).toBeInTheDocument(); + expect(screen.getByText("model-4")).toBeInTheDocument(); + expect(screen.queryByText("model-5")).toBeNull(); + expect(screen.getByText("+1 more")).toBeInTheDocument(); + }); + + it("formats sub-cent costs and zero-token days", () => { + renderChart( + usage([ + day(0, { + costUsd: 0.004, + totalTokens: 0, + inputTokens: 0, + cachedInputTokens: 0, + outputTokens: 0, + requests: 0, + lineItems: [], + models: [], + }), + ]), + ); + const detail = document.querySelector(".openai-usage__detail") as HTMLElement; + expect(within(detail).getByText("<$0.01")).toBeInTheDocument(); + expect(within(detail).queryByText("OpenAIChartLineItems")).toBeNull(); + expect(within(detail).queryByText("OpenAIChartModels")).toBeNull(); + }); + + it("draws only the latest 60 days and relabels the window to match", () => { + const days = Array.from({ length: 90 }, (_, i) => day(i)); + renderChart(usage(days, 90)); + expect(bars()).toHaveLength(60); + expect(screen.getByText("Last 60 days")).toBeInTheDocument(); + }); + + it("shows the empty message when the window has no days", () => { + renderChart(usage([])); + expect(screen.getByText("DetailChartEmpty")).toBeInTheDocument(); + expect(screen.queryByRole("listbox")).toBeNull(); + }); + + it("asks the host to re-measure when the detail panel changes", async () => { + const onLayoutChange = vi.fn(); + renderChart(usage([day(0), day(1)]), onLayoutChange); + await vi.waitFor(() => expect(onLayoutChange).toHaveBeenCalledTimes(1)); + fireEvent.mouseEnter(bars()[0]); + await vi.waitFor(() => expect(onLayoutChange).toHaveBeenCalledTimes(2)); + }); +}); diff --git a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx new file mode 100644 index 0000000000..9177f10e5d --- /dev/null +++ b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx @@ -0,0 +1,239 @@ +import { type KeyboardEvent, useEffect, useId, useMemo, useRef, useState } from "react"; +import type { LocaleKey } from "../i18n/keys"; +import type { + OpenAiApiDailyUsageSnapshot, + OpenAiApiUsageSnapshot, +} from "../types/bridge"; +import { BarChart } from "./charts/BarChart"; +import { providerCostColor } from "./charts/chartPalette"; + +/** + * Per-UTC-day chart for the OpenAI Admin API (upstream 0.66.0). Cost or token + * bars, one per day; hovering or focusing a bar selects the day and fills the + * detail panel (requests, token split, line items, models). USD only. + */ + +type Metric = "cost" | "tokens"; +type T = (key: LocaleKey) => string; + +const METRICS: readonly Metric[] = ["cost", "tokens"]; +/** Bars stay at least 1px wide inside the shared 280px chart geometry. */ +const MAX_CHART_DAYS = 60; +const MAX_DETAIL_ROWS = 5; + +const usdFormat = new Intl.NumberFormat("en-US", { + style: "currency", + currency: "USD", +}); +const compactFormat = new Intl.NumberFormat("en-US", { + notation: "compact", + maximumFractionDigits: 1, +}); +const countFormat = new Intl.NumberFormat("en-US"); + +function formatUsd(value: number): string { + if (value > 0 && value < 0.01) return "<$0.01"; + return usdFormat.format(value); +} + +function dayLabel(day: OpenAiApiDailyUsageSnapshot): string { + return new Date(day.startTime * 1000).toISOString().slice(0, 10); +} + +function metricValue(day: OpenAiApiDailyUsageSnapshot, metric: Metric): number { + return metric === "cost" ? day.costUsd : day.totalTokens; +} + +function formatMetric(value: number, metric: Metric): string { + return metric === "cost" ? formatUsd(value) : compactFormat.format(value); +} + +function windowLabel(days: number, t: T): string { + return days === 1 + ? t("OpenAIChartWindowToday") + : t("OpenAIChartWindowDays").replace("{}", String(days)); +} + +interface DetailRow { + key: string; + name: string; + value: string; +} + +function DetailList({ + title, + rows, + t, +}: { + title: string; + rows: DetailRow[]; + t: T; +}) { + if (rows.length === 0) return null; + const hidden = rows.length - MAX_DETAIL_ROWS; + return ( +
+
{title}
+
    + {rows.slice(0, MAX_DETAIL_ROWS).map((row) => ( +
  • + {row.name} + {row.value} +
    • + ))} + {hidden > 0 && ( +
  • + {t("OpenAIChartMore").replace("{}", String(hidden))} +
    • + )} +
+
+ ); +} + +function DayDetail({ day, t }: { day: OpenAiApiDailyUsageSnapshot; t: T }) { + const stats: Array<[LocaleKey, string]> = [ + ["OpenAIChartMetricCost", formatUsd(day.costUsd)], + ["OpenAIChartRequests", countFormat.format(day.requests)], + ["OpenAIChartInputTokens", countFormat.format(day.inputTokens)], + ["OpenAIChartCachedTokens", countFormat.format(day.cachedInputTokens)], + ["OpenAIChartOutputTokens", countFormat.format(day.outputTokens)], + ]; + const lineItems = day.lineItems.map((item, index) => ({ + key: `${item.name}-${index}`, + name: item.name, + value: formatUsd(item.costUsd), + })); + const models = day.models.map((model, index) => ({ + key: `${model.name}-${index}`, + name: model.name, + value: `${compactFormat.format(model.totalTokens)} · ${countFormat.format(model.requests)} ${t("UsageSpendRequests")}`, + })); + return ( +
+
{dayLabel(day)}
+
+ {stats.map(([labelKey, value]) => ( +
+
{t(labelKey)}
+
{value}
+
+ ))} +
+ + +
+ ); +} + +export interface OpenAIApiUsageChartProps { + usage: OpenAiApiUsageSnapshot; + animations: boolean; + t: T; + /** Popover hosts resize to content; called after the detail panel height may change. */ + onLayoutChange?: () => void; +} + +export function OpenAIApiUsageChart({ + usage, + animations, + t, + onLayoutChange, +}: OpenAIApiUsageChartProps) { + const [metric, setMetric] = useState("cost"); + // Selection is keyed by bucket start so it survives a refresh; null = latest day. + const [selectedStart, setSelectedStart] = useState(null); + const baseId = useId(); + const tabRefs = useRef>([]); + const days = useMemo(() => usage.daily.slice(-MAX_CHART_DAYS), [usage.daily]); + const points = useMemo( + () => days.map((day) => ({ label: dayLabel(day), value: metricValue(day, metric) })), + [days, metric], + ); + + const layoutKey = `${metric}:${selectedStart ?? "latest"}:${days.length}`; + const layoutCallback = useRef(onLayoutChange); + layoutCallback.current = onLayoutChange; + useEffect(() => { + const frame = requestAnimationFrame(() => layoutCallback.current?.()); + return () => cancelAnimationFrame(frame); + }, [layoutKey]); + + if (days.length === 0) { + return ( +
+
{t("DetailChartEmpty")}
+
+ ); + } + + const found = days.findIndex((day) => day.startTime === selectedStart); + const selectedIndex = found >= 0 ? found : days.length - 1; + const total = days.reduce((sum, day) => sum + metricValue(day, metric), 0); + const trimmed = days.length < usage.daily.length; + const windowDays = trimmed ? days.length : usage.historyDays; + const tabLabel = (m: Metric) => + t(m === "cost" ? "OpenAIChartMetricCost" : "OpenAIChartMetricTokens"); + const tabId = (m: Metric) => `${baseId}-tab-${m}`; + const panelId = `${baseId}-panel`; + + const onTabKeyDown = (e: KeyboardEvent) => { + const current = METRICS.indexOf(metric); + let next: number; + if (e.key === "ArrowLeft") next = (current + METRICS.length - 1) % METRICS.length; + else if (e.key === "ArrowRight") next = (current + 1) % METRICS.length; + else if (e.key === "Home") next = 0; + else if (e.key === "End") next = METRICS.length - 1; + else return; + e.preventDefault(); + setMetric(METRICS[next]); + tabRefs.current[next]?.focus(); + }; + + return ( +
+
+ {METRICS.map((m, index) => ( + + ))} +
+
+ formatMetric(v, metric)} + animations={animations} + emptyMessage={t("DetailChartEmpty")} + selection={{ + index: selectedIndex, + onSelect: (index) => setSelectedStart(days[index]?.startTime ?? null), + }} + /> +
{t("OpenAIChartHint")}
+ +
+ {windowLabel(windowDays, t)} + {formatMetric(total, metric)} +
+
+
+ ); +} diff --git a/apps/desktop-tauri/src/components/charts/BarChart.test.tsx b/apps/desktop-tauri/src/components/charts/BarChart.test.tsx index 12934d08e3..cb8e202caa 100644 --- a/apps/desktop-tauri/src/components/charts/BarChart.test.tsx +++ b/apps/desktop-tauri/src/components/charts/BarChart.test.tsx @@ -1,5 +1,5 @@ -import { render } from "@testing-library/react"; -import { describe, expect, it } from "vitest"; +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; import { BarChart } from "./BarChart"; describe("BarChart calendar slots", () => { @@ -48,3 +48,73 @@ describe("BarChart calendar slots", () => { expect((labels[1] as HTMLElement).style.transform).toBe(""); }); }); + +describe("BarChart controlled selection", () => { + const data = [ + { label: "a", value: 1 }, + { label: "b", value: 2 }, + { label: "c", value: 3 }, + ]; + + function setup(index = 2) { + const onSelect = vi.fn(); + const utils = render( + , + ); + return { onSelect, ...utils }; + } + + it("exposes a listbox of options with a roving tabindex", () => { + setup(1); + expect(screen.getByRole("listbox")).toHaveAttribute("aria-orientation", "horizontal"); + const options = screen.getAllByRole("option"); + expect(options.map((o) => o.getAttribute("aria-selected"))).toEqual(["false", "true", "false"]); + expect(options.map((o) => o.getAttribute("tabindex"))).toEqual(["-1", "0", "-1"]); + expect(options[1]).toHaveAttribute("data-selected", "true"); + expect(options[1]).toHaveAttribute("opacity", "1"); + expect(options[0]).toHaveAttribute("opacity", "0.6"); + }); + + it("reports hover and focus as a selection change", () => { + const { onSelect } = setup(); + const options = screen.getAllByRole("option"); + fireEvent.mouseEnter(options[0]); + expect(onSelect).toHaveBeenLastCalledWith(0); + fireEvent.focus(options[1]); + expect(onSelect).toHaveBeenLastCalledWith(1); + }); + + it("navigates with the arrow, Home and End keys and clamps at the ends", () => { + const { onSelect } = setup(1); + const list = screen.getByRole("listbox"); + fireEvent.keyDown(list, { key: "ArrowLeft" }); + expect(onSelect).toHaveBeenLastCalledWith(0); + fireEvent.keyDown(list, { key: "ArrowRight" }); + expect(onSelect).toHaveBeenLastCalledWith(2); + fireEvent.keyDown(list, { key: "Home" }); + expect(onSelect).toHaveBeenLastCalledWith(0); + fireEvent.keyDown(list, { key: "End" }); + expect(onSelect).toHaveBeenLastCalledWith(2); + onSelect.mockClear(); + fireEvent.keyDown(list, { key: "Tab" }); + expect(onSelect).not.toHaveBeenCalled(); + }); + + it("does not open the hover tooltip in selection mode", () => { + const { container } = setup(); + fireEvent.mouseMove(container.querySelector("svg") as SVGElement, { clientX: 10, clientY: 10 }); + expect(container.querySelector(".chart__tooltip")).toBeNull(); + }); + + it("stays a plain image chart without a selection prop", () => { + const { container } = render(); + expect(screen.queryByRole("listbox")).toBeNull(); + expect(screen.queryAllByRole("option")).toHaveLength(0); + expect(container.querySelector("[data-selected]")).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/components/charts/BarChart.tsx b/apps/desktop-tauri/src/components/charts/BarChart.tsx index 0151811f57..65608c8395 100644 --- a/apps/desktop-tauri/src/components/charts/BarChart.tsx +++ b/apps/desktop-tauri/src/components/charts/BarChart.tsx @@ -1,4 +1,4 @@ -import { useMemo, useRef, useState } from "react"; +import { type KeyboardEvent, useMemo, useRef, useState } from "react"; import { useChartAnimation } from "./useChartAnimation"; import { WIDTH, @@ -25,6 +25,16 @@ export interface BarChartPoint { value: number | null; } +/** + * Controlled bar selection. When set, the bars form a roving-tabindex listbox: + * hover or focus selects a bar, Left/Right move the selection, Home/End jump to + * the first/last bar, and the hover tooltip is left to the caller's detail view. + */ +export interface BarChartSelection { + index: number; + onSelect: (index: number) => void; +} + export interface BarChartProps { data: BarChartPoint[]; color?: string; @@ -35,6 +45,7 @@ export interface BarChartProps { animations?: boolean; /** Optional empty-state message rendered when `data.length === 0`. */ emptyMessage?: string; + selection?: BarChartSelection; } const DEFAULT_COLOR = "var(--chart-cost)"; @@ -48,10 +59,12 @@ export function BarChart({ ariaLabel, animations = true, emptyMessage, + selection, }: BarChartProps) { const fmt = valueFormatter ?? ((v: number) => v.toFixed(2)); const containerRef = useRef(null); const [hover, setHover] = useState<{ i: number; x: number; y: number } | null>(null); + const barRefs = useRef>([]); const anim = useChartAnimation(data.length, animations, [ data.length, @@ -91,6 +104,23 @@ export function BarChart({ }; const onLeave = () => setHover(null); + const selectedIndex = selection + ? Math.min(Math.max(selection.index, 0), data.length - 1) + : -1; + const onKeyDown = (e: KeyboardEvent) => { + if (!selection) return; + let next: number; + if (e.key === "ArrowLeft") next = selectedIndex - 1; + else if (e.key === "ArrowRight") next = selectedIndex + 1; + else if (e.key === "Home") next = 0; + else if (e.key === "End") next = data.length - 1; + else return; + e.preventDefault(); + next = Math.min(Math.max(next, 0), data.length - 1); + selection.onSelect(next); + barRefs.current[next]?.focus(); + }; + return (
{data.map((p, i) => { const base = p.value == null ? 1 : p.value === 0 ? 1 : Math.max(3, (p.value / max) * plotHeight); @@ -111,6 +143,7 @@ export function BarChart({ const bodyH = isPeak ? Math.max(0, barH - CAP_HEIGHT) : barH; const bodyY = isPeak ? y + CAP_HEIGHT : y; const isHovered = hover?.i === i; + const isSelected = i === selectedIndex; return ( @@ -120,11 +153,37 @@ export function BarChart({ width={barWidth} height={bodyH} fill={color} - opacity={p.value == null ? 0 : p.value === 0 ? 0.25 : isHovered ? 1 : 0.9} + opacity={ + p.value == null + ? 0 + : p.value === 0 + ? 0.25 + : selection + ? isSelected ? 1 : 0.6 + : isHovered ? 1 : 0.9 + } rx={1} className="chart__bar" - onMouseMove={p.value == null ? undefined : (e) => onMove(e, i)} - onMouseLeave={onLeave} + {...(selection + ? { + ref: (node: SVGRectElement | null) => { + barRefs.current[i] = node; + }, + role: "option", + "aria-selected": isSelected, + "aria-label": p.value == null ? p.label : `${p.label}: ${fmt(p.value)}`, + tabIndex: isSelected ? 0 : -1, + "data-selected": isSelected ? "true" : "false", + onMouseEnter: () => selection.onSelect(i), + onFocus: () => selection.onSelect(i), + } + : { + onMouseMove: + p.value == null + ? undefined + : (e: React.MouseEvent) => onMove(e, i), + onMouseLeave: onLeave, + })} > {p.value == null ? p.label : `${p.label}: ${fmt(p.value)}`} @@ -157,7 +216,7 @@ export function BarChart({ {data[data.length - 1].label} - {hover && !anim.running && ( + {!selection && hover && !anim.running && (
div { + display: flex; + flex-direction: column; + gap: 1px; + min-width: 0; +} + +.openai-usage__stats dt { + font-size: 10px; + color: var(--text-secondary); +} + +.openai-usage__stats dd { + margin: 0; + font-variant-numeric: tabular-nums; +} + +.openai-usage__list-title { + font-size: 10px; + color: var(--text-secondary); + margin-bottom: 2px; +} + +.openai-usage__list ul { + list-style: none; + margin: 0; + padding: 0; + display: flex; + flex-direction: column; + gap: 2px; +} + +.openai-usage__list li { + display: flex; + justify-content: space-between; + gap: 12px; + min-width: 0; +} + +.openai-usage__name { + min-width: 0; + overflow-wrap: anywhere; +} + +.openai-usage__amount { + flex-shrink: 0; + font-variant-numeric: tabular-nums; + color: var(--text-secondary); +} + +.openai-usage__list li.openai-usage__more { + color: var(--text-muted); +} + +.openai-usage__footer { + display: flex; + justify-content: space-between; + gap: 12px; + padding-top: 6px; + border-top: 1px solid var(--panel-border); + font-size: 11px; + color: var(--text-secondary); +} + +.openai-usage__footer strong { + color: var(--text-primary); + font-variant-numeric: tabular-nums; +} + @media (prefers-reduced-motion: reduce) { .chart__bar, .chart__line, diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx index 0dadfe492d..fc00d44dcd 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx @@ -409,6 +409,7 @@ export function ProviderDetailPane({ providerId={detail.id} accountEmail={detail.email} accentColor={providerAccentColors[detail.id]} + openAiApiUsage={detail.openAiApiUsage} t={t} /> diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.test.tsx index ade4cdcd3b..0268b6e616 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/charts/ChartsSection.test.tsx @@ -2,7 +2,7 @@ import { describe, expect, it, vi } from "vitest"; import { render, screen, waitFor } from "@testing-library/react"; import { ChartsSection } from "./ChartsSection"; import { getProviderChartData } from "../../../../../lib/tauri"; -import type { ProviderChartData } from "../../../../../types/bridge"; +import type { OpenAiApiUsageSnapshot, ProviderChartData } from "../../../../../types/bridge"; vi.mock("../../../../../lib/tauri", () => ({ getProviderChartData: vi.fn(), @@ -93,3 +93,75 @@ describe("ChartsSection tokens mode (upstream 0.50.0 #2930)", () => { expect(screen.queryByText("DetailChartTokens")).toBeNull(); }); }); + +describe("ChartsSection OpenAI API daily usage (upstream 0.66.0)", () => { + const openAiUsage: OpenAiApiUsageSnapshot = { + historyDays: 30, + projectId: null, + daily: [ + { + startTime: 1_788_220_800, + endTime: 1_788_307_200, + costUsd: 2.5, + requests: 12, + inputTokens: 1000, + cachedInputTokens: 100, + outputTokens: 400, + totalTokens: 1400, + lineItems: [{ name: "Text tokens", costUsd: 2.5 }], + models: [], + }, + ], + }; + + it("draws the per-day chart for openaiapi instead of the local-log tabs", async () => { + mockChart.mockResolvedValue(chartData({ providerId: "openaiapi" })); + render( + key} + />, + ); + expect(screen.getByText("OpenAIChartTitle", { selector: ".provider-detail-chart__title" })).toBeTruthy(); + expect(screen.getByRole("listbox")).toBeTruthy(); + expect(screen.getByText("Text tokens")).toBeTruthy(); + await waitFor(() => expect(mockChart).toHaveBeenCalled()); + expect(screen.queryByText("DetailChartCost")).toBeNull(); + }); + + it("renders nothing for openaiapi without usage or with an empty window", () => { + mockChart.mockResolvedValue(chartData({ providerId: "openaiapi" })); + const { container, rerender } = render( + key} />, + ); + expect(container.firstChild).toBeNull(); + rerender( + key} + />, + ); + expect(container.firstChild).toBeNull(); + }); + + it("ignores an OpenAI snapshot handed to another provider", async () => { + mockChart.mockResolvedValue(chartData({ providerId: "claude" })); + render( + key} + />, + ); + await waitFor(() => { + expect(screen.getByRole("tab", { selected: true }).textContent).toBe("DetailChartCost"); + }); + expect(screen.queryByText("OpenAIChartTitle")).toBeNull(); + expect(screen.queryByRole("listbox")).toBeNull(); + }); +}); 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 8257e3069c..30bf0b3a06 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 @@ -1,8 +1,13 @@ import { type CSSProperties, useEffect, useState } from "react"; import { getProviderChartData, getSettingsSnapshot } from "../../../../../lib/tauri"; import { providerSupportsChartData } from "../../../../../lib/providerCharts"; -import type { ProviderChartData, SettingsSnapshot } from "../../../../../types/bridge"; +import type { + OpenAiApiUsageSnapshot, + ProviderChartData, + SettingsSnapshot, +} from "../../../../../types/bridge"; import type { useLocale } from "../../../../../hooks/useLocale"; +import { OpenAIApiUsageChart } from "../../../../../components/OpenAIApiUsageChart"; import { CostHistoryChart } from "./CostHistoryChart"; import { CreditsHistoryChart } from "./CreditsHistoryChart"; import { TokensHistoryChart } from "./TokensHistoryChart"; @@ -15,6 +20,8 @@ interface Props { accountEmail: string | null; /** Per-provider accent color override (hex); applied as CSS --provider-accent. */ accentColor?: string; + /** Per-day OpenAI Admin API history (`openaiapi` only); drawn instead of the local-log tabs. */ + openAiApiUsage?: OpenAiApiUsageSnapshot | null; t: T; } @@ -29,7 +36,13 @@ type TabKey = "tokens" | "cost" | "credits" | "usage"; * Phase 10: fetches the latest settings snapshot so the animation flag feeds * through to each chart component. */ -export function ChartsSection({ providerId, accountEmail, accentColor, t }: Props) { +export function ChartsSection({ + providerId, + accountEmail, + accentColor, + openAiApiUsage, + t, +}: Props) { const [data, setData] = useState(null); const [active, setActive] = useState(null); const [animations, setAnimations] = useState(true); @@ -71,6 +84,19 @@ export function ChartsSection({ providerId, accountEmail, accentColor, t }: Prop }; }, [providerId]); + if (providerId === "openaiapi") { + if (!openAiApiUsage || openAiApiUsage.daily.length === 0) return null; + return ( +
+
{t("OpenAIChartTitle")}
+ +
+ ); + } + if (!data) return null; const hasCost = data.costHistory.length > 0; diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index ee6bdef0f4..01560daed3 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -1024,6 +1024,8 @@ export interface ProviderDetail { cost: CostSnapshotBridge | null; pace: PaceSnapshot | null; + /** Per-UTC-day OpenAI Admin API history; only the `openaiapi` Admin path sets it. */ + openAiApiUsage?: OpenAiApiUsageSnapshot | null; lastError: string | null; errorState: ProviderStateKind | null; diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..dff8d022c8 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -906,6 +906,19 @@ locale_keys! { DetailChartCredits, DetailChartUsageBreakdown, DetailChartEmpty, + OpenAIChartTitle, + OpenAIChartMetricCost, + OpenAIChartMetricTokens, + OpenAIChartHint, + OpenAIChartRequests, + OpenAIChartInputTokens, + OpenAIChartCachedTokens, + OpenAIChartOutputTokens, + OpenAIChartLineItems, + OpenAIChartModels, + OpenAIChartMore, + OpenAIChartWindowToday, + OpenAIChartWindowDays, DetailUpdatedPrefix, PanelAllProviders, PanelAllProvidersShort, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..b3738e3bb9 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -569,6 +569,19 @@ DetailChartRefreshing = Refreshing… DetailChartCredits = Credits used (30 days) DetailChartUsageBreakdown = Usage by service (30 days) DetailChartEmpty = No chart data yet. +OpenAIChartTitle = Daily API usage +OpenAIChartMetricCost = Cost +OpenAIChartMetricTokens = Tokens +OpenAIChartHint = Hover or focus a bar for details +OpenAIChartRequests = Requests +OpenAIChartInputTokens = Input tokens +OpenAIChartCachedTokens = Cached input tokens +OpenAIChartOutputTokens = Output tokens +OpenAIChartLineItems = Cost by line item +OpenAIChartModels = Models +OpenAIChartMore = +{ "{}" } more +OpenAIChartWindowToday = Today +OpenAIChartWindowDays = Last { "{}" } days DetailUpdatedPrefix = Updated PanelAllProviders = All providers PanelAllProvidersShort = All diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index 38dc0b05ae..d3b6cb5eb7 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -525,6 +525,19 @@ DetailChartRefreshing = Actualizando… DetailChartCredits = Créditos usados (30 días) DetailChartUsageBreakdown = Uso por servicio (30 días) DetailChartEmpty = Sin datos de gráfico aún. +OpenAIChartTitle = Uso diario de la API +OpenAIChartMetricCost = Costo +OpenAIChartMetricTokens = Tokens +OpenAIChartHint = Pasa el cursor o enfoca una barra para ver detalles +OpenAIChartRequests = Solicitudes +OpenAIChartInputTokens = Tokens de entrada +OpenAIChartCachedTokens = Tokens de entrada en caché +OpenAIChartOutputTokens = Tokens de salida +OpenAIChartLineItems = Costo por concepto +OpenAIChartModels = Modelos +OpenAIChartMore = +{ "{}" } más +OpenAIChartWindowToday = Hoy +OpenAIChartWindowDays = Últimos { "{}" } días DetailUpdatedPrefix = Actualizado PanelAllProviders = Todos los proveedores PanelAllProvidersShort = Todos diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index b6c2dc0fc8..b9a86c6fb0 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -507,6 +507,19 @@ DetailChartRefreshing = 更新中… DetailChartCredits = 使用クレジット（30日間） DetailChartUsageBreakdown = サービス別使用量（30日間） DetailChartEmpty = まだチャートデータはありません。 +OpenAIChartTitle = 日別 API 使用量 +OpenAIChartMetricCost = コスト +OpenAIChartMetricTokens = トークン +OpenAIChartHint = バーにカーソルを合わせるかフォーカスすると詳細を表示します +OpenAIChartRequests = リクエスト +OpenAIChartInputTokens = 入力トークン +OpenAIChartCachedTokens = キャッシュされた入力トークン +OpenAIChartOutputTokens = 出力トークン +OpenAIChartLineItems = 項目別コスト +OpenAIChartModels = モデル +OpenAIChartMore = 他 { "{}" } 件 +OpenAIChartWindowToday = 今日 +OpenAIChartWindowDays = 過去 { "{}" } 日間 DetailUpdatedPrefix = 更新 PanelAllProviders = すべてのプロバイダー PanelAllProvidersShort = すべて diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 299dcf1a54..f1001fb6f7 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -512,6 +512,19 @@ DetailChartRefreshing = 새로 고치는 중… DetailChartCredits = 사용 크레딧 (30일) DetailChartUsageBreakdown = 서비스별 사용량 (30일) DetailChartEmpty = 아직 차트 데이터가 없습니다. +OpenAIChartTitle = 일별 API 사용량 +OpenAIChartMetricCost = 비용 +OpenAIChartMetricTokens = 토큰 +OpenAIChartHint = 막대에 마우스를 올리거나 포커스하면 세부 정보가 표시됩니다 +OpenAIChartRequests = 요청 +OpenAIChartInputTokens = 입력 토큰 +OpenAIChartCachedTokens = 캐시된 입력 토큰 +OpenAIChartOutputTokens = 출력 토큰 +OpenAIChartLineItems = 항목별 비용 +OpenAIChartModels = 모델 +OpenAIChartMore = +{ "{}" }개 더 +OpenAIChartWindowToday = 오늘 +OpenAIChartWindowDays = 최근 { "{}" }일 DetailUpdatedPrefix = 업데이트됨 PanelAllProviders = 모든 제공자 PanelAllProvidersShort = 전체 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index f978d540e7..378349e707 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -491,6 +491,19 @@ DetailChartRefreshing = Обновление… DetailChartCredits = Использовано кредитов (30 дней) DetailChartUsageBreakdown = Использование службой (30 дней) DetailChartEmpty = Данных диаграммы пока нет. +OpenAIChartTitle = Ежедневное использование API +OpenAIChartMetricCost = Стоимость +OpenAIChartMetricTokens = Токены +OpenAIChartHint = Наведите курсор на столбец или выберите его, чтобы увидеть подробности +OpenAIChartRequests = Запросы +OpenAIChartInputTokens = Входные токены +OpenAIChartCachedTokens = Кэшированные входные токены +OpenAIChartOutputTokens = Выходные токены +OpenAIChartLineItems = Стоимость по статьям +OpenAIChartModels = Модели +OpenAIChartMore = ещё { "{}" } +OpenAIChartWindowToday = Сегодня +OpenAIChartWindowDays = Последние { "{}" } дн. DetailUpdatedPrefix = Обновлено PanelAllProviders = Все провайдеры PanelAllProvidersShort = Все diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index 08de23e2e3..b63298695c 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -526,6 +526,19 @@ DetailChartCost = Maliyet (30 gün) DetailChartCredits = Kullanılan krediler (30 gün) DetailChartUsageBreakdown = Hizmete göre kullanım (30 gün) DetailChartEmpty = Henüz grafik verisi yok. +OpenAIChartTitle = Günlük API kullanımı +OpenAIChartMetricCost = Maliyet +OpenAIChartMetricTokens = Token +OpenAIChartHint = Ayrıntılar için bir çubuğun üzerine gelin veya odaklanın +OpenAIChartRequests = İstekler +OpenAIChartInputTokens = Girdi token'ları +OpenAIChartCachedTokens = Önbelleğe alınan girdi token'ları +OpenAIChartOutputTokens = Çıktı token'ları +OpenAIChartLineItems = Kalem bazında maliyet +OpenAIChartModels = Modeller +OpenAIChartMore = +{ "{}" } daha +OpenAIChartWindowToday = Bugün +OpenAIChartWindowDays = Son { "{}" } gün DetailUpdatedPrefix = Güncellendi PanelAllProviders = Tüm sağlayıcılar PanelAllProvidersShort = Tümü diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index daf8e3c9af..6032c75474 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -506,6 +506,19 @@ DetailChartRefreshing = 刷新中… DetailChartCredits = 已用额度（30 天） DetailChartUsageBreakdown = 按服务划分的用量（30 天） DetailChartEmpty = 暂无图表数据。 +OpenAIChartTitle = 每日 API 用量 +OpenAIChartMetricCost = 费用 +OpenAIChartMetricTokens = Token +OpenAIChartHint = 悬停或聚焦柱形以查看详情 +OpenAIChartRequests = 请求 +OpenAIChartInputTokens = 输入 Token +OpenAIChartCachedTokens = 缓存输入 Token +OpenAIChartOutputTokens = 输出 Token +OpenAIChartLineItems = 按项目费用 +OpenAIChartModels = 模型 +OpenAIChartMore = 另有 { "{}" } 项 +OpenAIChartWindowToday = 今天 +OpenAIChartWindowDays = 最近 { "{}" } 天 DetailUpdatedPrefix = 更新于 PanelAllProviders = 所有提供者 PanelAllProvidersShort = 全部 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 31bb3bcf02..2af948efb9 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -506,6 +506,19 @@ DetailChartRefreshing = 重新整理中… DetailChartCredits = 已用額度（30 天） DetailChartUsageBreakdown = 按服務劃分的用量（30 天） DetailChartEmpty = 暫無圖表資料。 +OpenAIChartTitle = 每日 API 用量 +OpenAIChartMetricCost = 費用 +OpenAIChartMetricTokens = Token +OpenAIChartHint = 將滑鼠移到或聚焦長條即可查看詳細資料 +OpenAIChartRequests = 請求 +OpenAIChartInputTokens = 輸入 Token +OpenAIChartCachedTokens = 快取輸入 Token +OpenAIChartOutputTokens = 輸出 Token +OpenAIChartLineItems = 依項目費用 +OpenAIChartModels = 模型 +OpenAIChartMore = 另有 { "{}" } 項 +OpenAIChartWindowToday = 今天 +OpenAIChartWindowDays = 最近 { "{}" } 天 DetailUpdatedPrefix = 更新於 PanelAllProviders = 所有提供者 PanelAllProvidersShort = 全部 From 518230ffa861ac04f3699ba74cbd11cc6cfe3dbf Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:19:56 +0700 Subject: [PATCH 116/259] Port upstream 0.64.0: LiteLLM key-bound identity, user/team budgets, private-network HTTP --- rust/src/locale/en-US.ftl | 2 +- rust/src/providers/litellm/endpoint.rs | 172 +++++++++++++++ rust/src/providers/litellm/info.rs | 271 +++++++++++++++++++++++ rust/src/providers/litellm/mod.rs | 219 ++++++------------- rust/src/providers/litellm/tests.rs | 278 ++++++++++++++++++++++++ rust/src/settings/provider_workspace.rs | 28 ++- 6 files changed, 812 insertions(+), 158 deletions(-) create mode 100644 rust/src/providers/litellm/endpoint.rs create mode 100644 rust/src/providers/litellm/info.rs create mode 100644 rust/src/providers/litellm/tests.rs diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..ec1baafc9b 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -784,7 +784,7 @@ OpenAiProjectIdHelp = Leave blank for organization-wide usage. Set a project ID LiteLlmApiTitle = LiteLLM API LiteLlmBaseUrlLabel = Base URL LiteLlmBaseUrlPlaceholder = https://litellm.example.com -LiteLlmBaseUrlHelp = Used with the saved API key for LiteLLM /key/info. +LiteLlmBaseUrlHelp = Used with the saved API key for LiteLLM key, user, and team info. Use HTTPS, or HTTP on a loopback or private-network address. DevinApiTitle = Devin API DevinOrganizationLabel = Organization DevinOrganizationPlaceholder = org/acme diff --git a/rust/src/providers/litellm/endpoint.rs b/rust/src/providers/litellm/endpoint.rs new file mode 100644 index 0000000000..b175ed137e --- /dev/null +++ b/rust/src/providers/litellm/endpoint.rs @@ -0,0 +1,172 @@ +//! LiteLLM base-URL policy and management-route URLs. +//! +//! Upstream `litellm.ts` declares the `LITELLM_BASE_URL` endpoint with the +//! `https-or-private-network-http` policy: HTTPS anywhere, plain HTTP only for +//! loopback, RFC 1918, link-local, IPv6 unique-local, and `.local` hosts. + +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + +use reqwest::Url; + +use crate::core::ProviderError; + +const INVALID_BASE: &str = "LiteLLM base URL must use HTTPS, or HTTP on a loopback or private-network address, without embedded credentials."; + +/// Validate a LiteLLM base URL. A scheme-less value is treated as HTTPS. +pub(crate) fn validated_base_url(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(ProviderError::Other("LiteLLM base URL is empty".into())); + } + let lower = trimmed.to_ascii_lowercase(); + if ["%2f", "%5c", "%3f", "%23", "%40", "%3a"] + .iter() + .any(|encoded| lower.contains(encoded)) + { + return Err(ProviderError::Other( + "LiteLLM base URL must not contain encoded host delimiters".into(), + )); + } + let candidate = if trimmed.contains("://") { + trimmed.to_string() + } else { + format!("https://{trimmed}") + }; + let url = Url::parse(&candidate) + .map_err(|e| ProviderError::Other(format!("Invalid LiteLLM base URL: {e}")))?; + let host = url + .host_str() + .ok_or_else(|| ProviderError::Other("LiteLLM base URL must include a host".into()))?; + let scheme_ok = match url.scheme() { + "https" => true, + "http" => is_private_network_host(host), + _ => false, + }; + if !scheme_ok + || !url.username().is_empty() + || url.password().is_some() + || host.contains('%') + || host.chars().any(|c| c.is_control() || c.is_whitespace()) + { + return Err(ProviderError::Other(INVALID_BASE.into())); + } + Ok(url) +} + +/// Build `{base}/{path}` for a management route. A trailing `/v1` on the base +/// is dropped, and the base path and query are otherwise preserved. `query` +/// replaces the base query when given. +pub(super) fn management_url( + base: &str, + path: &str, + query: Option<(&str, &str)>, +) -> Result { + let mut url = validated_base_url(base)?; + let trimmed = url.path().trim_end_matches('/'); + let root = trimmed.strip_suffix("/v1").unwrap_or(trimmed).to_string(); + url.set_path(&format!("{root}/{path}")); + url.set_fragment(None); + if let Some((key, value)) = query { + url.query_pairs_mut().clear().append_pair(key, value); + } + Ok(url) +} + +fn is_private_network_host(host: &str) -> bool { + let normalized = host.trim_end_matches('.').to_ascii_lowercase(); + if normalized == "localhost" + || normalized.ends_with(".localhost") + || normalized.ends_with(".local") + { + return true; + } + let ip_candidate = normalized + .strip_prefix('[') + .and_then(|value| value.strip_suffix(']')) + .unwrap_or(&normalized); + match ip_candidate.parse::() { + Ok(IpAddr::V4(ip)) => is_private_ipv4(ip), + Ok(IpAddr::V6(ip)) => is_private_ipv6(ip), + Err(_) => false, + } +} + +fn is_private_ipv4(ip: Ipv4Addr) -> bool { + ip.is_loopback() || ip.is_private() || ip.is_link_local() +} + +fn is_private_ipv6(ip: Ipv6Addr) -> bool { + if let Some(mapped) = ip.to_ipv4_mapped() { + return is_private_ipv4(mapped); + } + ip.is_loopback() + || (ip.segments()[0] & 0xfe00) == 0xfc00 + || (ip.segments()[0] & 0xffc0) == 0xfe80 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn allows_https_anywhere_and_private_network_http() { + for value in [ + "https://litellm.example.com", + "litellm.example.com", + "http://localhost:4000", + "http://127.0.0.1:4000", + "http://[::1]:4000", + "http://10.1.2.3", + "http://172.16.0.9", + "http://192.168.1.20:4000", + "http://169.254.10.10", + "http://[fd12:3456::1]", + "http://[fe80::1]", + "http://proxy.local:4000", + ] { + assert!(validated_base_url(value).is_ok(), "rejected {value}"); + } + } + + #[test] + fn rejects_public_http_credentials_and_encoded_delimiters() { + for value in [ + "", + "http://litellm.example.com", + "http://8.8.8.8", + "http://172.32.0.1", + "http://[2001:db8::1]", + "http://example.com.evil.test", + "ftp://10.0.0.1", + "https://user:pass@litellm.example.com", + "http://user@10.0.0.1", + "https://example.com%2f.evil.test", + ] { + assert!(validated_base_url(value).is_err(), "accepted {value}"); + } + } + + #[test] + fn management_url_strips_v1_and_keeps_subpath() { + let url = management_url("https://h.example.com/litellm/v1/", "key/info", None).unwrap(); + assert_eq!(url.as_str(), "https://h.example.com/litellm/key/info"); + let url = management_url("http://10.0.0.2:4000/v1", "key/info", None).unwrap(); + assert_eq!(url.as_str(), "http://10.0.0.2:4000/key/info"); + } + + #[test] + fn management_url_encodes_query_and_replaces_base_query() { + let url = management_url( + "https://h.example.com?token=abc", + "user/info", + Some(("user_id", "a b&c")), + ) + .unwrap(); + assert_eq!( + url.as_str(), + "https://h.example.com/user/info?user_id=a+b%26c" + ); + let kept = management_url("https://h.example.com?token=abc", "key/info", None).unwrap(); + assert_eq!(kept.as_str(), "https://h.example.com/key/info?token=abc"); + } +} diff --git a/rust/src/providers/litellm/info.rs b/rust/src/providers/litellm/info.rs new file mode 100644 index 0000000000..c308c41ede --- /dev/null +++ b/rust/src/providers/litellm/info.rs @@ -0,0 +1,271 @@ +//! LiteLLM management-route payloads and their projection into a usage result. +//! +//! Wire shapes follow upstream `litellm.ts`: `/key/info` names the key's +//! `user_id` / `team_id`, then `/user/info` or `/team/info` supplies the +//! budgets. Returned IDs must match the key's IDs before anything is shown. + +use chrono::{DateTime, NaiveDateTime, Utc}; +use serde::Deserialize; +use serde_json::Value; + +use crate::core::{ + CostSnapshot, ProviderError, ProviderFetchResult, RateWindow, SubscriptionMetadata, + UsageSnapshot, +}; + +#[derive(Deserialize)] +pub(super) struct KeyInfoResponse { + info: KeyInfo, +} + +#[derive(Deserialize)] +struct KeyInfo { + user_id: Option, + team_id: Option, + expires: Option, +} + +#[derive(Deserialize)] +pub(super) struct UserInfoResponse { + user_id: Option, + user_info: UserInfo, + teams: Option>, +} + +#[derive(Deserialize)] +struct UserInfo { + user_id: Option, + user_email: Option, + user_alias: Option, + spend: Option, + max_budget: Option, + budget_reset_at: Option, + metadata: Option, +} + +#[derive(Deserialize)] +struct UserMetadata { + preferred_username: Option, +} + +#[derive(Deserialize)] +pub(super) struct TeamInfoResponse { + team_id: Option, + team_info: Budget, +} + +#[derive(Deserialize)] +struct Budget { + team_id: Option, + team_alias: Option, + spend: Option, + max_budget: Option, + budget_reset_at: Option, +} + +/// Identity and routing data read from `/key/info`. +pub(super) struct KeyBinding { + pub user_id: Option, + pub team_id: Option, + expires: Option>, +} + +/// A spend/budget pair with its optional reset instant. +struct Spend { + spend: f64, + limit: Option, + reset: Option>, +} + +impl Spend { + fn budget(&self) -> Option { + self.limit.filter(|limit| *limit > 0.0) + } + + fn window(&self, label: Option<&str>) -> Option { + let limit = self.budget()?; + let mut window = RateWindow::new(self.spend / limit * 100.0); + window.resets_at = self.reset; + let detail = format!("${:.2} / ${limit:.2}", self.spend); + window.reset_description = Some(match label { + Some(label) => format!("{label}: {detail}"), + None => detail, + }); + Some(window) + } +} + +struct TeamBudget { + alias: Option, + spend: Spend, +} + +impl TeamBudget { + fn from_wire(budget: &Budget) -> Self { + Self { + alias: budget.team_alias.clone(), + spend: Spend { + spend: budget.spend.unwrap_or(0.0), + limit: budget.max_budget, + reset: parse_date(budget.budget_reset_at.as_deref()), + }, + } + } + + fn window(&self) -> Option { + let label = match &self.alias { + Some(alias) => format!("Team {alias}"), + None => "Team".to_string(), + }; + self.spend.window(Some(&label)) + } +} + +pub(super) fn parse_error(message: impl std::fmt::Display) -> ProviderError { + ProviderError::Parse(format!("LiteLLM parse error: {message}")) +} + +pub(super) fn bind_key(response: KeyInfoResponse) -> Result { + let info = response.info; + let user_id = nonempty(info.user_id); + let team_id = nonempty(info.team_id); + if user_id.is_none() && team_id.is_none() { + return Err(parse_error( + "LiteLLM key info did not include a user_id or team_id.", + )); + } + Ok(KeyBinding { + user_id, + team_id, + expires: parse_date(info.expires.as_deref()), + }) +} + +/// Project a user-bound key: personal budget plus the key's matching team. +pub(super) fn result_from_user( + key: &KeyBinding, + user_id: &str, + response: UserInfoResponse, +) -> Result { + let user = response.user_info; + let response_id = user.user_id.as_deref().or(response.user_id.as_deref()); + if response_id.is_some_and(|id| id != user_id) { + return Err(parse_error("user_id did not match /key/info")); + } + let preferred = user + .metadata + .and_then(|metadata| metadata.preferred_username) + .and_then(|value| value.as_str().map(str::to_owned)); + let email = nonempty(user.user_email) + .or_else(|| nonempty(user.user_alias)) + .or_else(|| nonempty(preferred)); + let mut team = None; + for wire in response.teams.unwrap_or_default() { + let id = wire + .team_id + .as_deref() + .ok_or_else(|| parse_error("missing team_id"))?; + if team.is_none() && key.team_id.as_deref() == Some(id) { + team = Some(TeamBudget::from_wire(&wire)); + } + } + let personal = Spend { + spend: user.spend.unwrap_or(0.0), + limit: user.max_budget, + reset: parse_date(user.budget_reset_at.as_deref()), + }; + let primary = personal + .window(None) + .unwrap_or_else(|| RateWindow::new(0.0)); + let mut snapshot = UsageSnapshot::new(primary); + if let Some(email) = email { + snapshot = snapshot.with_email(email); + } + if let Some(team) = &team { + if let Some(alias) = &team.alias { + snapshot = snapshot.with_organization(alias); + } + if let Some(window) = team.window() { + snapshot = snapshot.with_extra_rate_window("team", "Team budget", window); + } + } + Ok(finish(snapshot, key, &personal, "Personal")) +} + +/// Project a team-only key: the team budget is the sole usage window. +pub(super) fn result_from_team( + key: &KeyBinding, + team_id: &str, + response: TeamInfoResponse, +) -> Result { + let response_id = response + .team_info + .team_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()) + .or(response + .team_id + .as_deref() + .filter(|id| !id.trim().is_empty())); + if response_id.is_some_and(|id| id != team_id) { + return Err(parse_error("team_id did not match /key/info")); + } + let team = TeamBudget::from_wire(&response.team_info); + let window = team.window().unwrap_or_else(|| RateWindow::new(0.0)); + let mut snapshot = UsageSnapshot::new(window).with_primary_label("Team budget"); + if let Some(alias) = &team.alias { + snapshot = snapshot.with_organization(alias); + } + Ok(finish(snapshot, key, &team.spend, "Team")) +} + +fn finish( + snapshot: UsageSnapshot, + key: &KeyBinding, + spend: &Spend, + scope: &str, +) -> ProviderFetchResult { + let mut snapshot = snapshot.with_login_method("api"); + if key.expires.is_some() { + snapshot = + snapshot.with_subscription(Some(SubscriptionMetadata::new(None, key.expires, None))); + } + let mut result = ProviderFetchResult::new(snapshot, "api"); + let limit = spend.budget(); + if spend.spend > 0.0 || limit.is_some() { + let kind = if limit.is_some() { "budget" } else { "spend" }; + let mut cost = CostSnapshot::new(spend.spend, "USD", format!("{scope} {kind}")); + if let Some(limit) = limit { + cost = cost.with_limit(limit); + } + if let Some(reset) = spend.reset { + cost = cost.with_resets_at(reset); + } + result = result.with_cost(cost); + } + result +} + +fn nonempty(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Parse an ISO-8601 instant; naive timestamps are read as UTC and anything +/// unparseable is dropped, matching upstream's tolerant `date()` helper. +fn parse_date(value: Option<&str>) -> Option> { + let raw = value?.trim(); + if raw.is_empty() { + return None; + } + DateTime::parse_from_rfc3339(raw) + .map(|date| date.with_timezone(&Utc)) + .ok() + .or_else(|| { + NaiveDateTime::parse_from_str(raw, "%Y-%m-%dT%H:%M:%S%.f") + .ok() + .map(|date| date.and_utc()) + }) +} diff --git a/rust/src/providers/litellm/mod.rs b/rust/src/providers/litellm/mod.rs index d27a573668..c4e479a67a 100644 --- a/rust/src/providers/litellm/mod.rs +++ b/rust/src/providers/litellm/mod.rs @@ -1,13 +1,27 @@ use async_trait::async_trait; -use reqwest::{Client, Url}; -use serde_json::Value; +use reqwest::{Client, StatusCode, Url}; +use serde::de::DeserializeOwned; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, + SourceMode, +}; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +mod endpoint; +mod info; +#[cfg(test)] +mod tests; + +use endpoint::management_url; +pub(crate) use endpoint::validated_base_url; +use info::{ + KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, parse_error, result_from_team, + result_from_user, }; const CREDENTIAL_TARGET: &str = "codexbar-litellm"; +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; pub struct LiteLLMProvider { metadata: ProviderMetadata, @@ -38,6 +52,40 @@ impl LiteLLMProvider { } } +impl LiteLLMProvider { + async fn get_json(&self, url: Url, key: &str) -> Result { + let route = url.path().to_string(); + let response = self + .client + .get(url) + .bearer_auth(key) + .header("Accept", "application/json") + .send() + .await?; + let status = response.status(); + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + if status == StatusCode::TOO_MANY_REQUESTS { + return Err(ProviderError::Other( + "LiteLLM rate limited the request (HTTP 429).".into(), + )); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "LiteLLM {route} returned status {status}" + ))); + } + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => parse_error("response too large"), + BoundedBodyError::Read(error) => ProviderError::Network(error), + })?; + serde_json::from_slice(&body).map_err(|e| parse_error(format!("{route}: {e}"))) + } +} + impl Default for LiteLLMProvider { fn default() -> Self { Self::new() @@ -58,28 +106,23 @@ impl Provider for LiteLLMProvider { match ctx.source_mode { SourceMode::Auto | SourceMode::OAuth => { let (base, key) = resolve_base_and_key(ctx)?; - let response = self - .client - .get(management_url(&base, "key/info")?) - .bearer_auth(key) - .header("Accept", "application/json") - .send() + let key_info: KeyInfoResponse = self + .get_json(management_url(&base, "key/info", None)?, &key) .await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); + let binding = bind_key(key_info)?; + if let Some(user_id) = binding.user_id.as_deref() { + let url = management_url(&base, "user/info", Some(("user_id", user_id)))?; + let response: UserInfoResponse = self.get_json(url, &key).await?; + result_from_user(&binding, user_id, response) + } else if let Some(team_id) = binding.team_id.as_deref() { + let url = management_url(&base, "team/info", Some(("team_id", team_id)))?; + let response: TeamInfoResponse = self.get_json(url, &key).await?; + result_from_team(&binding, team_id, response) + } else { + Err(parse_error( + "LiteLLM key info did not include a user_id or team_id.", + )) } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "LiteLLM key/info returned status {}", - response.status() - ))); - } - let value: Value = response.json().await.map_err(|e| { - ProviderError::Parse(format!("Failed to parse LiteLLM key/info: {e}")) - })?; - Ok(result_from_key_info(&value)) } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -117,131 +160,3 @@ fn resolve_base_and_key(ctx: &FetchContext) -> Result<(String, String), Provider })?; Ok((base, key)) } - -fn management_url(base: &str, path: &str) -> Result { - let mut url = crate::providers::validated_https_url(base, "LiteLLM base")?; - if url.path().trim_end_matches('/').ends_with("/v1") { - let stripped = url - .path() - .trim_end_matches('/') - .trim_end_matches("/v1") - .to_string(); - url.set_path(&stripped); - } - url.join(path) - .map_err(|e| ProviderError::Other(format!("Invalid LiteLLM URL: {e}"))) -} - -fn result_from_key_info(value: &Value) -> ProviderFetchResult { - let root = value - .get("info") - .or_else(|| value.get("key")) - .unwrap_or(value); - let spend = number(root, &["spend", "spend_usd", "spendUSD"]).unwrap_or(0.0); - let limit = number(root, &["max_budget", "maxBudget", "budget", "limit"]); - let percent = limit - .filter(|v| *v > 0.0) - .map_or(0.0, |limit| spend / limit * 100.0); - let mut primary = RateWindow::new(percent); - if let Some(limit) = limit.filter(|value| *value > 0.0) { - primary.reset_description = Some(budget_detail(spend, limit)); - } - let mut snapshot = UsageSnapshot::new(primary).with_login_method(format!("Spend ${spend:.2}")); - if let Some(team) = root.get("team_info").or_else(|| root.get("teamInfo")) - && let Some(team_spend) = number(team, &["spend", "team_spend", "teamSpend"]) - { - let team_limit = number(team, &["max_budget", "budget", "limit"]); - let team_percent = team_limit - .filter(|v| *v > 0.0) - .map_or(0.0, |limit| team_spend / limit * 100.0); - let mut team_window = RateWindow::new(team_percent); - if let Some(team_limit) = team_limit.filter(|value| *value > 0.0) { - let alias = string(team, &["team_alias", "teamAlias", "alias"]) - .map(|value| format!("Team {value}: ")) - .unwrap_or_default(); - team_window.reset_description = - Some(format!("{alias}{}", budget_detail(team_spend, team_limit))); - } - snapshot = snapshot.with_extra_rate_window("team", "Team budget", team_window); - } - let mut result = ProviderFetchResult::new(snapshot, "api"); - if spend > 0.0 { - let mut cost = CostSnapshot::new(spend, "USD", "Spend"); - if let Some(limit) = limit { - cost = cost.with_limit(limit); - } - result = result.with_cost(cost); - } - result -} - -fn number(value: &Value, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_f64)) -} - -fn string(value: &Value, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) -} - -fn budget_detail(spend: f64, budget: f64) -> String { - format!("${spend:.2} / ${budget:.2}") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_spend_budget() { - let result = - result_from_key_info(&serde_json::json!({"info":{"spend":25.0,"max_budget":100.0}})); - assert_eq!(result.usage.primary.used_percent, 25.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("$25.00 / $100.00") - ); - } - - #[test] - fn preserves_team_budget_detail_with_alias() { - let result = result_from_key_info(&serde_json::json!({ - "info": { - "team_info": { - "team_alias": "Platform", - "spend": 70.0, - "max_budget": 1000.0 - } - } - })); - assert_eq!(result.usage.extra_rate_windows.len(), 1); - assert_eq!( - result.usage.extra_rate_windows[0] - .window - .reset_description - .as_deref(), - Some("Team Platform: $70.00 / $1000.00") - ); - } - - #[test] - fn saved_base_url_uses_only_app_saved_key() { - let mut ctx = FetchContext { - workspace_id: Some("https://litellm.example.com".to_string()), - ..Default::default() - }; - assert!(matches!( - resolve_base_and_key(&ctx), - Err(ProviderError::AuthRequired) - )); - - ctx.api_key = Some("sk-app".to_string()); - let (base, key) = resolve_base_and_key(&ctx).unwrap(); - assert_eq!(base, "https://litellm.example.com"); - assert_eq!(key, "sk-app"); - } -} diff --git a/rust/src/providers/litellm/tests.rs b/rust/src/providers/litellm/tests.rs new file mode 100644 index 0000000000..3cd74e01ae --- /dev/null +++ b/rust/src/providers/litellm/tests.rs @@ -0,0 +1,278 @@ +use serde_json::{Value, json}; + +use super::info::{ + KeyBinding, KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, result_from_team, + result_from_user, +}; +use super::*; + +fn binding(info: Value) -> KeyBinding { + bind_key(serde_json::from_value::(json!({ "info": info })).unwrap()).unwrap() +} + +fn user_result(key: &KeyBinding, body: Value) -> Result { + let user_id = key.user_id.clone().unwrap(); + result_from_user( + key, + &user_id, + serde_json::from_value::(body).unwrap(), + ) +} + +fn team_result(key: &KeyBinding, body: Value) -> Result { + let team_id = key.team_id.clone().unwrap(); + result_from_team( + key, + &team_id, + serde_json::from_value::(body).unwrap(), + ) +} + +fn assert_parse_error(result: Result, expected: &str) { + match result { + Err(ProviderError::Parse(message)) => { + assert!(message.contains(expected), "unexpected message: {message}") + } + Err(other) => panic!("expected parse error, got {other}"), + Ok(_) => panic!("expected parse error"), + } +} + +#[test] +fn key_info_without_user_or_team_id_fails() { + let response: KeyInfoResponse = + serde_json::from_value(json!({"info": {"user_id": " ", "spend": 1.0}})).unwrap(); + match bind_key(response) { + Err(ProviderError::Parse(message)) => assert!( + message.contains("LiteLLM key info did not include a user_id or team_id."), + "unexpected message: {message}" + ), + _ => panic!("expected parse error"), + } +} + +#[test] +fn key_info_requires_info_object() { + assert!(serde_json::from_value::(json!({"user_id": "u"})).is_err()); +} + +#[test] +fn personal_budget_is_primary_with_identity() { + let key = binding(json!({"user_id": "user-1", "expires": "2026-12-31T00:00:00Z"})); + let result = user_result( + &key, + json!({ + "user_id": "user-1", + "user_info": { + "user_id": "user-1", + "user_email": "dev@example.com", + "spend": 25.0, + "max_budget": 100.0, + "budget_reset_at": "2026-10-01T00:00:00Z" + }, + "teams": [] + }), + ) + .unwrap(); + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("$25.00 / $100.00") + ); + assert!(result.usage.primary.resets_at.is_some()); + assert_eq!( + result.usage.account_email.as_deref(), + Some("dev@example.com") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("api")); + assert!(result.usage.extra_rate_windows.is_empty()); + assert!( + result + .usage + .subscription + .as_ref() + .is_some_and(|sub| sub.expires_at.is_some()) + ); + let cost = result.cost.expect("personal cost"); + assert_eq!(cost.used, 25.0); + assert_eq!(cost.limit, Some(100.0)); + assert_eq!(cost.period, "Personal budget"); +} + +#[test] +fn identity_falls_back_to_alias_then_preferred_username() { + let key = binding(json!({"user_id": "user-1"})); + let alias = user_result( + &key, + json!({"user_info": {"user_alias": "alias", "metadata": {"preferred_username": "pref"}}}), + ) + .unwrap(); + assert_eq!(alias.usage.account_email.as_deref(), Some("alias")); + let pref = user_result( + &key, + json!({"user_info": {"user_email": " ", "metadata": {"preferred_username": "pref"}}}), + ) + .unwrap(); + assert_eq!(pref.usage.account_email.as_deref(), Some("pref")); +} + +#[test] +fn matching_team_budget_is_a_separate_row() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-b"})); + let result = user_result( + &key, + json!({ + "user_info": {"user_id": "user-1", "spend": 3.0}, + "teams": [ + {"team_id": "team-a", "team_alias": "Other", "spend": 1.0, "max_budget": 10.0}, + {"team_id": "team-b", "team_alias": "Platform", "spend": 70.0, "max_budget": 1000.0} + ] + }), + ) + .unwrap(); + assert_eq!(result.usage.extra_rate_windows.len(), 1); + let team = &result.usage.extra_rate_windows[0].window; + assert!((team.used_percent - 7.0).abs() < 1e-9); + assert_eq!( + team.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1000.00") + ); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Platform") + ); + let cost = result.cost.expect("spend-only cost"); + assert_eq!(cost.period, "Personal spend"); + assert_eq!(cost.limit, None); +} + +#[test] +fn team_without_a_matching_entry_is_omitted() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-x"})); + let result = user_result( + &key, + json!({ + "user_info": {"spend": 1.0}, + "teams": [{"team_id": "team-a", "spend": 1.0, "max_budget": 10.0}] + }), + ) + .unwrap(); + assert!(result.usage.extra_rate_windows.is_empty()); + assert_eq!(result.usage.account_organization, None); +} + +#[test] +fn mismatched_user_id_is_rejected() { + let key = binding(json!({"user_id": "user-1"})); + assert_parse_error( + user_result(&key, json!({"user_info": {"user_id": "user-2"}})), + "user_id did not match /key/info", + ); + assert_parse_error( + user_result(&key, json!({"user_id": "user-2", "user_info": {}})), + "user_id did not match /key/info", + ); +} + +#[test] +fn team_entries_without_team_id_are_rejected() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-a"})); + assert_parse_error( + user_result(&key, json!({"user_info": {}, "teams": [{"spend": 1.0}]})), + "missing team_id", + ); +} + +#[test] +fn wrongly_typed_fields_fail_to_parse() { + assert!( + serde_json::from_value::(json!({"user_info": {"spend": "12"}})).is_err() + ); + assert!(serde_json::from_value::(json!({"teams": []})).is_err()); +} + +#[test] +fn team_only_key_shows_team_budget_as_sole_window() { + let key = binding(json!({"team_id": "team-a"})); + let result = team_result( + &key, + json!({ + "team_id": "team-a", + "team_info": { + "team_id": "team-a", + "team_alias": "Platform", + "spend": 70.0, + "max_budget": 1000.0, + "budget_reset_at": "2026-10-01T00:00:00" + } + }), + ) + .unwrap(); + assert!((result.usage.primary.used_percent - 7.0).abs() < 1e-9); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1000.00") + ); + assert!(result.usage.primary.resets_at.is_some()); + assert_eq!(result.usage.primary_label.as_deref(), Some("Team budget")); + assert!(result.usage.extra_rate_windows.is_empty()); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Platform") + ); + assert_eq!(result.usage.account_email, None); + assert_eq!(result.cost.expect("team cost").period, "Team budget"); +} + +#[test] +fn mismatched_team_id_is_rejected() { + let key = binding(json!({"team_id": "team-a"})); + assert_parse_error( + team_result( + &key, + json!({"team_id": "team-b", "team_info": {"spend": 1.0}}), + ), + "team_id did not match /key/info", + ); + assert_parse_error( + team_result(&key, json!({"team_info": {"team_id": "team-b"}})), + "team_id did not match /key/info", + ); +} + +#[test] +fn spend_above_budget_clamps_percent_and_zero_budget_is_spend_only() { + let key = binding(json!({"user_id": "user-1"})); + let over = user_result( + &key, + json!({"user_info": {"spend": 150.0, "max_budget": 100.0}}), + ) + .unwrap(); + assert_eq!(over.usage.primary.used_percent, 100.0); + let unbudgeted = user_result( + &key, + json!({"user_info": {"spend": 4.0, "max_budget": 0.0}}), + ) + .unwrap(); + assert_eq!(unbudgeted.usage.primary.reset_description, None); + assert_eq!(unbudgeted.cost.expect("cost").period, "Personal spend"); + let empty = user_result(&key, json!({"user_info": {}})).unwrap(); + assert!(empty.cost.is_none()); +} + +#[test] +fn saved_base_url_uses_only_app_saved_key() { + let mut ctx = FetchContext { + workspace_id: Some("https://litellm.example.com".to_string()), + ..Default::default() + }; + assert!(matches!( + resolve_base_and_key(&ctx), + Err(ProviderError::AuthRequired) + )); + + ctx.api_key = Some("sk-app".to_string()); + let (base, key) = resolve_base_and_key(&ctx).unwrap(); + assert_eq!(base, "https://litellm.example.com"); + assert_eq!(key, "sk-app"); +} diff --git a/rust/src/settings/provider_workspace.rs b/rust/src/settings/provider_workspace.rs index 338747c8ef..e0e48fb4be 100644 --- a/rust/src/settings/provider_workspace.rs +++ b/rust/src/settings/provider_workspace.rs @@ -55,12 +55,19 @@ pub fn validate_provider_workspace_value( Err("Helmcode tenant must be 'helmcode' or 'nanBuilders'".to_string()) } } - ProviderId::LiteLLM => validate_token_endpoint(trimmed, "LiteLLM base URL", |_| true), + ProviderId::LiteLLM => validate_litellm_base_url(trimmed), ProviderId::Sub2Api => validate_sub2api_base_url(trimmed), _ => Ok(trimmed.to_string()), } } +fn validate_litellm_base_url(raw: &str) -> Result { + match crate::providers::litellm::validated_base_url(raw) { + Ok(url) => Ok(url.to_string().trim_end_matches('/').to_string()), + Err(err) => Err(err.to_string()), + } +} + fn validate_sub2api_base_url(raw: &str) -> Result { match crate::providers::sub2api::validated_sub2api_base_url(raw) { Ok(url) => Ok(url.to_string().trim_end_matches('/').to_string()), @@ -195,7 +202,7 @@ mod tests { } #[test] - fn validates_token_endpoint_hosts() { + fn validates_litellm_base_url_policy() { assert_eq!( validate_provider_workspace_value( ProviderId::LiteLLM, @@ -204,12 +211,23 @@ mod tests { .unwrap(), "https://litellm.example.com/v1" ); + for value in [ + "http://127.0.0.1:4000", + "http://10.0.0.5:4000", + "http://192.168.1.4", + "http://[::1]:4000", + "http://proxy.local:4000", + "https://10.0.0.5", + ] { + assert!( + validate_provider_workspace_value(ProviderId::LiteLLM, value).is_ok(), + "rejected {value}" + ); + } for value in [ "http://litellm.example.com", + "http://8.8.8.8", "https://user@litellm.example.com", - "https://127.0.0.1", - "https://10.0.0.5", - "https://[::1]", "https://example.com%2f.evil.test", ] { assert!( From 3e63802c616e7d3e4705c70a817d5e9d5eeb3116 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:22:04 +0700 Subject: [PATCH 117/259] Port upstream 0.61.0: OpenRouter diagnostics and Activity summary --- rust/src/providers/openrouter/activity.rs | 107 +++++++++- rust/src/providers/openrouter/diagnostics.rs | 190 +++++++++++++++++ rust/src/providers/openrouter/mod.rs | 203 ++++++++++--------- rust/src/providers/openrouter/tests.rs | 164 +++++++++++++++ 4 files changed, 560 insertions(+), 104 deletions(-) create mode 100644 rust/src/providers/openrouter/diagnostics.rs diff --git a/rust/src/providers/openrouter/activity.rs b/rust/src/providers/openrouter/activity.rs index 9fcd2bec9f..2ee7dd79db 100644 --- a/rust/src/providers/openrouter/activity.rs +++ b/rust/src/providers/openrouter/activity.rs @@ -1,4 +1,4 @@ -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, HashMap, HashSet}; use chrono::{DateTime, Duration, Utc}; use serde_json::Value; @@ -10,15 +10,49 @@ const MAX_ACTIVITY_ROWS: usize = 20_000; const MAX_DISTINCT_ROWS: usize = 10_000; const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991; +/// Aggregate of the deduplicated, in-window Activity rows (upstream +/// `activityDetails`: Tokens = prompt + completion, Requests, distinct Models). +/// Reasoning tokens are validated but never added to `tokens` a second time. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct ActivitySummary { + pub(super) tokens: u64, + pub(super) requests: u64, + pub(super) models: usize, +} + +#[derive(Debug, Clone)] +pub(super) struct ActivityReport { + pub(super) cost: CostSnapshot, + pub(super) summary: ActivitySummary, +} + +/// Add `amount` to `total`, rejecting aggregates beyond the JS safe-integer +/// range like upstream (`Number.isSafeInteger`). +fn checked_aggregate(total: &mut u64, amount: u64) -> Result<(), ProviderError> { + *total = total + .checked_add(amount) + .filter(|sum| *sum <= MAX_SAFE_INTEGER) + .ok_or_else(|| { + ProviderError::Parse( + "OpenRouter Activity aggregate must be within the safe integer range".into(), + ) + })?; + Ok(()) +} + pub(super) fn parse_activity_cost( payloads: &[Value], now: DateTime, -) -> Result { +) -> Result { let latest_completed = now.date_naive() - Duration::days(1); let cutoff = latest_completed - Duration::days(29); let mut seen: HashMap = HashMap::new(); let mut daily: BTreeMap = BTreeMap::new(); let mut total = 0.0; + let mut tokens = 0u64; + let mut reasoning_tokens = 0u64; + let mut requests_total = 0u64; + let mut models: HashSet = HashSet::new(); let mut rows_seen = 0usize; for payload in payloads { @@ -138,6 +172,12 @@ pub(super) fn parse_activity_cost( "OpenRouter activity.data exceeds 10000 distinct rows".into(), )); } + checked_aggregate(&mut tokens, prompt + completion)?; + checked_aggregate(&mut reasoning_tokens, reasoning)?; + checked_aggregate(&mut requests_total, requests)?; + if !model.is_empty() { + models.insert(model.to_string()); + } total += cost; *daily.entry(day.to_string()).or_default() += cost; } @@ -148,14 +188,22 @@ pub(super) fn parse_activity_cost( "OpenRouter Activity spend overflowed".into(), )); } - Ok(CostSnapshot::new(total, "USD", "Last 30 days (UTC)") + let cost = CostSnapshot::new(total, "USD", "Last 30 days (UTC)") .with_daily( daily .into_iter() .map(|(day, amount)| CostDailyPoint { day, amount }) .collect(), ) - .always_visible()) + .always_visible(); + Ok(ActivityReport { + cost, + summary: ActivitySummary { + tokens, + requests: requests_total, + models: models.len(), + }, + }) } fn normalize_activity_day(raw: &str) -> Option<&str> { @@ -241,7 +289,17 @@ mod tests { let latest_completed = serde_json::json!({"data":[ {"date":"2026-08-21","model":"m1","prompt_tokens":10,"completion_tokens":5,"reasoning_tokens":2,"requests":1,"usage":1.25,"byok_usage_inference":0.25} ]}); - let cost = parse_activity_cost(&[history, latest_completed], now()).unwrap(); + let report = parse_activity_cost(&[history, latest_completed], now()).unwrap(); + let cost = report.cost; + // The duplicated latest-completed row is counted once. + assert_eq!( + report.summary, + ActivitySummary { + tokens: 22, + requests: 2, + models: 2 + } + ); assert!((cost.used - 2.0).abs() < 1e-12); assert_eq!(cost.daily.len(), 2); assert_eq!(cost.period, "Last 30 days (UTC)"); @@ -254,10 +312,12 @@ mod tests { "completion_tokens":2,"reasoning_tokens":8,"requests":1,"usage":1.0} ]}); - let cost = parse_activity_cost(&[payload], now()).unwrap(); + let report = parse_activity_cost(&[payload], now()).unwrap(); - assert_eq!(cost.used, 1.0); - assert_eq!(cost.daily.len(), 1); + assert_eq!(report.cost.used, 1.0); + assert_eq!(report.cost.daily.len(), 1); + // Tokens stay prompt + completion; reasoning is not added again. + assert_eq!(report.summary.tokens, 12); } #[test] @@ -277,7 +337,7 @@ mod tests { {"date":"2026-07-22","model":"old","prompt_tokens":10,"completion_tokens":5,"requests":1,"usage":99.0}, {"date":"2026-07-23","model":"in","prompt_tokens":10,"completion_tokens":5,"requests":1,"usage":1.0} ]}); - let cost = parse_activity_cost(&[payload], now()).unwrap(); + let cost = parse_activity_cost(&[payload], now()).unwrap().cost; assert_eq!(cost.used, 1.0); } @@ -287,7 +347,7 @@ mod tests { let payload = serde_json::json!({"data":[ {"date":date,"model":"m","prompt_tokens":10,"completion_tokens":5,"requests":1,"usage":1.0} ]}); - let cost = parse_activity_cost(&[payload], now()).unwrap(); + let cost = parse_activity_cost(&[payload], now()).unwrap().cost; assert_eq!(cost.daily.len(), 1); assert_eq!(cost.daily[0].day, "2026-08-21"); } @@ -314,4 +374,31 @@ mod tests { assert!(error.to_string().contains("completed UTC day")); } + + #[test] + fn rows_without_a_model_do_not_count_as_models() { + let payload = serde_json::json!({"data":[ + {"date":"2026-08-21","prompt_tokens":1,"completion_tokens":1,"requests":1,"usage":0.1}, + {"date":"2026-08-21","model":" ","endpoint_id":"e","prompt_tokens":1,"completion_tokens":1,"requests":1,"usage":0.1}, + {"date":"2026-08-21","model_permaslug":"a/b","prompt_tokens":1,"completion_tokens":1,"requests":1,"usage":0.1} + ]}); + + let summary = parse_activity_cost(&[payload], now()).unwrap().summary; + + assert_eq!(summary.models, 1); + assert_eq!(summary.requests, 3); + } + + #[test] + fn aggregate_beyond_the_safe_integer_range_is_rejected() { + let big = MAX_SAFE_INTEGER / 2 + 1; + let payload = serde_json::json!({"data":[ + {"date":"2026-08-21","model":"a","prompt_tokens":big,"completion_tokens":0,"requests":1,"usage":0.1}, + {"date":"2026-08-20","model":"b","prompt_tokens":big,"completion_tokens":0,"requests":1,"usage":0.1} + ]}); + + let error = parse_activity_cost(&[payload], now()).unwrap_err(); + + assert!(error.to_string().contains("safe integer")); + } } diff --git a/rust/src/providers/openrouter/diagnostics.rs b/rust/src/providers/openrouter/diagnostics.rs new file mode 100644 index 0000000000..850d7370ac --- /dev/null +++ b/rust/src/providers/openrouter/diagnostics.rs @@ -0,0 +1,190 @@ +//! Optional-request diagnostics and display rows for the OpenRouter provider +//! (upstream `openrouter.js` at v0.61.0: safe degradation reasons plus the +//! Credits, API key, and Activity detail sections). +//! +//! Reasons are fixed strings chosen from the failure class, never raw +//! response bodies, headers, or keys. + +use reqwest::StatusCode; + +use super::activity::ActivitySummary; +use super::{CreditsData, KeyData, OpenRouterProvider}; +use crate::core::{ProviderDisplayDetail, ProviderError}; + +const UNAVAILABLE: &str = "Unavailable right now"; +pub(super) const ACTIVITY_NOT_CONFIGURED: &str = "Management API key not configured"; +pub(super) const ACTIVITY_KEY_REQUIRED: &str = "Management API key required"; + +/// One optional request that did not produce usable data: the typed error the +/// refresh pipeline acts on, plus the safe reason shown to the user. +#[derive(Debug)] +pub(super) struct Degraded { + pub(super) error: ProviderError, + pub(super) reason: String, +} + +impl Degraded { + /// A non-success HTTP status. Statuses in `auth_statuses` are typed as + /// [`ProviderError::AuthRequired`]; the reason always names the status. + pub(super) fn http(label: &str, status: StatusCode, auth_statuses: &[StatusCode]) -> Self { + let error = if auth_statuses.contains(&status) { + ProviderError::AuthRequired + } else { + ProviderError::Other(format!("OpenRouter {label} request returned HTTP {status}")) + }; + Self { + error, + reason: format!("Request returned HTTP {}", status.as_u16()), + } + } + + /// A successful HTTP response whose body failed parsing or validation. + pub(super) fn invalid(error: ProviderError) -> Self { + Self { + error, + reason: "Response was invalid".to_string(), + } + } + + /// A response body that failed to decode, unless the body read itself + /// timed out (then it is a timeout, not an invalid response). + pub(super) fn body(label: &str, error: reqwest::Error) -> Self { + if error.is_timeout() { + return Self::from(error); + } + Self::invalid(ProviderError::Parse(format!( + "OpenRouter {label} response was invalid: {error}" + ))) + } + + pub(super) fn with_reason(mut self, reason: &str) -> Self { + self.reason = reason.to_string(); + self + } +} + +impl From for Degraded { + fn from(error: reqwest::Error) -> Self { + let reason = if error.is_timeout() { + "Request timed out" + } else { + "Request failed" + }; + Self { + error: ProviderError::Network(error), + reason: reason.to_string(), + } + } +} + +/// What each optional source reported, with the safe reason when it did not. +pub(super) struct Observations<'a> { + pub(super) credits: &'a Result, + pub(super) key: &'a Result, + pub(super) activity: &'a Result, +} + +pub(super) fn build_display_details(observed: &Observations<'_>) -> Vec { + let mut rows = Vec::new(); + match observed.credits { + Ok(credits) => { + rows.push(row( + "credits-remaining", + "Credits remaining", + currency(credits.balance()), + )); + rows.push(row( + "credits-used", + "Credits used", + currency(credits.total_usage), + )); + rows.push(row( + "credits-total", + "Credits total added", + currency(credits.total_credits), + )); + } + Err(reason) => rows.push(unavailable("credits-balance", "Credits balance", reason)), + } + match observed.key { + Ok(key) => rows.extend(key_rows(key)), + Err(reason) => rows.push(unavailable("key-limit", "API key limit", reason)), + } + match observed.activity { + Ok(summary) => { + rows.push( + ProviderDisplayDetail::new( + "activity-tokens", + "Activity tokens", + summary.tokens.to_string(), + ) + .and_then(|row| row.with_secondary_value("Last 30 completed UTC days")), + ); + rows.push(row( + "activity-requests", + "Activity requests", + summary.requests.to_string(), + )); + rows.push(row( + "activity-models", + "Activity models", + summary.models.to_string(), + )); + } + Err(reason) => rows.push(unavailable( + "spend-history", + "Spend history (last 30 days)", + reason, + )), + } + rows.into_iter().flatten().collect() +} + +fn key_rows(key: &KeyData) -> Vec> { + let mut rows = Vec::new(); + match key.limit.filter(|limit| *limit > 0.0) { + Some(limit) => { + rows.push( + ProviderDisplayDetail::new("key-limit", "API key limit", currency(limit)) + .and_then(|row| row.with_secondary_value("Spending cap, not balance")), + ); + if let Some((_, used, limit)) = OpenRouterProvider::key_quota_metrics(key) { + rows.push(row( + "key-remaining", + "API key remaining", + currency((limit - used).max(0.0)), + )); + } + if let Some(usage) = key.usage { + rows.push(row("key-used", "API key used", currency(usage))); + } + } + None => rows.push(row( + "key-limit", + "API key limit", + "No limit configured".to_string(), + )), + } + if let Some(window) = key + .limit_reset + .as_deref() + .map(str::trim) + .filter(|window| !window.is_empty()) + { + rows.push(row("key-reset-window", "Reset window", window.to_string())); + } + rows +} + +fn row(id: &str, title: &str, value: String) -> Option { + ProviderDisplayDetail::new(id, title, value) +} + +fn unavailable(id: &str, title: &str, reason: &str) -> Option { + ProviderDisplayDetail::new(id, title, UNAVAILABLE) + .and_then(|row| row.with_secondary_value(reason)) +} + +fn currency(value: f64) -> String { + format!("${:.2}", value.max(0.0)) +} diff --git a/rust/src/providers/openrouter/mod.rs b/rust/src/providers/openrouter/mod.rs index 1c15164c82..e0186bfcfa 100755 --- a/rust/src/providers/openrouter/mod.rs +++ b/rust/src/providers/openrouter/mod.rs @@ -4,6 +4,7 @@ //! Requires API key for authentication mod activity; +mod diagnostics; #[cfg(test)] #[path = "tests.rs"] @@ -18,6 +19,8 @@ use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; +use activity::ActivityReport; +use diagnostics::{Degraded, Observations}; /// OpenRouter API base URL — the bare `/api/v1` prefix, matching upstream /// (steipete/CodexBar `OpenRouterSettingsReader.apiURL`). @@ -26,17 +29,31 @@ use crate::core::{ /// The fork's original bug baked `/auth` into the base (`.../api/v1/auth`), /// which turned the credits call into `/api/v1/auth/credits` -> 404. const OPENROUTER_API_BASE: &str = "https://openrouter.ai/api/v1"; -const OPENROUTER_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); -/// Optional key-quota enrichment joins on a one-second fast deadline -/// (upstream 0.49.0 #2778) so a slow `/key` endpoint can never stall the -/// refresh; degraded enrichment is logged and skipped, never fatal. -const OPENROUTER_KEY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1); +/// Per-request deadline for `/credits`, `/key`, and Activity. Upstream +/// `openrouter.js` (v0.61.0) gives each optional request four seconds; a +/// degraded request is reported with a safe reason, never fatal on its own. +const OPENROUTER_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(4); const OPENROUTER_ACTIVITY_URL: &str = "https://openrouter.ai/api/v1/activity"; const OPENROUTER_MANAGEMENT_ENV: &str = "OPENROUTER_MANAGEMENT_API_KEY"; /// Windows Credential Manager target for OpenRouter API token const OPENROUTER_CREDENTIAL_TARGET: &str = "codexbar-openrouter"; +/// Statuses that mean the credential itself was rejected. +const AUTH_REJECTED: &[reqwest::StatusCode] = &[ + reqwest::StatusCode::UNAUTHORIZED, + reqwest::StatusCode::FORBIDDEN, +]; + +/// Snapshot one optional request for the display rows: the usable value, or +/// the safe reason it degraded. +fn observe(result: &Result, value: impl FnOnce(&T) -> U) -> Result { + match result { + Ok(ok) => Ok(value(ok)), + Err(degraded) => Err(degraded.reason.clone()), + } +} + /// OpenRouter /credits response #[derive(Debug, Clone, Deserialize)] struct CreditsResponse { @@ -205,30 +222,33 @@ impl OpenRouterProvider { async fn fetch_usage_api( &self, ctx: &FetchContext, - ) -> Result<(UsageSnapshot, Option), ProviderError> { + ) -> Result { let api_key = Self::get_api_token(ctx.api_key.as_deref())?; - let client = Self::build_client(OPENROUTER_TIMEOUT)?; + let client = Self::build_client()?; // OpenRouter can reject the account-level `/credits` request while // still returning the selected key's current spend/quota from `/key`. // Fetch both independent sources together, then let the pure resolver // choose the stable primary/secondary lanes. - let (credits_result, key_data_result) = tokio::join!( + let (credits_result, key_result) = tokio::join!( Self::fetch_credits(&client, &api_key), - Self::fetch_key_data(&api_key), + Self::fetch_key_data(&client, &api_key), ); - if let Err(error) = &credits_result { - tracing::debug!(error = %error, "OpenRouter credits endpoint degraded"); + let credits_observed = observe(&credits_result, |credits| credits.data.clone()); + let key_observed = observe(&key_result, Clone::clone); + if let Err(degraded) = &credits_result { + tracing::debug!(reason = %degraded.reason, "OpenRouter credits endpoint degraded"); } - let key_data = match key_data_result { + let key_data = match key_result { Ok(key_data) => Some(key_data), - Err(error) => { + Err(degraded) => { tracing::debug!( - error = %error, + reason = %degraded.reason, "OpenRouter key endpoint degraded; preserving independent credits data" ); None } }; + let credits_result = credits_result.map_err(|degraded| degraded.error); let fallback_cost = Self::build_uncapped_cost(key_data.as_ref(), credits_result.as_ref().ok()); let usage = Self::resolve_usage(credits_result, key_data.clone())?; @@ -240,43 +260,64 @@ impl OpenRouterProvider { .is_some_and(|key_data| key_data.is_management_key == Some(true)) .then_some(api_key.as_str()) }); - let activity_cost = match activity_key { - Some(key) => match Self::fetch_activity_cost(key).await { - Ok(cost) => Some(cost), - Err(error) => { - tracing::debug!( - error = %error, - "OpenRouter management Activity degraded; preserving credits/quota" - ); - None - } - }, + let activity_result = match activity_key { + Some(key) => Some(Self::fetch_activity(&client, key).await), + None => None, + }; + let activity_observed = match &activity_result { + Some(result) => observe(result, |report| report.summary), + None => Err(diagnostics::ACTIVITY_NOT_CONFIGURED.to_string()), + }; + let activity_cost = match activity_result { + Some(Ok(report)) => Some(report.cost), + Some(Err(degraded)) => { + tracing::debug!( + reason = %degraded.reason, + "OpenRouter management Activity degraded; preserving credits/quota" + ); + None + } None => None, }; - Ok((usage, activity_cost.or(fallback_cost))) + let mut result = ProviderFetchResult::new(usage, "api"); + if let Some(cost) = activity_cost.or(fallback_cost) { + result = result.with_cost(cost); + } + let details = diagnostics::build_display_details(&Observations { + credits: &credits_observed, + key: &key_observed, + activity: &activity_observed, + }); + for detail in details { + result = result.with_display_detail(Some(detail)); + } + Ok(result) } - async fn fetch_activity_cost(management_key: &str) -> Result { - let client = Self::build_client(OPENROUTER_KEY_TIMEOUT)?; + async fn fetch_activity( + client: &reqwest::Client, + management_key: &str, + ) -> Result { let now = Utc::now(); let latest_completed = (now.date_naive() - chrono::Duration::days(1)) .format("%Y-%m-%d") .to_string(); let (history_result, latest_completed_result) = tokio::join!( - Self::fetch_activity_payload(&client, management_key, None), - Self::fetch_activity_payload(&client, management_key, Some(&latest_completed)), + Self::fetch_activity_payload(client, management_key, None), + Self::fetch_activity_payload(client, management_key, Some(&latest_completed)), ); let history = history_result?; let latest_completed_payload = latest_completed_result?; activity::parse_activity_cost(&[history, latest_completed_payload], now) + .map_err(Degraded::invalid) } async fn fetch_activity_payload( client: &reqwest::Client, management_key: &str, date: Option<&str>, - ) -> Result { + ) -> Result { let mut request = client .get(OPENROUTER_ACTIVITY_URL) .header("Authorization", format!("Bearer {management_key}")) @@ -285,25 +326,23 @@ impl OpenRouterProvider { request = request.query(&[("date", date)]); } let response = request.send().await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); + let status = response.status(); + if status == reqwest::StatusCode::FORBIDDEN { + return Err(Degraded::http("Activity", status, AUTH_REJECTED) + .with_reason(diagnostics::ACTIVITY_KEY_REQUIRED)); } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "OpenRouter Activity request returned HTTP {}", - response.status() - ))); + if !status.is_success() { + return Err(Degraded::http("Activity", status, AUTH_REJECTED)); } - response.json::().await.map_err(|error| { - ProviderError::Parse(format!("Invalid OpenRouter Activity response: {error}")) - }) + response + .json::() + .await + .map_err(|error| Degraded::body("Activity", error)) } - fn build_client(timeout: std::time::Duration) -> Result { + fn build_client() -> Result { crate::core::credentialed_http_client_builder() - .timeout(timeout) + .timeout(OPENROUTER_REQUEST_TIMEOUT) .build() .map_err(|e| ProviderError::Other(e.to_string())) } @@ -311,7 +350,7 @@ impl OpenRouterProvider { async fn fetch_credits( client: &reqwest::Client, api_key: &str, - ) -> Result { + ) -> Result { let credits_url = format!("{}/credits", OPENROUTER_API_BASE); let resp = client .get(&credits_url) @@ -320,21 +359,19 @@ impl OpenRouterProvider { .send() .await?; - if resp.status() == reqwest::StatusCode::UNAUTHORIZED { - return Err(ProviderError::AuthRequired); - } - if !resp.status().is_success() { - return Err(ProviderError::Other(format!( - "OpenRouter credits request returned HTTP {}", - resp.status() - ))); + return Err(Degraded::http( + "credits", + resp.status(), + &[reqwest::StatusCode::UNAUTHORIZED], + )); } - let response = resp.json::().await.map_err(|error| { - ProviderError::Parse(format!("OpenRouter credits response was invalid: {error}")) - })?; - response.data.validate()?; + let response = resp + .json::() + .await + .map_err(|error| Degraded::body("credits", error))?; + response.data.validate().map_err(Degraded::invalid)?; Ok(response) } @@ -413,40 +450,25 @@ impl OpenRouterProvider { Some(usage) } - async fn fetch_key_data(api_key: &str) -> Result { - let key_client = Self::build_client(OPENROUTER_KEY_TIMEOUT)?; - let key_resp = Self::send_key_request(&key_client, api_key).await?; - - if key_resp.status() == reqwest::StatusCode::UNAUTHORIZED - || key_resp.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); - } - if !key_resp.status().is_success() { - return Err(ProviderError::Other(format!( - "OpenRouter key request returned HTTP {}", - key_resp.status() - ))); - } - - let response = key_resp.json::().await.map_err(|error| { - ProviderError::Parse(format!("OpenRouter key response was invalid: {error}")) - })?; - response.data.validate()?; - Ok(response.data) - } - - async fn send_key_request( - client: &reqwest::Client, - api_key: &str, - ) -> Result { + async fn fetch_key_data(client: &reqwest::Client, api_key: &str) -> Result { let key_url = format!("{}/key", OPENROUTER_API_BASE); - client + let key_resp = client .get(&key_url) .header("Authorization", format!("Bearer {}", api_key)) .header("Accept", "application/json") .send() + .await?; + + if !key_resp.status().is_success() { + return Err(Degraded::http("key", key_resp.status(), AUTH_REJECTED)); + } + + let response = key_resp + .json::() .await + .map_err(|error| Degraded::body("key", error))?; + response.data.validate().map_err(Degraded::invalid)?; + Ok(response.data) } fn apply_key_lanes(usage: &mut UsageSnapshot, key_data: &KeyData, quota_suffix: &str) { @@ -565,14 +587,7 @@ impl Provider for OpenRouterProvider { tracing::debug!("Fetching OpenRouter usage"); match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => { - let (usage, cost) = self.fetch_usage_api(ctx).await?; - let mut result = ProviderFetchResult::new(usage, "api"); - if let Some(cost) = cost { - result = result.with_cost(cost); - } - Ok(result) - } + SourceMode::Auto | SourceMode::OAuth => self.fetch_usage_api(ctx).await, SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) } diff --git a/rust/src/providers/openrouter/tests.rs b/rust/src/providers/openrouter/tests.rs index a7728bde20..bc8e08ae30 100644 --- a/rust/src/providers/openrouter/tests.rs +++ b/rust/src/providers/openrouter/tests.rs @@ -387,3 +387,167 @@ fn parsed_key_wire_fields_decode() { assert_eq!(parsed.data.limit_remaining, Some(12.5)); assert_eq!(parsed.data.limit_reset.as_deref(), Some("monthly")); } + +// ── 0.61.0 (#3272, #3733): optional-request diagnostics and detail rows ── + +use super::activity::ActivitySummary; +use super::diagnostics::{ + ACTIVITY_KEY_REQUIRED, ACTIVITY_NOT_CONFIGURED, Observations, build_display_details, +}; + +type DetailRow = (String, String, Option); + +fn detail_rows( + credits: Result, + key: Result, + activity: Result, +) -> Vec { + build_display_details(&Observations { + credits: &credits, + key: &key, + activity: &activity, + }) + .iter() + .map(|row| { + ( + row.title().to_string(), + row.value().to_string(), + row.secondary_value().map(str::to_string), + ) + }) + .collect() +} + +fn row_of<'a>(rows: &'a [DetailRow], title: &str) -> &'a DetailRow { + rows.iter() + .find(|row| row.0 == title) + .unwrap_or_else(|| panic!("missing row {title}: {rows:?}")) +} + +#[test] +fn optional_request_deadline_is_four_seconds() { + assert_eq!( + OPENROUTER_REQUEST_TIMEOUT, + std::time::Duration::from_secs(4) + ); +} + +#[test] +fn successful_sources_render_credits_key_and_activity_rows() { + let rows = detail_rows( + Ok(CreditsData { + total_credits: 5.0, + total_usage: 3.1, + }), + Ok(key_data( + Some(30.0), + Some(30.0), + Some(" monthly "), + Some(0.0), + None, + None, + None, + )), + Ok(ActivitySummary { + tokens: 22, + requests: 2, + models: 2, + }), + ); + + assert_eq!(row_of(&rows, "Credits remaining").1, "$1.90"); + assert_eq!(row_of(&rows, "Credits used").1, "$3.10"); + assert_eq!(row_of(&rows, "Credits total added").1, "$5.00"); + let limit = row_of(&rows, "API key limit"); + assert_eq!(limit.1, "$30.00"); + assert_eq!(limit.2.as_deref(), Some("Spending cap, not balance")); + assert_eq!(row_of(&rows, "API key remaining").1, "$30.00"); + assert_eq!(row_of(&rows, "API key used").1, "$0.00"); + assert_eq!(row_of(&rows, "Reset window").1, "monthly"); + assert_eq!(row_of(&rows, "Activity tokens").1, "22"); + assert_eq!(row_of(&rows, "Activity requests").1, "2"); + assert_eq!(row_of(&rows, "Activity models").1, "2"); + assert!( + rows.iter() + .all(|row| row.0 != "Spend history (last 30 days)") + ); +} + +#[test] +fn uncapped_key_reports_no_limit_and_omits_remaining() { + let rows = detail_rows( + Err("Request failed".into()), + Ok(key_data(None, None, None, Some(1.0), None, None, None)), + Err(ACTIVITY_NOT_CONFIGURED.into()), + ); + + assert_eq!(row_of(&rows, "API key limit").1, "No limit configured"); + assert!(rows.iter().all(|row| row.0 != "API key remaining")); + assert!(rows.iter().all(|row| row.0 != "Reset window")); +} + +#[test] +fn degraded_sources_keep_safe_reasons_beside_the_unavailable_marker() { + let rows = detail_rows( + Err("Request returned HTTP 503".into()), + Err("Request timed out".into()), + Err(ACTIVITY_KEY_REQUIRED.into()), + ); + + for (title, reason) in [ + ("Credits balance", "Request returned HTTP 503"), + ("API key limit", "Request timed out"), + ("Spend history (last 30 days)", ACTIVITY_KEY_REQUIRED), + ] { + let row = row_of(&rows, title); + assert_eq!(row.1, "Unavailable right now"); + assert_eq!(row.2.as_deref(), Some(reason)); + } + assert_eq!(rows.len(), 3); +} + +#[test] +fn http_failures_keep_their_status_and_auth_typing() { + use reqwest::StatusCode; + + let rejected = Degraded::http("key", StatusCode::FORBIDDEN, AUTH_REJECTED); + assert!(matches!(rejected.error, ProviderError::AuthRequired)); + assert_eq!(rejected.reason, "Request returned HTTP 403"); + + // Credits only treats 401 as a rejected credential. + let unavailable = Degraded::http( + "credits", + StatusCode::SERVICE_UNAVAILABLE, + &[StatusCode::UNAUTHORIZED], + ); + assert!(matches!(unavailable.error, ProviderError::Other(_))); + assert_eq!(unavailable.reason, "Request returned HTTP 503"); + + let activity = Degraded::http("Activity", StatusCode::FORBIDDEN, AUTH_REJECTED) + .with_reason(ACTIVITY_KEY_REQUIRED); + assert!(matches!(activity.error, ProviderError::AuthRequired)); + assert_eq!(activity.reason, "Management API key required"); +} + +#[test] +fn invalid_bodies_are_labelled_without_leaking_the_payload() { + let degraded = Degraded::invalid(ProviderError::Parse("secret-body".into())); + assert_eq!(degraded.reason, "Response was invalid"); +} + +#[tokio::test] +async fn slow_response_reports_a_timeout_and_other_transport_errors_a_failure() { + // A bound listener that never accepts or replies: the request stalls. + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/", listener.local_addr().unwrap()); + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_millis(100)) + .build() + .unwrap(); + let stalled = Degraded::from(client.get(&url).send().await.unwrap_err()); + assert_eq!(stalled.reason, "Request timed out"); + + // A non-timeout transport error (unsupported scheme) is a plain failure. + let failed = Degraded::from(client.get("ftp://127.0.0.1/").send().await.unwrap_err()); + assert_eq!(failed.reason, "Request failed"); +} From 86dd3fbb005ee2ebd14775858f5f5ff42b52297e Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:25:17 +0700 Subject: [PATCH 118/259] Port upstream 0.67.0: fetch Bedrock daily cost history for period selection --- .../src-tauri/src/commands/usage_spend.rs | 5 +- .../src/commands/usage_spend/tests.rs | 36 ++++ rust/src/providers/bedrock/daily.rs | 182 ++++++++++++++++++ rust/src/providers/bedrock/mod.rs | 166 +++++++--------- 4 files changed, 291 insertions(+), 98 deletions(-) create mode 100644 rust/src/providers/bedrock/daily.rs 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 4297905f10..3e9e2ec118 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -833,6 +833,7 @@ fn cached_spend( } else { let today = chrono::Utc::now().date_naive(); let seven_cutoff = today - chrono::Duration::days(6); + let thirty_cutoff = today - chrono::Duration::days(29); let mut seven = 0.0; let mut thirty = 0.0; let mut saw_seven = false; @@ -841,7 +842,9 @@ fn cached_spend( let Ok(day) = chrono::NaiveDate::parse_from_str(&point.day, "%Y-%m-%d") else { continue; }; - if day > today { + // Providers with long daily history (Bedrock keeps 14 months) must + // not widen the fixed 30-day column. + if day > today || day < thirty_cutoff { continue; } thirty += point.amount; diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs index 20b39eadf2..b42400eeb3 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs @@ -107,6 +107,42 @@ fn daily_period_cost_follows_the_selected_window() { ); } +#[test] +fn long_provider_daily_history_does_not_widen_the_fixed_columns() { + // Bedrock reports 14 months of daily spend; the 7d / 30d columns keep + // their windows while the selected period still sees the whole history. + let today = Utc::now().date_naive(); + let day = |ago: i64| { + serde_json::json!({ + "day": (today - chrono::Duration::days(ago)).format("%Y-%m-%d").to_string(), + "amount": 1.0, + }) + }; + let cost: super::super::bridge::CostSnapshotBridge = + serde_json::from_value(serde_json::json!({ + "used": 99.0, + "period": "Monthly", + "daily": [day(0), day(6), day(7), day(29), day(30), day(200)], + })) + .expect("cost fixture"); + let mut snapshot = ProviderUsageSnapshot::from_error( + codexbar::core::ProviderId::Bedrock, + codexbar::core::instantiate_provider(codexbar::core::ProviderId::Bedrock).metadata(), + "unused".to_string(), + codexbar::core::ProviderStateKind::Unknown, + ); + snapshot.cost = Some(cost); + + let spend = cached_spend( + Some(&snapshot), + CostReportingPeriod::AllAvailable, + Utc::now(), + ); + assert_eq!(spend.seven_day, Some(2.0)); + assert_eq!(spend.thirty_day, Some(4.0)); + assert_eq!(spend.period_cost, Some(6.0)); +} + #[test] fn pi_history_is_an_alternate_view_not_a_shared_overview_source() { assert!(!include_in_shared_overview("pi", true, true)); diff --git a/rust/src/providers/bedrock/daily.rs b/rust/src/providers/bedrock/daily.rs new file mode 100644 index 0000000000..56b388fe41 --- /dev/null +++ b/rust/src/providers/bedrock/daily.rs @@ -0,0 +1,182 @@ +//! Cost Explorer date ranges and daily Bedrock cost parsing (upstream +//! `BedrockUsageFetcher.dailyRange` / `parseDailyResponse`). + +use std::collections::BTreeMap; + +use chrono::{DateTime, Duration, Months, NaiveDate, Utc}; +use serde_json::Value; + +use super::bedrock_group_amounts; +use crate::core::CostDailyPoint; +use crate::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; + +/// Cost Explorer `TimePeriod` for daily buckets (upstream `dailyRange`). +/// +/// Cost Explorer buckets are UTC, so the month is resolved in UTC. It exposes +/// the current month plus thirteen earlier months, so `since` never reaches +/// further back than that; the exclusive end is tomorrow. An all-available +/// request passes any early `since` and gets the whole exposed range. +pub(super) fn daily_range(since: NaiveDate, now: DateTime) -> (String, String) { + let month_start = utc_month_start(now); + let earliest = month_start + .checked_sub_months(Months::new(13)) + .unwrap_or(month_start); + let tomorrow = now.date_naive() + Duration::days(1); + ( + since.max(earliest).format("%Y-%m-%d").to_string(), + tomorrow.format("%Y-%m-%d").to_string(), + ) +} + +fn utc_month_start(now: DateTime) -> NaiveDate { + CostReportingPeriod::MonthToDate + .bounds(now, CostTimeZone::Named(chrono_tz::UTC), None) + .start +} + +/// Current-month range: month to date through tomorrow (exclusive). +pub(super) fn current_month_range() -> (String, String) { + let now = Utc::now(); + daily_range(utc_month_start(now), now) +} + +/// Every month Cost Explorer exposes, through tomorrow (exclusive). +pub(super) fn all_available_range() -> (String, String) { + daily_range(NaiveDate::MIN, Utc::now()) +} + +/// Daily Bedrock spend from `GetCostAndUsage` DAILY pages. +/// +/// Days without positive Bedrock spend are omitted and the same day across +/// pages is summed, as upstream merges its daily reports. +pub(super) fn parse_daily_costs(pages: &[Value]) -> Vec { + let mut by_day: BTreeMap = BTreeMap::new(); + for result in pages + .iter() + .filter_map(|page| page.get("ResultsByTime").and_then(Value::as_array)) + .flatten() + { + let Some(day) = result + .get("TimePeriod") + .and_then(|period| period.get("Start")) + .and_then(Value::as_str) + .filter(|day| NaiveDate::parse_from_str(day, "%Y-%m-%d").is_ok()) + else { + continue; + }; + let cost: f64 = bedrock_group_amounts(result) + .filter(|amount| amount.is_finite() && *amount > 0.0) + .sum(); + if cost > 0.0 { + *by_day.entry(day.to_string()).or_default() += cost; + } + } + by_day + .into_iter() + .map(|(day, amount)| CostDailyPoint { day, amount }) + .collect() +} + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use serde_json::json; + + use super::*; + + fn utc(y: i32, m: u32, d: u32, h: u32) -> DateTime { + Utc.with_ymd_and_hms(y, m, d, h, 0, 0).single().unwrap() + } + + fn date(y: i32, m: u32, d: u32) -> NaiveDate { + NaiveDate::from_ymd_opt(y, m, d).unwrap() + } + + #[test] + fn daily_range_month_to_date_starts_at_utc_month_start() { + let now = utc(2026, 5, 15, 12); + let (start, end) = daily_range(utc_month_start(now), now); + assert_eq!(start, "2026-05-01"); + assert_eq!(end, "2026-05-16"); + } + + #[test] + fn daily_range_month_start_uses_utc_not_local_time() { + // 23:30 UTC on the last day of April is still April in Cost Explorer. + let now = Utc + .with_ymd_and_hms(2026, 4, 30, 23, 30, 0) + .single() + .unwrap(); + let (start, end) = daily_range(utc_month_start(now), now); + assert_eq!(start, "2026-04-01"); + assert_eq!(end, "2026-05-01"); + } + + #[test] + fn daily_range_all_is_capped_at_current_month_plus_thirteen() { + let now = utc(2026, 5, 15, 12); + let (start, end) = daily_range(date(2000, 1, 1), now); + assert_eq!(start, "2025-04-01"); + assert_eq!(end, "2026-05-16"); + assert_eq!(daily_range(NaiveDate::MIN, now).0, "2025-04-01"); + } + + #[test] + fn daily_range_keeps_a_recent_since() { + let now = utc(2026, 5, 15, 12); + let (start, _) = daily_range(date(2026, 3, 10), now); + assert_eq!(start, "2026-03-10"); + } + + fn day(start: &str, groups: &[(&str, &str)]) -> Value { + json!({ + "TimePeriod": { "Start": start, "End": "ignored" }, + "Groups": groups.iter().map(|(service, amount)| json!({ + "Keys": [service], + "Metrics": { "UnblendedCost": { "Amount": amount, "Unit": "USD" } } + })).collect::>() + }) + } + + #[test] + fn daily_costs_keep_only_positive_bedrock_days() { + let page = json!({ "ResultsByTime": [ + day("2026-03-01", &[("Amazon Bedrock", "7.00"), ("Amazon S3", "99.00")]), + day("2026-03-02", &[("Amazon Bedrock", "0.00")]), + day("2026-03-03", &[("Amazon S3", "5.00")]), + day("2026-03-04", &[("Claude Sonnet (Amazon Bedrock Edition)", "1.25"), ("Amazon Bedrock", "0.75")]), + ]}); + assert_eq!( + parse_daily_costs(&[page]), + vec![ + CostDailyPoint { + day: "2026-03-01".into(), + amount: 7.0 + }, + CostDailyPoint { + day: "2026-03-04".into(), + amount: 2.0 + }, + ] + ); + } + + #[test] + fn daily_costs_merge_the_same_day_across_pages_and_skip_bad_rows() { + let first = json!({ "ResultsByTime": [day("2026-03-01", &[("Amazon Bedrock", "3.00")])]}); + let second = json!({ "ResultsByTime": [ + day("2026-03-01", &[("Amazon Bedrock", "4.00")]), + day("not-a-date", &[("Amazon Bedrock", "9.00")]), + day("2026-03-02", &[("Amazon Bedrock", "NaN")]), + day("2026-03-05", &[("Amazon Bedrock", "-2.00")]), + ]}); + let points = parse_daily_costs(&[first, second, json!({})]); + assert_eq!( + points, + vec![CostDailyPoint { + day: "2026-03-01".into(), + amount: 7.0 + }] + ); + } +} diff --git a/rust/src/providers/bedrock/mod.rs b/rust/src/providers/bedrock/mod.rs index 4068eab03b..23aa92ea37 100644 --- a/rust/src/providers/bedrock/mod.rs +++ b/rust/src/providers/bedrock/mod.rs @@ -3,15 +3,19 @@ //! Fetches current-month Bedrock spend from AWS Cost Explorer using SigV4. use async_trait::async_trait; -use chrono::{DateTime, Datelike, Duration, Months, NaiveDate, TimeZone, Utc}; +use chrono::{Datelike, Duration, TimeZone, Utc}; use reqwest::Client; use serde_json::{Value, json}; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, hex, hmac_sha256, sha256_hex, + CostDailyPoint, CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, + ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, hex, hmac_sha256, + sha256_hex, }; -use crate::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; + +mod daily; + +use daily::{all_available_range, current_month_range, parse_daily_costs}; const COST_EXPLORER_URL: &str = "https://ce.us-east-1.amazonaws.com"; const COST_EXPLORER_TARGET: &str = "AWSInsightsIndexService.GetCostAndUsage"; @@ -300,27 +304,59 @@ impl BedrockProvider { credentials: &AwsCredentials, ) -> Result { let (start_date, end_date) = current_month_range(); - let mut total = 0.0; + let pages = self + .fetch_cost_pages(credentials, &start_date, &end_date, "MONTHLY") + .await?; + Ok(pages.iter().map(parse_bedrock_cost).sum()) + } + + /// Daily Bedrock spend over every month Cost Explorer exposes, so a + /// month-to-date or all-available selection can be answered from it. + async fn fetch_daily_spend( + &self, + credentials: &AwsCredentials, + ) -> Result, ProviderError> { + let (start_date, end_date) = all_available_range(); + let pages = self + .fetch_cost_pages(credentials, &start_date, &end_date, "DAILY") + .await?; + Ok(parse_daily_costs(&pages)) + } + + async fn fetch_cost_pages( + &self, + credentials: &AwsCredentials, + start_date: &str, + end_date: &str, + granularity: &str, + ) -> Result, ProviderError> { + let mut pages = Vec::new(); + let mut seen_tokens = std::collections::HashSet::new(); let mut next_page_token: Option = None; loop { let page = self .fetch_cost_page( credentials, - &start_date, - &end_date, + start_date, + end_date, + granularity, next_page_token.as_deref(), ) .await?; - total += parse_bedrock_cost(&page); next_page_token = extract_next_page_token(&page); - - if next_page_token.is_none() { - break; + pages.push(page); + + match &next_page_token { + None => return Ok(pages), + Some(token) if !seen_tokens.insert(token.clone()) => { + return Err(ProviderError::Parse( + "Cost Explorer returned repeated NextPageToken".to_string(), + )); + } + Some(_) => {} } } - - Ok(total) } async fn fetch_claude_activity( @@ -384,9 +420,10 @@ impl BedrockProvider { credentials: &AwsCredentials, start_date: &str, end_date: &str, + granularity: &str, next_page_token: Option<&str>, ) -> Result { - let body_bytes = cost_request_body(start_date, end_date, next_page_token)?; + let body_bytes = cost_request_body(start_date, end_date, granularity, next_page_token)?; let body_hash = sha256_hex(&body_bytes); let now = Utc::now(); let amz_date = now.format("%Y%m%dT%H%M%SZ").to_string(); @@ -461,6 +498,10 @@ impl BedrockProvider { } let mut cost = CostSnapshot::new(spend, "USD", "Monthly"); + match self.fetch_daily_spend(&credentials).await { + Ok(daily) => cost = cost.with_daily(daily), + Err(error) => tracing::debug!(%error, "Bedrock daily cost history unavailable"), + } if let Some(limit) = budget { cost = cost.with_limit(limit); } @@ -528,6 +569,7 @@ impl Provider for BedrockProvider { fn cost_request_body( start_date: &str, end_date: &str, + granularity: &str, next_page_token: Option<&str>, ) -> Result, ProviderError> { let mut body = json!({ @@ -535,7 +577,7 @@ fn cost_request_body( "Start": start_date, "End": end_date, }, - "Granularity": "MONTHLY", + "Granularity": granularity, "Metrics": ["UnblendedCost"], "GroupBy": [ { "Type": "DIMENSION", "Key": "SERVICE" } @@ -657,36 +699,6 @@ fn parse_aws_profile_credentials(stdout: &[u8]) -> Result) -> (String, String) { - let month_start = utc_month_start(now); - let earliest = month_start - .checked_sub_months(Months::new(13)) - .unwrap_or(month_start); - let tomorrow = now.date_naive() + Duration::days(1); - ( - since.max(earliest).format("%Y-%m-%d").to_string(), - tomorrow.format("%Y-%m-%d").to_string(), - ) -} - -fn utc_month_start(now: DateTime) -> NaiveDate { - CostReportingPeriod::MonthToDate - .bounds(now, CostTimeZone::Named(chrono_tz::UTC), None) - .start -} - -/// Current-month range: month to date through tomorrow (exclusive). -fn current_month_range() -> (String, String) { - let now = Utc::now(); - daily_range(utc_month_start(now), now) -} - fn end_of_current_month() -> Option> { let now = Utc::now(); let (year, month) = if now.month() == 12 { @@ -697,18 +709,13 @@ fn end_of_current_month() -> Option> { Utc.with_ymd_and_hms(year, month, 1, 0, 0, 0).single() } -fn parse_bedrock_cost(page: &Value) -> f64 { - page.get("ResultsByTime") +/// Amounts of the Bedrock service groups in one `ResultsByTime` entry. +fn bedrock_group_amounts(result: &Value) -> impl Iterator + '_ { + result + .get("Groups") .and_then(|v| v.as_array()) .into_iter() .flatten() - .flat_map(|result| { - result - .get("Groups") - .and_then(|v| v.as_array()) - .into_iter() - .flatten() - }) .filter(|group| { group .get("Keys") @@ -725,6 +732,14 @@ fn parse_bedrock_cost(page: &Value) -> f64 { .and_then(|v| v.as_str()) .and_then(|amount| amount.parse::().ok()) }) +} + +fn parse_bedrock_cost(page: &Value) -> f64 { + page.get("ResultsByTime") + .and_then(|v| v.as_array()) + .into_iter() + .flatten() + .flat_map(bedrock_group_amounts) .sum() } @@ -823,49 +838,6 @@ fn sanitized_body(body: &str) -> String { mod tests { use super::*; - fn utc(y: i32, m: u32, d: u32, h: u32) -> DateTime { - Utc.with_ymd_and_hms(y, m, d, h, 0, 0).single().unwrap() - } - - fn date(y: i32, m: u32, d: u32) -> NaiveDate { - NaiveDate::from_ymd_opt(y, m, d).unwrap() - } - - #[test] - fn daily_range_month_to_date_starts_at_utc_month_start() { - let now = utc(2026, 5, 15, 12); - let (start, end) = daily_range(utc_month_start(now), now); - assert_eq!(start, "2026-05-01"); - assert_eq!(end, "2026-05-16"); - } - - #[test] - fn daily_range_month_start_uses_utc_not_local_time() { - // 23:30 UTC on the last day of April is still April in Cost Explorer. - let now = Utc - .with_ymd_and_hms(2026, 4, 30, 23, 30, 0) - .single() - .unwrap(); - let (start, end) = daily_range(utc_month_start(now), now); - assert_eq!(start, "2026-04-01"); - assert_eq!(end, "2026-05-01"); - } - - #[test] - fn daily_range_all_is_capped_at_current_month_plus_thirteen() { - let now = utc(2026, 5, 15, 12); - let (start, end) = daily_range(date(2000, 1, 1), now); - assert_eq!(start, "2025-04-01"); - assert_eq!(end, "2026-05-16"); - } - - #[test] - fn daily_range_keeps_a_recent_since() { - let now = utc(2026, 5, 15, 12); - let (start, _) = daily_range(date(2026, 3, 10), now); - assert_eq!(start, "2026-03-10"); - } - #[test] fn parses_bedrock_cost_only() { let page = json!({ From f292961c521d8337e29fc53118f8aa19310e9755 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:34:58 +0700 Subject: [PATCH 119/259] Port upstream 0.64.0: Antigravity cost follow-ups (lower bounds, withheld reads, --refresh) Partial Antigravity history never replaces a previously complete read, contradicted or budget-truncated reads are withheld, and trustworthy partial reads are published as lower bounds (tokensAreLowerBound / costIsLowerBound). Adds cost --provider antigravity --refresh for one bounded pricing refresh; routine reads never download pricing. Ports upstream steipete/CodexBar PR #3757 (7eebfd4f1) in part. --- .../src-tauri/src/commands/usage_spend.rs | 35 ++++- .../src/lib/usageSpendSharing.test.ts | 8 + .../src/lib/usageSpendSharing.ts | 5 +- .../surfaces/settings/tabs/UsageSpendTab.tsx | 2 + apps/desktop-tauri/src/types/bridge.ts | 3 + docs/CLI.md | 2 + rust/src/cli/cost.rs | 33 +++- rust/src/cli/serve/data.rs | 2 + rust/src/providers/antigravity/cost.rs | 62 +++++++- .../providers/antigravity/local_history.rs | 107 ++++++++++++- .../providers/antigravity/local_sessions.rs | 5 +- .../antigravity/local_sessions_reader.rs | 34 +++- .../src/providers/antigravity/local_sqlite.rs | 74 ++++++--- .../antigravity/local_sqlite_tests.rs | 64 ++++++++ rust/src/providers/muse/local_usage/mod.rs | 1 + rust/src/spend_contract.rs | 109 +------------ rust/src/spend_contract/local_history.rs | 147 ++++++++++++++++++ rust/src/spend_contract/tests.rs | 81 ++++++++++ 18 files changed, 623 insertions(+), 151 deletions(-) create mode 100644 rust/src/spend_contract/local_history.rs 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 4114cccc53..5608e8c833 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -34,6 +34,11 @@ pub struct UsageSpendRow { pub thirty_day_estimate: Option, pub seven_day_tokens: Option, pub thirty_day_tokens: Option, + /// The token figure is a floor from an incomplete scan ("at least N"). + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub seven_day_tokens_lower_bound: bool, + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub thirty_day_tokens_lower_bound: bool, pub currency: String, pub source: String, /// Included in the shared Overview spend denominator. @@ -547,6 +552,8 @@ fn build_usage_spend_summary_in_containment_proof( thirty_day_estimate, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, + seven_day_tokens_lower_bound: seven.lower_bound, + thirty_day_tokens_lower_bound: thirty.lower_bound, currency, source: spend.source, included_in_overview: true, @@ -709,6 +716,7 @@ fn build_usage_spend_summary( .unwrap_or_else(|| provider_id.clone()); let mut local_cost_estimates = None; + let mut token_lower_bounds = (false, false); let spend = match provider_id.as_str() { "codex" => SpendValues { seven_day: codex_7_contract.known_cost_usd, @@ -803,6 +811,7 @@ fn build_usage_spend_summary( let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); let spend = antigravity_spend_values(cached_spend(cached_snapshot), &seven, &thirty); + token_lower_bounds = (seven.lower_bound, thirty.lower_bound); local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } @@ -837,6 +846,8 @@ fn build_usage_spend_summary( thirty_day_estimate, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, + seven_day_tokens_lower_bound: token_lower_bounds.0, + thirty_day_tokens_lower_bound: token_lower_bounds.1, currency, source: spend.source, included_in_overview: include_in_shared_overview( @@ -922,10 +933,9 @@ fn antigravity_spend_values( 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); + // Exact for a complete scan, a floor for a lower bound, unknown otherwise. + spend.seven_day_tokens = seven.published_tokens(); + spend.thirty_day_tokens = thirty.published_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() { @@ -1031,7 +1041,9 @@ mod cache_key_tests { unpriced, ..Default::default() }, + ..Default::default() }, + ..Default::default() } } @@ -1107,6 +1119,21 @@ mod cache_key_tests { assert!(spend.source.contains("known API list-price subtotal")); } + #[test] + fn antigravity_lower_bound_history_publishes_floors_not_exact_totals() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let mut seven = local_history(100, LocalHistoryCoverage::Partial, Some(1.25), 0); + seven.lower_bound = true; + let withheld = local_history(0, LocalHistoryCoverage::Partial, None, 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &withheld); + + assert_eq!(spend.seven_day, None); + assert_eq!(spend.seven_day_tokens, Some(100)); + assert_eq!(spend.thirty_day, None); + assert_eq!(spend.thirty_day_tokens, None); + } + #[test] fn antigravity_complete_empty_history_is_a_known_zero() { use codexbar::spend_contract::LocalHistoryCoverage; diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index 39c95c9a3f..c1eac606b1 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("marks a lower-bound token count with a floor sign and leaves exact counts bare", () => { + const floor = formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125, true); + expect(floor).toContain("≥1,500 tokens"); + expect(floor).toMatch(/^≥.* known/); + const exact = formatSpendMetric(null, 1_500, "USD", "tokens", null, false); + expect(exact).toBe("1,500 tokens"); + }); + it("renders a complete known-zero total instead of a subtotal", () => { const metric = formatSpendMetric(0, 0, "USD", "tokens", 9); expect(metric).not.toBe("—"); diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 4045eb9199..a792d0c53b 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -131,6 +131,7 @@ export function formatSpendMetric( currency: string, tokenLabel: string, knownSubtotal?: number | null, + tokensAreLowerBound = false, ): string { const parts: string[] = []; if (cost != null && Number.isFinite(cost)) { @@ -139,7 +140,7 @@ export function formatSpendMetric( parts.push(`≥${formatUsd(knownSubtotal, currency)} known`); } if (tokens != null && Number.isFinite(tokens)) { - parts.push(`${Math.max(0, tokens).toLocaleString()} ${tokenLabel}`); + parts.push(`${tokensAreLowerBound ? "≥" : ""}${Math.max(0, tokens).toLocaleString()} ${tokenLabel}`); } return parts.length > 0 ? parts.join(" · ") : "—"; } @@ -213,6 +214,7 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri row.currency, "tokens", row.sevenDayEstimate?.knownSubtotalUsd, + row.sevenDayTokensLowerBound, ), formatSpendMetric( row.thirtyDay, @@ -220,6 +222,7 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri row.currency, "tokens", row.thirtyDayEstimate?.knownSubtotalUsd, + row.thirtyDayTokensLowerBound, ), row.currency || "USD", row.source, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index e7ee282ece..bd3ec36c19 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -241,6 +241,7 @@ export default function UsageSpendTab(_props: TabProps) { row.currency, t("UsageSpendTokens"), row.sevenDayEstimate?.knownSubtotalUsd, + row.sevenDayTokensLowerBound, )} @@ -250,6 +251,7 @@ export default function UsageSpendTab(_props: TabProps) { row.currency, t("UsageSpendTokens"), row.thirtyDayEstimate?.knownSubtotalUsd, + row.thirtyDayTokensLowerBound, )} {row.currency || "USD"} diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index fd8dd55109..0ce0c4abe7 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -387,6 +387,9 @@ export interface UsageSpendRow { thirtyDayEstimate?: LocalCostEstimate; sevenDayTokens?: number | null; thirtyDayTokens?: number | null; + /** The token figure is a floor from an incomplete scan ("at least N"). */ + sevenDayTokensLowerBound?: boolean; + thirtyDayTokensLowerBound?: boolean; currency: string; source: string; includedInOverview: boolean; diff --git a/docs/CLI.md b/docs/CLI.md index a9c25bcdcc..9fea657f54 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -68,6 +68,8 @@ codexbar cost -p codex --format json --summary-only --provider-native-only --day Claude/Codex costs come from local session logs. Antigravity exposes local **token history only** through `cost`; dollar cost remains unknown rather than becoming a false `$0`. Other providers may differ; do not assume upstream Cursor dashboard cost behavior unless implemented in this tree. +Antigravity history that stopped short is never shown as exact. A scan whose decoded rows are trustworthy reports `tokensAreLowerBound` / `costIsLowerBound` (text: "at least N"); a scan cut off by a hard limit, or one whose sources contradict each other, is withheld and publishes no total. A later partial read does not replace an earlier complete read of the same window and roots within one process. Routine reads never download pricing; `codexbar cost --provider antigravity --refresh` runs one bounded models.dev pricing refresh when a recorded model has no known public price, then rescans. + `--remote` adds one separate native Codex report fetched through non-interactive SSH; overlapping local and remote histories are never combined. `--summary-only` emits the versioned, path-free JSON contract used by the remote comparison and accepts only `--provider codex --format json`. Both modes reject session grouping and other provider selections. Codex local-history scans use a 60-second scanner-side debounce for ordinary disk-cache reads. This is separate from the desktop provider refresh setting. With Adaptive refresh off, **Manual** (`refresh_interval_secs = 0`) disables the recurring desktop refresh timer, but it does not forbid startup/stale-aware reads, explicit refreshes, or pending Codex catch-up scans. Low Power Mode floors recurring automatic refreshes to 30 minutes; explicit/manual work remains immediate. diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 819b7abb59..1ff9788bc8 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -62,6 +62,12 @@ pub struct CostArgs { /// Emit the versioned native Codex summary contract as JSON. #[arg(long = "summary-only")] pub summary_only: bool, + + /// Antigravity only: when a recorded model has no known public price, run + /// one bounded models.dev pricing refresh and rescan. Routine reads never + /// download pricing. + #[arg(long)] + pub refresh: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -169,9 +175,13 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { display_name: provider.display_name().to_string(), summary: CostSummary::default(), supported: true, - token_history: Some(crate::providers::antigravity::local_sessions::summarize( - args.days, - )), + token_history: Some( + crate::providers::antigravity::local_sessions::summarize_with_pricing_refresh( + args.days, + args.refresh, + ) + .await, + ), }); } ProviderId::Muse => { @@ -375,6 +385,13 @@ 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) { use crate::spend_contract::LocalHistoryCoverage; match history.coverage { + LocalHistoryCoverage::Partial if history.lower_bound => { + println!( + " Tokens: at least {} (lower bound; scan stopped short)", + format_number(history.total_tokens) + ); + println!(" Sessions: at least {}", history.session_count); + } LocalHistoryCoverage::Complete if history.total_tokens == 0 => { println!(" No token usage in the last {days} days (scan complete)"); } @@ -395,7 +412,7 @@ fn print_local_token_history(history: &crate::spend_contract::LocalTokenHistoryS history.cost_estimate.coverage.unpriced ); } else { - println!(" Known API list-price subtotal: ${cost:.2} (history incomplete)"); + println!(" Known API list-price subtotal: at least ${cost:.2} (history incomplete)"); } } else { println!(" Local token history; dollar costs unavailable"); @@ -639,6 +656,7 @@ mod tests { session_count: 2, coverage: LocalHistoryCoverage::Complete, cost_estimate: Default::default(), + ..Default::default() }, 30, ); @@ -654,6 +672,7 @@ mod tests { session_count: 1, coverage: LocalHistoryCoverage::Partial, cost_estimate: Default::default(), + ..Default::default() }, 30, ); @@ -677,7 +696,9 @@ mod tests { estimated: 1, ..Default::default() }, + ..Default::default() }, + ..Default::default() }, 30, ); @@ -704,7 +725,9 @@ mod tests { estimated: 1, ..Default::default() }, + ..Default::default() }, + ..Default::default() }, 30, ); @@ -754,7 +777,9 @@ mod tests { unpriced: 1, ..Default::default() }, + ..Default::default() }, + ..Default::default() }, 30, ); diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index 53e0f62d07..10af1b74d2 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -164,6 +164,7 @@ mod tests { session_count: 1, coverage: LocalHistoryCoverage::Complete, cost_estimate: Default::default(), + ..Default::default() }, 30, ); @@ -178,6 +179,7 @@ mod tests { session_count: 1, coverage: LocalHistoryCoverage::Partial, cost_estimate: Default::default(), + ..Default::default() }, 30, ); diff --git a/rust/src/providers/antigravity/cost.rs b/rust/src/providers/antigravity/cost.rs index f8eb38f24a..7254245ca4 100644 --- a/rust/src/providers/antigravity/cost.rs +++ b/rust/src/providers/antigravity/cost.rs @@ -1,4 +1,42 @@ +use std::collections::HashSet; + use crate::core::CostUsagePricing; +use crate::spend_contract::LocalCostEstimate; + +/// Antigravity records routing variants of a vendor model that bill at the base +/// model's public price. The alias stays provider-local so shared Claude pricing +/// keeps reporting unknown Claude variants as unpriced. +fn pricing_base_model(model: &str) -> Option<&str> { + let lowered = model.to_ascii_lowercase(); + ["-tiered", "-low", "-thinking"] + .iter() + .find(|suffix| lowered.ends_with(**suffix)) + .map(|suffix| &model[..model.len() - suffix.len()]) + .filter(|base| !base.is_empty()) +} + +/// Models worth fetching for an explicit pricing refresh: each unpriced model +/// and, for a routing variant, its base model. +fn refresh_model_ids(estimate: &LocalCostEstimate) -> HashSet { + estimate + .unpriced_models + .iter() + .flat_map(|model| { + [Some(model.as_str()), pricing_base_model(model)] + .into_iter() + .flatten() + .map(str::to_string) + }) + .collect() +} + +/// Run one bounded models.dev refresh for unpriced models. Returns true when a +/// rescan may now price them. Empty or fully priced history never downloads. +pub(super) async fn refresh_unpriced_model_pricing(estimate: &LocalCostEstimate) -> bool { + let model_ids = refresh_model_ids(estimate); + !model_ids.is_empty() + && crate::core::refresh_unknown_models_if_needed("anthropic", &model_ids).await +} pub(super) fn estimate_cost_usd( model: Option<&str>, @@ -16,18 +54,26 @@ pub(super) fn estimate_cost_usd( 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) - }) + resolve(model).or_else(|| pricing_base_model(model).and_then(resolve)) } #[cfg(test)] mod tests { - use super::estimate_cost_usd; + use super::{estimate_cost_usd, refresh_model_ids}; + use crate::spend_contract::LocalCostEstimate; + + #[test] + fn refresh_targets_unpriced_models_and_their_routing_base() { + let mut estimate = LocalCostEstimate::default(); + estimate.record_list_price(Some("claude-future-9-thinking"), None); + estimate.record_list_price(Some("claude-future-9-thinking"), None); + estimate.record_list_price(None, None); + let ids = refresh_model_ids(&estimate); + assert_eq!(ids.len(), 2); + assert!(ids.contains("claude-future-9-thinking")); + assert!(ids.contains("claude-future-9")); + assert!(refresh_model_ids(&LocalCostEstimate::default()).is_empty()); + } #[test] fn prices_known_models_and_provider_local_routing_variants() { diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs index 464201feab..39de8dd482 100644 --- a/rust/src/providers/antigravity/local_history.rs +++ b/rust/src/providers/antigravity/local_history.rs @@ -1,10 +1,12 @@ use super::{local_sessions_reader as local_sessions, local_sqlite}; +use std::collections::HashMap; use std::fs; use std::path::{Path, PathBuf}; +use std::sync::{LazyLock, Mutex}; use chrono::{DateTime, Utc}; -use crate::spend_contract::LocalTokenHistorySummary; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; fn clean_env_path(value: Option<&str>) -> Option { value @@ -51,6 +53,42 @@ pub(super) fn summarize_local_usage_from_explicit_roots( }) } +/// Last complete summary per scan scope (roots plus window). A later partial +/// or withheld read of the same scope must not replace previously complete +/// history; only a newer complete read does. +#[derive(Default)] +struct CompleteHistoryRetention { + complete: Mutex>, +} + +impl CompleteHistoryRetention { + fn resolve(&self, scope: &str, fresh: LocalTokenHistorySummary) -> LocalTokenHistorySummary { + let mut complete = self + .complete + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + match fresh.coverage { + LocalHistoryCoverage::Complete => { + complete.insert(scope.to_string(), fresh.clone()); + fresh + } + LocalHistoryCoverage::Partial => complete.get(scope).cloned().unwrap_or(fresh), + LocalHistoryCoverage::Unavailable => fresh, + } + } +} + +static COMPLETE_HISTORY: LazyLock = LazyLock::new(Default::default); + +fn retention_scope(roots: &[PathBuf], tokscale_sessions: &Path, days: u32) -> String { + let roots = roots + .iter() + .map(|root| root.to_string_lossy()) + .collect::>() + .join("|"); + format!("{days}|{roots}|{}", tokscale_sessions.to_string_lossy()) +} + pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { let now = Utc::now(); let Some(home) = dirs::home_dir() else { @@ -58,7 +96,29 @@ 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_explicit_roots(&roots, &tokscale_sessions, now, days) + let fresh = summarize_local_usage_from_explicit_roots(&roots, &tokscale_sessions, now, days); + COMPLETE_HISTORY.resolve(&retention_scope(&roots, &tokscale_sessions, days), fresh) +} + +/// CLI `cost --refresh`: when a recorded model has no known public price, run +/// one bounded pricing refresh and rescan. Routine reads never download pricing. +pub async fn summarize_local_usage_with_pricing_refresh( + days: u32, + refresh: bool, +) -> LocalTokenHistorySummary { + let first = summarize_local_usage(days); + if !refresh || !super::cost::refresh_unpriced_model_pricing(&first.cost_estimate).await { + return first; + } + let rescanned = summarize_local_usage(days); + // A pricing download must not replace a complete scan with one that became + // partial in the meantime. + if first.coverage == LocalHistoryCoverage::Complete + && rescanned.coverage != LocalHistoryCoverage::Complete + { + return first; + } + rescanned } /// Count local Antigravity conversation artifacts for the quota provider's @@ -107,7 +167,6 @@ fn count_extension(root: &Path, extension: &str) -> usize { #[cfg(test)] mod tests { use super::*; - use crate::spend_contract::LocalHistoryCoverage; use chrono::TimeZone; use rusqlite::Connection; @@ -225,4 +284,46 @@ mod tests { fs::write(cache.join("one.jsonl"), b"{}\n").unwrap(); assert_eq!(offline_conversation_count_in(dir.path()), 1); } + + #[test] + fn partial_read_never_replaces_previously_complete_history() { + let retention = CompleteHistoryRetention::default(); + let complete = LocalTokenHistorySummary { + total_tokens: 500, + session_count: 3, + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }; + let partial = LocalTokenHistorySummary { + total_tokens: 20, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + lower_bound: true, + ..Default::default() + }; + + // Nothing complete yet: the partial read is reported as it is. + assert_eq!(retention.resolve("scope", partial.clone()), partial); + assert_eq!(retention.resolve("scope", complete.clone()), complete); + // A later partial or withheld read keeps the complete history. + assert_eq!(retention.resolve("scope", partial.clone()), complete); + assert_eq!( + retention.resolve("scope", LocalTokenHistorySummary::withheld()), + complete + ); + // Another scope is unaffected. + assert_eq!(retention.resolve("other", partial.clone()), partial); + // A newer complete read replaces the retained one. + let newer = LocalTokenHistorySummary { + total_tokens: 700, + ..complete.clone() + }; + assert_eq!(retention.resolve("scope", newer.clone()), newer); + assert_eq!(retention.resolve("scope", partial), newer); + // Source absence is reported honestly, not masked by old history. + assert_eq!( + retention.resolve("scope", LocalTokenHistorySummary::default()), + LocalTokenHistorySummary::default() + ); + } } diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index d2d1be7b74..ee3e0c8a2b 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -1,4 +1,7 @@ -pub use super::local_history::{offline_conversation_count, summarize_local_usage as summarize}; +pub use super::local_history::{ + offline_conversation_count, summarize_local_usage as summarize, + summarize_local_usage_with_pricing_refresh as summarize_with_pricing_refresh, +}; 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 index cbc3297cb4..6b494f628e 100644 --- a/rust/src/providers/antigravity/local_sessions_reader.rs +++ b/rust/src/providers/antigravity/local_sessions_reader.rs @@ -145,11 +145,15 @@ fn summarize_paths_with_budget( let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); let mut complete = !truncated; + // A hard discovery or byte budget stops the read; such a truncated scan is + // withheld instead of being published as a lower bound. + let mut exhausted = truncated || paths.len() > MAX_SESSION_FILES; let mut remaining_total_bytes = total_byte_budget; for path in paths.iter().take(MAX_SESSION_FILES) { if remaining_total_bytes == 0 { complete = false; + exhausted = true; break; } let file = match File::open(path) { @@ -181,6 +185,7 @@ fn summarize_paths_with_budget( } Ok(Some(BoundedJsonlLine::Truncated)) => { complete = false; + exhausted = true; break; } Ok(None) => break, @@ -266,13 +271,16 @@ fn summarize_paths_with_budget( continue; }; total_tokens = next_total_tokens; - cost_estimate.record_list_price(estimate_cost_usd( + cost_estimate.record_list_price( model.as_deref(), - input, - cache_read, - cache_write, - output.saturating_add(reasoning), - )); + estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + ), + ); path_had_usage = true; } if path_had_usage { @@ -280,6 +288,10 @@ fn summarize_paths_with_budget( } } + if exhausted { + return LocalTokenHistorySummary::withheld(); + } + LocalTokenHistorySummary { total_tokens, session_count: sessions_with_usage.len(), @@ -291,7 +303,9 @@ fn summarize_paths_with_budget( LocalHistoryCoverage::Partial }, cost_estimate, + lower_bound: false, } + .with_lower_bound_if_partial() } fn read_bounded_jsonl_line( @@ -536,6 +550,9 @@ mod tests { assert_eq!(summary.total_tokens, 15); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + // Decoded rows survive as a floor, never as an exact total. + assert!(summary.lower_bound); + assert_eq!(summary.published_tokens(), Some(15)); } #[test] @@ -555,8 +572,11 @@ mod tests { let summary = summarize_paths_with_budget(&[first_path, second_path], now, 7, false, first.len()); - assert_eq!(summary.total_tokens, 10); + // A hard budget stop is withheld: no total is published from it. + assert_eq!(summary.total_tokens, 0); assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert!(!summary.lower_bound); + assert_eq!(summary.published_tokens(), None); } #[test] diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index bc71f794d7..cbe52ec19a 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -60,6 +60,9 @@ struct Budget { bytes: usize, schema_bytes: usize, deadline: Instant, + /// A hard scan limit (databases, rows, bytes, duration) stopped the read. + /// A truncated read is withheld rather than published as a lower bound. + exhausted: bool, } impl Budget { @@ -75,6 +78,7 @@ impl Budget { bytes: 0, schema_bytes: 0, deadline, + exhausted: false, } } @@ -151,11 +155,13 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL for path in &paths { if !budget.check() { complete = false; + budget.exhausted = true; break; } budget.databases += 1; if budget.databases > MAX_DATABASES { complete = false; + budget.exhausted = true; break; } match read_database(path, &mut budget) { @@ -178,6 +184,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL complete &= budget.check(); if budget.rows >= MAX_ROWS || budget.bytes >= MAX_TOTAL_BYTES { complete = false; + budget.exhausted = true; break; } } @@ -186,6 +193,13 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL return SQLiteScan::Unsupported; } + // A truncated read or two sources disagreeing about one request means the + // surviving rows may be wrong, not merely incomplete: withhold them. + if budget.exhausted || !budget.check() { + return SQLiteScan::Summary(LocalTokenHistorySummary::withheld()); + } + + let mut contradicted = false; let mut total_tokens = 0_u64; let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions = HashSet::new(); @@ -212,6 +226,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL if let Some(prior) = rows.get(&row_key) { if prior != &event { complete = false; + contradicted = true; } continue; } @@ -226,6 +241,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL if let Some(prior) = responses.get(&response_key) { if prior.turn != event.turn { complete = false; + contradicted = true; } else { rows.insert(row_key, event); } @@ -254,15 +270,15 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } } if event.total > 0 { + 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 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) { @@ -271,21 +287,29 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL None } }); - cost_estimate.record_list_price(estimated_cost); + cost_estimate.record_list_price(model, estimated_cost); } sessions.insert(event.session); } - SQLiteScan::Summary(LocalTokenHistorySummary { - total_tokens, - session_count: sessions.len(), - coverage: if complete { - LocalHistoryCoverage::Complete - } else { - LocalHistoryCoverage::Partial - }, - cost_estimate, - }) + if contradicted { + return SQLiteScan::Summary(LocalTokenHistorySummary::withheld()); + } + + SQLiteScan::Summary( + LocalTokenHistorySummary { + total_tokens, + session_count: sessions.len(), + coverage: if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, + cost_estimate, + lower_bound: false, + } + .with_lower_bound_if_partial(), + ) } fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, bool) { @@ -294,6 +318,7 @@ fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, for root in roots { if !budget.check() { + budget.exhausted = true; return (paths, false); } let resolved_root = match fs::canonicalize(root) { @@ -320,10 +345,12 @@ fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, }; for entry in entries { if !budget.check() { + budget.exhausted = true; return (paths, false); } budget.directory_entries += 1; if budget.directory_entries > MAX_DIRECTORY_ENTRIES { + budget.exhausted = true; return (paths, false); } let entry = match entry { @@ -363,6 +390,7 @@ fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, } } if paths.len() >= MAX_DATABASES { + budget.exhausted = true; return (paths, false); } paths.push(resolved); @@ -509,12 +537,14 @@ fn read_generation_rows( while let Some(row) = query.next()? { if !budget.check() { complete = false; + budget.exhausted = true; break; } database_rows += 1; budget.rows += 1; if database_rows > MAX_ROWS_PER_DATABASE || budget.rows > MAX_ROWS { complete = false; + budget.exhausted = true; break; } @@ -534,6 +564,7 @@ fn read_generation_rows( Some(value) if value <= MAX_DATABASE_BYTES => value, _ => { complete = false; + budget.exhausted = true; break; } }; @@ -541,6 +572,7 @@ fn read_generation_rows( Some(value) if value <= MAX_TOTAL_BYTES => value, _ => { complete = false; + budget.exhausted = true; break; } }; @@ -639,11 +671,13 @@ fn read_step_timestamps( while let Some(row) = query.next()? { if !budget.check() { complete = false; + budget.exhausted = true; break; } budget.rows += 1; if budget.rows > MAX_ROWS { complete = false; + budget.exhausted = true; break; } @@ -663,6 +697,7 @@ fn read_step_timestamps( Some(value) if value <= MAX_DATABASE_BYTES => value, _ => { complete = false; + budget.exhausted = true; break; } }; @@ -670,6 +705,7 @@ fn read_step_timestamps( Some(value) if value <= MAX_TOTAL_BYTES => value, _ => { complete = false; + budget.exhausted = true; break; } }; diff --git a/rust/src/providers/antigravity/local_sqlite_tests.rs b/rust/src/providers/antigravity/local_sqlite_tests.rs index 1846bd5e26..73f7bd2591 100644 --- a/rust/src/providers/antigravity/local_sqlite_tests.rs +++ b/rust/src/providers/antigravity/local_sqlite_tests.rs @@ -278,3 +278,67 @@ fn schema_entry_budget_is_incomplete_not_foreign() { SchemaInspection::Incomplete ); } + +#[test] +fn undecodable_row_beside_valid_rows_yields_a_lower_bound() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).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, ?1)", + [valid_turn_blob(100, timestamp)], + ) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(2, ?1)", + params!["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); + assert!(summary.total_tokens > 0); + assert!(summary.lower_bound); + assert_eq!(summary.published_tokens(), Some(summary.total_tokens)); + assert_eq!(summary.total_usd(), None); +} + +#[test] +fn contradicting_rows_for_one_index_are_withheld_not_published() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + for input in [100_u64, 900_u64] { + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .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); + assert_eq!(summary.total_tokens, 0); + assert!(!summary.lower_bound); + assert_eq!(summary.published_tokens(), None); +} diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 24e84b2473..004079a624 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -65,6 +65,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { session_count: report.session_count, coverage: report.coverage, cost_estimate: Default::default(), + ..Default::default() } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index 722acec630..282919c7d3 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -1,8 +1,13 @@ //! Unified Usage & Spend accounting contract for upstream 0.53 parity. //! Accounting semantics live here so UI/CLI never infer unknown vs zero. +mod local_history; mod opencodex; +pub use local_history::{ + LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, local_token_history_json, +}; + use std::collections::{BTreeMap, HashMap, HashSet}; use std::fs; use std::path::PathBuf; @@ -70,110 +75,6 @@ impl CostProvenance { } } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub enum LocalHistoryCoverage { - Complete, - Partial, - #[default] - Unavailable, -} - -#[derive(Debug, Clone, Default, PartialEq)] -pub struct LocalTokenHistorySummary { - pub total_tokens: u64, - pub session_count: usize, - pub coverage: LocalHistoryCoverage, - 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 { - /// 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 { - fn complete_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( - provider: &str, - history: &LocalTokenHistorySummary, - days: u32, -) -> serde_json::Value { - let complete = history.coverage == LocalHistoryCoverage::Complete; - 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, - "days_scanned": days, - "cost": { - "total_usd": total_usd, - "known_subtotal_usd": known_subtotal_usd, - "currency": total_usd.or(known_subtotal_usd).map(|_| "USD"), - "pricingCoverage": &history.cost_estimate.coverage, - }, - "daily": [], - "tokens": {"total": complete.then_some(history.total_tokens)}, - "sessions_count": complete.then_some(history.session_count), - "historyCoverage": match history.coverage { - LocalHistoryCoverage::Complete => "complete", - LocalHistoryCoverage::Partial => "partial", - LocalHistoryCoverage::Unavailable => "unavailable", - }, - "knownZero": complete && history.total_tokens == 0, - "note": note, - }) -} #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct CostCoverageCounts { diff --git a/rust/src/spend_contract/local_history.rs b/rust/src/spend_contract/local_history.rs new file mode 100644 index 0000000000..745b45f676 --- /dev/null +++ b/rust/src/spend_contract/local_history.rs @@ -0,0 +1,147 @@ +use std::collections::BTreeSet; + +use serde::{Deserialize, Serialize}; + +use super::CostCoverageCounts; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum LocalHistoryCoverage { + Complete, + Partial, + #[default] + Unavailable, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct LocalTokenHistorySummary { + pub total_tokens: u64, + pub session_count: usize, + pub coverage: LocalHistoryCoverage, + pub cost_estimate: LocalCostEstimate, + /// The scan stopped short but every row it decoded is trustworthy, so the + /// totals are floors ("at least N"), never exact. Readers that cannot tell a + /// trustworthy subset from contradicted evidence leave this `false`. + pub lower_bound: bool, +} + +impl LocalTokenHistorySummary { + /// A read that stopped short in a way that makes its rows untrustworthy + /// (hard budget exhausted, or sources contradicting each other). The scan is + /// known to be incomplete, but no total is published from it. + pub(crate) fn withheld() -> Self { + Self { + coverage: LocalHistoryCoverage::Partial, + ..Self::default() + } + } + + /// Mark a partial scan that still decoded rows as a lower bound. + pub(crate) fn with_lower_bound_if_partial(mut self) -> Self { + self.lower_bound = self.coverage == LocalHistoryCoverage::Partial && self.total_tokens > 0; + self + } + + /// Token total safe to publish: exact for a complete scan, a floor for a + /// marked lower bound, unknown otherwise. + pub fn published_tokens(&self) -> Option { + (self.coverage == LocalHistoryCoverage::Complete || self.lower_bound) + .then_some(self.total_tokens) + } + + /// 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 { + /// 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, + /// Recorded model names with no public price. Local-only input for an + /// explicit pricing refresh; never part of the wire contract. + #[serde(skip)] + pub unpriced_models: BTreeSet, +} + +impl LocalCostEstimate { + fn complete_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, model: Option<&str>, cost: Option) { + let Some(cost) = cost.filter(|value| value.is_finite() && *value >= 0.0) else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + if let Some(model) = model.map(str::trim).filter(|model| !model.is_empty()) { + self.unpriced_models.insert(model.to_string()); + } + 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( + provider: &str, + history: &LocalTokenHistorySummary, + days: u32, +) -> serde_json::Value { + let complete = history.coverage == LocalHistoryCoverage::Complete; + let total_usd = history.total_usd(); + let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; + let published_tokens = history.published_tokens(); + 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 (lower bound)" + } 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, + "days_scanned": days, + "cost": { + "total_usd": total_usd, + "known_subtotal_usd": known_subtotal_usd, + "currency": total_usd.or(known_subtotal_usd).map(|_| "USD"), + "pricingCoverage": &history.cost_estimate.coverage, + }, + "daily": [], + "tokens": {"total": published_tokens}, + "sessions_count": (complete || history.lower_bound).then_some(history.session_count), + "tokensAreLowerBound": history.lower_bound, + "costIsLowerBound": known_subtotal_usd.is_some() && total_usd.is_none(), + "historyCoverage": match history.coverage { + LocalHistoryCoverage::Complete => "complete", + LocalHistoryCoverage::Partial => "partial", + LocalHistoryCoverage::Unavailable => "unavailable", + }, + "knownZero": complete && history.total_tokens == 0, + "note": note, + }) +} diff --git a/rust/src/spend_contract/tests.rs b/rust/src/spend_contract/tests.rs index ac32e2e5eb..9595cfc8dc 100644 --- a/rust/src/spend_contract/tests.rs +++ b/rust/src/spend_contract/tests.rs @@ -8,12 +8,14 @@ fn local_history_total_requires_complete_scan_and_pricing() { estimated: 1, ..Default::default() }, + ..Default::default() }; let partial_history = LocalTokenHistorySummary { total_tokens: 100, session_count: 1, coverage: LocalHistoryCoverage::Partial, cost_estimate: priced.clone(), + ..Default::default() }; assert_eq!(partial_history.total_usd(), None); assert_eq!(partial_history.cost_estimate.known_subtotal_usd, Some(1.25)); @@ -29,7 +31,9 @@ fn local_history_total_requires_complete_scan_and_pricing() { unpriced: 1, ..Default::default() }, + ..Default::default() }, + ..Default::default() }; assert_eq!(mixed_pricing.total_usd(), None); assert_eq!(mixed_pricing.cost_estimate.known_subtotal_usd, Some(1.25)); @@ -39,6 +43,7 @@ fn local_history_total_requires_complete_scan_and_pricing() { session_count: 1, coverage: LocalHistoryCoverage::Complete, cost_estimate: priced, + ..Default::default() }; assert_eq!(complete.total_usd(), Some(1.25)); } @@ -571,3 +576,79 @@ fn coverage_for_models_counts_priced_rows_as_estimated() { assert_eq!(coverage.unpriced, 1); assert_eq!(coverage.total(), 3); } + +#[test] +fn lower_bound_history_publishes_floors_and_never_an_exact_total() { + let priced = LocalCostEstimate { + known_subtotal_usd: Some(0.5), + coverage: CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + ..Default::default() + }; + let scanned = LocalTokenHistorySummary { + total_tokens: 900, + session_count: 2, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: priced, + ..Default::default() + }; + let floor = scanned.with_lower_bound_if_partial(); + assert!(floor.lower_bound); + assert_eq!(floor.published_tokens(), Some(900)); + assert_eq!(floor.total_usd(), None); + + let payload = local_token_history_json("antigravity", &floor, 30); + assert_eq!(payload["tokens"]["total"], 900); + assert_eq!(payload["sessions_count"], 2); + assert_eq!(payload["tokensAreLowerBound"], true); + assert_eq!(payload["costIsLowerBound"], true); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.5); +} + +#[test] +fn withheld_history_is_partial_with_no_published_numbers() { + let withheld = LocalTokenHistorySummary::withheld(); + assert_eq!(withheld.coverage, LocalHistoryCoverage::Partial); + assert!(!withheld.lower_bound); + assert_eq!(withheld.published_tokens(), None); + assert!(!withheld.clone().with_lower_bound_if_partial().lower_bound); + + let payload = local_token_history_json("antigravity", &withheld, 30); + assert!(payload["tokens"]["total"].is_null()); + assert!(payload["sessions_count"].is_null()); + assert_eq!(payload["tokensAreLowerBound"], false); + assert_eq!(payload["costIsLowerBound"], false); + assert_eq!(payload["historyCoverage"], "partial"); +} + +#[test] +fn complete_history_is_never_marked_as_a_lower_bound() { + let complete = LocalTokenHistorySummary { + total_tokens: 10, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + } + .with_lower_bound_if_partial(); + assert!(!complete.lower_bound); + assert_eq!(complete.published_tokens(), Some(10)); +} + +#[test] +fn unpriced_model_names_are_recorded_but_not_serialized() { + let mut estimate = LocalCostEstimate::default(); + estimate.record_list_price(Some(" mystery-model "), None); + estimate.record_list_price(None, None); + estimate.record_list_price(Some("known"), Some(1.0)); + assert_eq!(estimate.coverage.unpriced, 2); + assert_eq!(estimate.coverage.estimated, 1); + assert_eq!( + estimate.unpriced_models.iter().collect::>(), + vec!["mystery-model"] + ); + let json = serde_json::to_value(&estimate).unwrap(); + assert!(json.get("unpricedModels").is_none()); +} From 5358eae0d5ec3f502e1fc7c083910752e871640e Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:38:21 +0700 Subject: [PATCH 120/259] Port upstream 0.61.0: show provider history provenance and token totals in CLI usage text --- CHANGELOG.md | 1 + rust/src/cli/usage/render.rs | 57 +++++++++++++- rust/src/cli/usage_tests.rs | 92 +++++++++++++++++++++++ rust/src/core/usage_snapshot.rs | 26 +++++++ rust/src/providers/openrouter/activity.rs | 62 ++++++++++++++- 5 files changed, 233 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c054768f2c..d9a864359b 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ menu-bar layout. - Grok: Settings and tray **Add account** flow matching Codex/Claude — isolated `grok login --oauth`, save current CLI login, switch, and remove without logging out the active session. - Replicate: cookie-authenticated monthly spend and optional prepaid credit balance from the billing page, with user and organization account isolation. - DeepSeek: show reported per-model spend in the provider details while preserving the billing currency, reporting period, zero values, and incomplete-total safeguards. +- CLI: show provider-supplied history in usage text with its source period, known zero, and cost provenance (reported, estimated, or includes estimates), plus input-plus-output token totals for OpenRouter Activity. ### Fixed - Claude: when Hide Personal Info is enabled, keep saved account rows distinguishable with stable localized `Account N` labels and matching redacted tooltips. diff --git a/rust/src/cli/usage/render.rs b/rust/src/cli/usage/render.rs index 62fd50310d..e78c5e8f2d 100644 --- a/rust/src/cli/usage/render.rs +++ b/rust/src/cli/usage/render.rs @@ -10,6 +10,7 @@ use crate::core::{ CostSnapshot, ProviderDisplayDetail, ProviderFetchResult, ProviderId, ProviderInventoryItem, RateWindow, UsagePace, UsageSnapshot, instantiate_provider, }; +use crate::spend_contract::CostProvenance; use crate::status::{ProviderStatus as StatusInfo, StatusLevel}; pub fn render_text_error(provider_id: ProviderId, error_msg: &str, use_color: bool) -> String { @@ -427,11 +428,12 @@ fn append_cost_line(lines: &mut Vec, cost: Option<&CostSnapshot>) { // Provider-supplied Activity history is a completed reporting window, // rather than the ordinary current-cost meter. Providers mark such - // snapshots `always_visible`; keep their source period and known zero - // visible in text output without adding a second generic cost line. The - // daily points remain available in the JSON cost payload. + // snapshots `always_visible`; keep their source period, known zero, token + // total and cost provenance visible in text output without adding a second + // generic cost line or recomputing totals from the daily points. The daily + // points remain available in the JSON cost payload. if cost.limit.is_none() && cost.always_visible { - lines.push(format!(" {}: {}", cost.period, cost.format_used())); + lines.push(format_history_line(cost)); return; } @@ -451,6 +453,53 @@ fn append_cost_line(lines: &mut Vec, cost: Option<&CostSnapshot>) { } } +/// One provider-history line: `: () · tokens`. +/// +/// Mirrors upstream `CLIRenderer.liveHistoryLine` (0.61.0). The spend and token +/// totals are the provider's own; nothing is summed from daily points. +fn format_history_line(cost: &CostSnapshot) -> String { + let spend = cost.format_used(); + let spend = match cost.provenance { + Some(CostProvenance::VendorMetered) => format!("{spend} (reported)"), + Some(CostProvenance::ListPriceEstimate) => format!("{spend} (estimated)"), + Some(CostProvenance::Mixed) => format!("{spend} (includes estimates)"), + Some(CostProvenance::Unknown) | None => spend, + }; + let mut values = vec![spend]; + if let Some(tokens) = cost.history_tokens { + let unit = if tokens == 1 { "token" } else { "tokens" }; + values.push(format!("{} {unit}", format_token_count(tokens))); + } + format!(" {}: {}", cost.period, values.join(" · ")) +} + +/// Compact token count: `999`, `1.2K`, `15K`, `2.5M`, `1B`. A unit is promoted +/// once the lower unit would round to 1000 (upstream `tokenCountString`). +fn format_token_count(tokens: u64) -> String { + const UNITS: [(u64, f64, &str); 3] = [ + (999_500_000, 1_000_000_000.0, "B"), + (999_500, 1_000_000.0, "M"), + (1_000, 1_000.0, "K"), + ]; + for (threshold, divisor, suffix) in UNITS { + if tokens >= threshold { + #[expect( + clippy::cast_precision_loss, + reason = "display rounding to at most two significant digits" + )] + let scaled = tokens as f64 / divisor; + let formatted = if scaled >= 10.0 { + format!("{scaled:.0}") + } else { + let one = format!("{scaled:.1}"); + one.strip_suffix(".0").unwrap_or(&one).to_string() + }; + return format!("{formatted}{suffix}"); + } + } + tokens.to_string() +} + /// Render usage as text (backwards compatible version) pub fn render_text(provider: ProviderId, result: &ProviderFetchResult, use_color: bool) -> String { render_text_with_status(provider, result, None, use_color) diff --git a/rust/src/cli/usage_tests.rs b/rust/src/cli/usage_tests.rs index d332131b81..44211d4028 100644 --- a/rust/src/cli/usage_tests.rs +++ b/rust/src/cli/usage_tests.rs @@ -324,3 +324,95 @@ fn json_inventory_is_additive_and_contains_no_redemption_token() { .contains("coupon-token-secret") ); } + +fn history_output(cost: CostSnapshot) -> String { + let result = fetch_result(UsageSnapshot::new(RateWindow::new(0.0))).with_cost(cost); + render_text_with_status(ProviderId::OpenRouter, &result, None, false) +} + +#[test] +fn history_line_shows_provenance_and_token_total() { + use crate::spend_contract::CostProvenance; + + let cases = [ + (CostProvenance::VendorMetered, "$1.25 (reported)"), + (CostProvenance::ListPriceEstimate, "$1.25 (estimated)"), + (CostProvenance::Mixed, "$1.25 (includes estimates)"), + (CostProvenance::Unknown, "$1.25"), + ]; + for (provenance, spend) in cases { + let output = history_output( + CostSnapshot::new(1.25, "USD", "Last 30 days (UTC)") + .with_history_tokens(15) + .with_provenance(provenance) + .always_visible(), + ); + let expected = format!(" Last 30 days (UTC): {spend} · 15 tokens"); + assert_eq!(output.matches(&expected).count(), 1, "{output}"); + } +} + +#[test] +fn history_line_preserves_known_zero_singular_token_and_currency() { + use crate::spend_contract::CostProvenance; + + let zero = history_output( + CostSnapshot::new(0.0, "USD", "Last 30 days (UTC)") + .with_history_tokens(0) + .with_provenance(CostProvenance::VendorMetered) + .always_visible(), + ); + assert!(zero.contains("Last 30 days (UTC): $0.00 (reported) · 0 tokens")); + + let one = history_output( + CostSnapshot::new(0.0, "USD", "Last 1 day") + .with_history_tokens(1) + .always_visible(), + ); + assert!(one.ends_with("Last 1 day: $0.00 · 1 token"), "{one}"); + + let euro = history_output(CostSnapshot::new(2.5, "EUR", "Billing period").always_visible()); + assert!(euro.contains("Billing period: €2.50")); + assert!(!euro.contains("token")); +} + +#[test] +fn history_line_compacts_large_token_totals() { + let cases = [ + (999, "999 tokens"), + (1_000, "1K tokens"), + (1_250, "1.2K tokens"), + (15_400, "15K tokens"), + (999_499, "999K tokens"), + (999_500, "1M tokens"), + (2_500_000, "2.5M tokens"), + (999_500_000, "1B tokens"), + ]; + for (tokens, expected) in cases { + let output = history_output( + CostSnapshot::new(1.0, "USD", "Last 30 days (UTC)") + .with_history_tokens(tokens) + .always_visible(), + ); + assert!( + output.contains(&format!("$1.00 · {expected}")), + "{tokens}: {output}" + ); + } +} + +#[test] +fn history_fields_do_not_change_the_cost_json_contract() { + use crate::spend_contract::CostProvenance; + + let result = fetch_result(UsageSnapshot::new(RateWindow::new(0.0))).with_cost( + CostSnapshot::new(1.25, "USD", "Last 30 days (UTC)") + .with_history_tokens(15) + .with_provenance(CostProvenance::VendorMetered) + .always_visible(), + ); + let json = render_json_result(ProviderId::OpenRouter, result, None); + let cost = json.get("cost").and_then(|cost| cost.as_object()).unwrap(); + assert!(!cost.contains_key("historyTokens") && !cost.contains_key("history_tokens")); + assert!(!cost.contains_key("provenance")); +} diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index ebbdcb8c20..09db5dfdaa 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize}; use super::ProviderDisplayDetail; use super::RateWindow; +use crate::spend_contract::CostProvenance; /// Subscription dates explicitly reported by an authenticated provider /// dashboard or subscription endpoint. /// @@ -414,6 +415,17 @@ pub struct CostSnapshot { /// usage signal and must remain visible when optional local summaries are hidden. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub always_visible: bool, + + /// Input-plus-output token total for the same reporting window as `used`, + /// when the provider supplies one. Text rendering only: it is not part of + /// the serialized cost payload, which keeps the ordinary usage JSON contract. + #[serde(skip)] + pub history_tokens: Option, + + /// How the `used` figure was produced, when the provider states it. Text + /// rendering only; not serialized. + #[serde(skip)] + pub provenance: Option, } impl CostSnapshot { @@ -432,6 +444,8 @@ impl CostSnapshot { account_id: None, daily: Vec::new(), always_visible: false, + history_tokens: None, + provenance: None, } } @@ -541,6 +555,18 @@ impl CostSnapshot { self } + /// Attach the provider-supplied token total for this spend window. + pub fn with_history_tokens(mut self, tokens: u64) -> Self { + self.history_tokens = Some(tokens); + self + } + + /// Record how the spend figure was produced (reported, estimated, mixed). + pub fn with_provenance(mut self, provenance: CostProvenance) -> Self { + self.provenance = Some(provenance); + self + } + /// Builder pattern: set currency symbol for localized rendering. pub fn with_currency_symbol(mut self, symbol: impl Into) -> Self { self.currency_symbol = Some(symbol.into()); diff --git a/rust/src/providers/openrouter/activity.rs b/rust/src/providers/openrouter/activity.rs index 9fcd2bec9f..67d6cf48e1 100644 --- a/rust/src/providers/openrouter/activity.rs +++ b/rust/src/providers/openrouter/activity.rs @@ -4,6 +4,7 @@ use chrono::{DateTime, Duration, Utc}; use serde_json::Value; use crate::core::{CostDailyPoint, CostSnapshot, ProviderError}; +use crate::spend_contract::CostProvenance; const MAX_ACTIVITY_ROWS: usize = 20_000; /// Distinct identity rows tracked for dedupe; bounds the `seen` map. @@ -19,6 +20,8 @@ pub(super) fn parse_activity_cost( let mut seen: HashMap = HashMap::new(); let mut daily: BTreeMap = BTreeMap::new(); let mut total = 0.0; + let mut estimated_total = 0.0; + let mut tokens = 0u64; let mut rows_seen = 0usize; for payload in payloads { @@ -139,16 +142,32 @@ pub(super) fn parse_activity_cost( )); } total += cost; + estimated_total += estimated; + // Prompt plus completion only: reasoning tokens are a separate counter. + tokens = tokens.saturating_add(prompt + completion); *daily.entry(day.to_string()).or_default() += cost; } } - if !total.is_finite() { + if !total.is_finite() || !estimated_total.is_finite() { return Err(ProviderError::Parse( "OpenRouter Activity spend overflowed".into(), )); } + // Matches upstream's plugin snapshot mapper: any BYOK estimate makes the + // window estimated, or mixed when metered spend is also present. + let provenance = if estimated_total > 0.0 { + if total - estimated_total > 0.0 { + CostProvenance::Mixed + } else { + CostProvenance::ListPriceEstimate + } + } else { + CostProvenance::VendorMetered + }; Ok(CostSnapshot::new(total, "USD", "Last 30 days (UTC)") + .with_history_tokens(tokens) + .with_provenance(provenance) .with_daily( daily .into_iter() @@ -314,4 +333,45 @@ mod tests { assert!(error.to_string().contains("completed UTC day")); } + + #[test] + fn records_token_total_and_cost_provenance() { + use crate::spend_contract::CostProvenance; + + let cases = [ + (1.25, 0.0, CostProvenance::VendorMetered), + (0.0, 0.75, CostProvenance::ListPriceEstimate), + (1.25, 0.75, CostProvenance::Mixed), + ]; + for (usage, byok, expected) in cases { + let history = serde_json::json!({"data":[ + {"date":"2026-08-17","model":"m","prompt_tokens":10,"completion_tokens":5, + "reasoning_tokens":2,"requests":1,"usage":usage,"byok_usage_inference":byok} + ]}); + let cost = parse_activity_cost(&[history], now()).unwrap(); + assert_eq!(cost.history_tokens, Some(15)); + assert_eq!(cost.provenance, Some(expected)); + assert_eq!(cost.used, usage + byok); + } + } + + #[test] + fn empty_activity_is_a_reported_zero_with_zero_tokens() { + use crate::spend_contract::CostProvenance; + + let cost = parse_activity_cost(&[serde_json::json!({"data":[]})], now()).unwrap(); + assert_eq!(cost.used, 0.0); + assert_eq!(cost.history_tokens, Some(0)); + assert_eq!(cost.provenance, Some(CostProvenance::VendorMetered)); + } + + #[test] + fn duplicate_rows_do_not_double_count_tokens() { + let row = serde_json::json!({"date":"2026-08-21","model":"m","prompt_tokens":10, + "completion_tokens":5,"requests":1,"usage":1.0}); + let history = serde_json::json!({"data":[row.clone()]}); + let latest = serde_json::json!({"data":[row]}); + let cost = parse_activity_cost(&[history, latest], now()).unwrap(); + assert_eq!(cost.history_tokens, Some(15)); + } } From fe4cbc7b59b2f8414ce2dbcd71fb39fe2acdecd1 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:40:21 +0700 Subject: [PATCH 121/259] Port upstream 0.60.4: keep Chrome DeepSeek balance through same-session transport failures --- .../src/commands/last_good_owner_tests.rs | 160 ++++++++ .../src-tauri/src/commands/mod.rs | 2 + .../src-tauri/src/commands/providers.rs | 83 +++- .../src-tauri/src/commands/tests.rs | 12 + apps/desktop-tauri/src-tauri/src/state.rs | 4 + rust/src/cli/diagnose.rs | 1 + rust/src/cli/hooks.rs | 1 + rust/src/core/last_good_owner.rs | 112 ++++++ rust/src/core/last_good_owner_tests.rs | 86 ++++ rust/src/core/mod.rs | 2 + rust/src/core/provider.rs | 11 + rust/src/core/provider_state.rs | 1 + rust/src/core/usage_snapshot.rs | 13 + rust/src/providers/deepseek/chrome_session.rs | 160 ++++++++ .../deepseek/chrome_session_tests.rs | 78 ++++ rust/src/providers/deepseek/mod.rs | 93 ++++- .../providers/deepseek/platform_balance.rs | 160 ++++++++ .../deepseek/platform_balance_tests.rs | 93 +++++ .../providers/deepseek/session_resolver.rs | 311 +++++++++++++++ .../deepseek/session_resolver_tests.rs | 366 ++++++++++++++++++ rust/src/providers/deepseek/tests.rs | 31 ++ 21 files changed, 1762 insertions(+), 18 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs create mode 100644 rust/src/core/last_good_owner.rs create mode 100644 rust/src/core/last_good_owner_tests.rs create mode 100644 rust/src/providers/deepseek/chrome_session.rs create mode 100644 rust/src/providers/deepseek/chrome_session_tests.rs create mode 100644 rust/src/providers/deepseek/platform_balance.rs create mode 100644 rust/src/providers/deepseek/platform_balance_tests.rs create mode 100644 rust/src/providers/deepseek/session_resolver.rs create mode 100644 rust/src/providers/deepseek/session_resolver_tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs b/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs new file mode 100644 index 0000000000..e9c7753d8c --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/last_good_owner_tests.rs @@ -0,0 +1,160 @@ +//! Owner-checked last-good retention in the provider refresh shell. +//! +//! A failure tied to a live session may keep the cached snapshot only when the +//! same session produced that snapshot. Scenarios use DeepSeek's Chrome +//! session balance, the provider that declares owned transport failures. + +use super::ProviderUsageSnapshot; +use super::providers::{preserve_last_good_transient_failure, record_last_good_owner}; +use crate::state::AppState; +use codexbar::core::{ + LastGoodOwner, ProviderError, ProviderFetchResult, ProviderId, ProviderStateKind, RateWindow, + UsageSnapshot, instantiate_provider, +}; + +const MEASURED_AT: &str = "2026-09-01T00:00:00Z"; + +fn owner(profile: &str, token: &str) -> LastGoodOwner { + LastGoodOwner::derive("deepseek-platform-balance", profile, token).expect("owner") +} + +fn cached_balance() -> ProviderUsageSnapshot { + let metadata = instantiate_provider(ProviderId::DeepSeek) + .metadata() + .clone(); + let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(0.0)), "web"); + let mut snapshot = + ProviderUsageSnapshot::from_fetch_result(ProviderId::DeepSeek, &metadata, &result, None); + snapshot.updated_at = MEASURED_AT.to_string(); + snapshot +} + +fn failed_refresh() -> ProviderUsageSnapshot { + let metadata = instantiate_provider(ProviderId::DeepSeek) + .metadata() + .clone(); + ProviderUsageSnapshot::from_error( + ProviderId::DeepSeek, + &metadata, + "Timeout".to_string(), + ProviderStateKind::Unknown, + ) +} + +fn state_with_cached_balance(cached_owner: Option) -> AppState { + let mut state = AppState::new(); + state.provider_cache.push(cached_balance()); + record_last_good_owner(&mut state, ProviderId::DeepSeek, cached_owner); + state +} + +fn refresh_failure(state: &mut AppState, error: &ProviderError) -> ProviderUsageSnapshot { + preserve_last_good_transient_failure(state, ProviderId::DeepSeek, failed_refresh(), error) +} + +#[test] +fn matching_session_keeps_the_balance_and_its_original_time() { + let session = owner("chrome:Default", "token-a"); + let mut state = state_with_cached_balance(Some(session.clone())); + + let kept = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(session)), + ); + + assert_eq!(kept.error, None); + assert_eq!(kept.updated_at, MEASURED_AT); +} + +#[test] +fn different_profile_or_token_shows_the_error() { + for failed_session in [ + owner("chrome:Profile 1", "token-a"), + owner("chrome:Default", "token-b"), + ] { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(failed_session)), + ); + assert_eq!(shown.error.as_deref(), Some("Timeout")); + } +} + +#[test] +fn failure_without_a_session_owner_fails_closed() { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + + let shown = refresh_failure(&mut state, &ProviderError::Timeout.with_failure_owner(None)); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn balance_without_an_owner_is_never_retained() { + // An API-key balance, a decoded cache, or a seeded snapshot has no owner. + let session = owner("chrome:Default", "token-a"); + let mut state = state_with_cached_balance(None); + + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(session)), + ); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn unattributed_transport_failure_does_not_retain_deepseek_balances() { + let mut state = state_with_cached_balance(Some(owner("chrome:Default", "token-a"))); + + let shown = refresh_failure(&mut state, &ProviderError::Timeout); + + assert_eq!(shown.error.as_deref(), Some("Timeout")); +} + +#[test] +fn fresh_snapshot_replaces_or_clears_the_owner() { + let first = owner("chrome:Default", "token-a"); + let second = owner("chrome:Default", "token-b"); + let mut state = state_with_cached_balance(Some(first.clone())); + + record_last_good_owner(&mut state, ProviderId::DeepSeek, Some(second.clone())); + assert_eq!( + state.last_good_owners.get(&ProviderId::DeepSeek), + Some(&second) + ); + let shown = refresh_failure( + &mut state, + &ProviderError::Timeout.with_failure_owner(Some(first)), + ); + assert_eq!(shown.error.as_deref(), Some("Timeout")); + + record_last_good_owner(&mut state, ProviderId::DeepSeek, None); + assert!(state.last_good_owners.is_empty()); +} + +#[test] +fn other_providers_keep_their_existing_failure_policy() { + let metadata = instantiate_provider(ProviderId::Codex).metadata().clone(); + let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "OAuth"); + let good = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); + let failed = ProviderUsageSnapshot::from_error( + ProviderId::Codex, + &metadata, + "Timeout".to_string(), + ProviderStateKind::Unknown, + ); + let mut state = AppState::new(); + state.provider_cache.push(good); + + let kept = preserve_last_good_transient_failure( + &mut state, + ProviderId::Codex, + failed, + &ProviderError::Timeout, + ); + + assert_eq!(kept.error, None); +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index fdfaeda967..e1798ee433 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -75,6 +75,8 @@ pub use surface::*; pub use system::*; pub(crate) use usage_items::*; +#[cfg(test)] +mod last_good_owner_tests; #[cfg(test)] mod tests; diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..7c55e07d07 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -39,6 +39,7 @@ pub(crate) fn invalidate_account_usage( state.is_refreshing = false; state.provider_refresh_started_at = None; state.transient_provider_failure_counts.remove(&id); + state.last_good_owners.remove(&id); state.auto_resume.clear_provider(id); state .provider_cache @@ -334,6 +335,9 @@ pub(crate) fn invalidate_provider_refresh_and_prune_disabled( guard .transient_provider_failure_counts .retain(|id, _| enabled_ids.contains(id)); + guard + .last_good_owners + .retain(|id, _| enabled_ids.contains(id)); guard .provider_cache_updated_at_by_provider .retain(|id, _| enabled_ids.contains(id)); @@ -565,7 +569,7 @@ async fn refresh_provider( token_account_id: Option, hooks_enabled: bool, ) { - let (snapshot, account_identity, failure_policy) = + let (snapshot, account_identity, retention) = fetch_provider_snapshot(id, ctx, token_account_id).await; let fresh_snapshot = snapshot.error.is_none(); @@ -584,8 +588,12 @@ async fn refresh_provider( &mut guard, id, snapshot, - failure_policy, + retention.policy, + &retention.failure_ownership, ); + if fresh_snapshot { + record_last_good_owner(&mut guard, id, retention.fresh_owner); + } // F6 (upstream 0.48.0): backfill missing reset timestamps from the // cached snapshot before persisting and publishing. let cached = guard @@ -762,7 +770,30 @@ pub(super) fn preserve_last_good_transient_failure( error: &codexbar::core::ProviderError, ) -> ProviderUsageSnapshot { let policy = instantiate_provider(id).last_good_failure_policy_for_error(error); - preserve_last_good_transient_failure_with_policy(guard, id, snapshot, Some(policy)) + preserve_last_good_transient_failure_with_policy( + guard, + id, + snapshot, + Some(policy), + &error.failure_ownership(), + ) +} + +/// Remember which live session produced the fresh snapshot now cached for +/// `id`, or forget any earlier owner when the fresh snapshot has none. +pub(super) fn record_last_good_owner( + guard: &mut AppState, + id: ProviderId, + owner: Option, +) { + match owner { + Some(owner) => { + guard.last_good_owners.insert(id, owner); + } + None => { + guard.last_good_owners.remove(&id); + } + } } fn preserve_last_good_transient_failure_with_policy( @@ -770,6 +801,7 @@ fn preserve_last_good_transient_failure_with_policy( id: ProviderId, snapshot: ProviderUsageSnapshot, policy: Option, + ownership: &codexbar::core::FailureOwnership, ) -> ProviderUsageSnapshot { let Some(error) = snapshot.error.as_deref() else { guard.transient_provider_failure_counts.remove(&id); @@ -781,6 +813,12 @@ fn preserve_last_good_transient_failure_with_policy( guard.transient_provider_failure_counts.remove(&id); return snapshot; } + // A failure tied to a live session may keep only a snapshot that the same + // session produced. Anything else shows the error. + if !ownership.allows_retention(guard.last_good_owners.get(&id)) { + guard.transient_provider_failure_counts.remove(&id); + return snapshot; + } let Some(mut previous) = guard .provider_cache @@ -841,20 +879,28 @@ fn preserve_last_good_transient_failure_with_policy( } } +/// What a refresh tells the shell about keeping or replacing the last good +/// snapshot. +#[derive(Default)] +struct RefreshRetention { + /// Failure handling when a prior good snapshot exists. + policy: Option, + /// Session the failed request belonged to, when the provider can tell. + failure_ownership: codexbar::core::FailureOwnership, + /// Session that produced a fresh snapshot. + fresh_owner: Option, +} + async fn fetch_provider_snapshot( id: ProviderId, ctx: FetchContext, token_account_id: Option, -) -> ( - ProviderUsageSnapshot, - Option, - Option, -) { +) -> (ProviderUsageSnapshot, Option, RefreshRetention) { let provider = instantiate_provider(id); let metadata = provider.metadata().clone(); let started = std::time::Instant::now(); - let (mut snapshot, account_identity, failure_policy) = + let (mut snapshot, account_identity, retention) = match tokio::time::timeout(provider_fetch_timeout(id, &ctx), provider.fetch_usage(&ctx)) .await { @@ -868,7 +914,10 @@ async fn fetch_provider_snapshot( token_account_id, ), account_identity, - None, + RefreshRetention { + fresh_owner: result.last_good_owner.clone(), + ..RefreshRetention::default() + }, ) } Ok(Err(e)) => { @@ -881,7 +930,11 @@ async fn fetch_provider_snapshot( provider.error_state_kind(&e), ), None, - Some(policy), + RefreshRetention { + policy: Some(policy), + failure_ownership: e.failure_ownership(), + fresh_owner: None, + }, ) } Err(_) => { @@ -895,13 +948,17 @@ async fn fetch_provider_snapshot( provider.error_state_kind(&error), ), None, - Some(policy), + RefreshRetention { + policy: Some(policy), + failure_ownership: error.failure_ownership(), + fresh_owner: None, + }, ) } }; record_provider_fetch_duration(id, &mut snapshot, started); - (snapshot, account_identity, failure_policy) + (snapshot, account_identity, retention) } fn record_provider_fetch_duration( diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index d25a0a064a..433efc9474 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -948,6 +948,7 @@ fn usage_item_descriptors_keep_raw_ids_and_redact_titles() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1172,6 +1173,7 @@ fn provider_cache_upsert_replaces_existing_provider() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut first = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1200,6 +1202,7 @@ fn provider_cache_prunes_disabled_providers() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let codex = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1236,6 +1239,7 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1273,6 +1277,7 @@ fn codex_transient_transport_failure_helper_uses_typed_policy() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Codex, &metadata, &result, None); @@ -1309,6 +1314,7 @@ fn claude_repeated_auth_failure_surfaces_error() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1352,6 +1358,7 @@ fn claude_cloudflare_challenge_retains_prior_usage_while_surfaceing_guidance() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1406,6 +1413,7 @@ fn claude_cloudflare_challenge_keeps_prior_usage_when_guidance_surfaces() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let mut good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1457,6 +1465,7 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() { has_successful_claude_cli_quota: true, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1504,6 +1513,7 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let good = ProviderUsageSnapshot::from_fetch_result(ProviderId::Claude, &metadata, &result, None); @@ -1708,6 +1718,7 @@ fn japanese_provider_snapshot_localizes_weekly_label() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let snapshot = @@ -1742,6 +1753,7 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, }; let snapshot = diff --git a/apps/desktop-tauri/src-tauri/src/state.rs b/apps/desktop-tauri/src-tauri/src/state.rs index 10d626ec43..289ab89d55 100644 --- a/apps/desktop-tauri/src-tauri/src/state.rs +++ b/apps/desktop-tauri/src-tauri/src/state.rs @@ -121,6 +121,9 @@ pub struct AppState { pub tray_anchor: Option, pub provider_cache: Vec, pub transient_provider_failure_counts: HashMap, + /// Live session behind each provider's cached good snapshot, for + /// owner-checked last-good retention. In memory only. + pub last_good_owners: HashMap, pub provider_cache_updated_at: Option, /// Per-provider freshness for scoped background refreshes. The aggregate /// timestamp cannot tell a dedicated auto-resume watcher whether its own @@ -193,6 +196,7 @@ impl AppState { tray_anchor: None, provider_cache: Vec::new(), transient_provider_failure_counts: HashMap::new(), + last_good_owners: HashMap::new(), provider_cache_updated_at: None, provider_cache_updated_at_by_provider: HashMap::new(), provider_refresh_started_at: None, diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 64d0b54d59..f26a261f1a 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -400,6 +400,7 @@ fn error_category(err: &ProviderError) -> &'static str { | ProviderError::NoCookies => "auth", ProviderError::OAuthTransient(_) => "api", ProviderError::Network(_) | ProviderError::Timeout => "network", + ProviderError::OwnedTransport { source, .. } => error_category(source), ProviderError::NotInstalled(_) | ProviderError::UnsupportedSource(_) => "config", ProviderError::Parse(_) => "parse", ProviderError::Other(message) => { diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 2f26498cf1..6c0d1f616c 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -409,6 +409,7 @@ fn hook_refresh_failure_status(error: &ProviderError) -> String { "network_error".into() } } + ProviderError::OwnedTransport { source, .. } => hook_refresh_failure_status(source), ProviderError::NotInstalled(_) => "error".into(), ProviderError::Parse(_) | ProviderError::UnsupportedSource(_) | ProviderError::Other(_) => { "error".into() diff --git a/rust/src/core/last_good_owner.rs b/rust/src/core/last_good_owner.rs new file mode 100644 index 0000000000..dbc0433dcb --- /dev/null +++ b/rust/src/core/last_good_owner.rs @@ -0,0 +1,112 @@ +//! Proof of which live session supplied a cached provider snapshot. +//! +//! A provider that reads a balance through a browser session can attach a +//! [`LastGoodOwner`] to its fetch result. The shell keeps that owner in memory +//! next to the cached snapshot and retains the snapshot through a transport +//! failure only when the failed request came from the same owner. The owner is +//! a salted digest of the session scope and secret, is never serialized, and +//! never appears in `Debug` output. + +use sha2::{Digest, Sha256}; + +use super::ProviderError; + +const OWNER_NAMESPACE: &str = "com.codexbar.last-good-owner.v1"; + +/// Opaque, in-memory identity of the session that produced a snapshot. +#[derive(Clone, PartialEq, Eq)] +pub struct LastGoodOwner(String); + +impl LastGoodOwner { + /// Derive an owner from a provider namespace, a scope such as a browser + /// profile id, and the session secret. Returns `None` when the scope or the + /// secret is empty after trimming, so an unidentified session cannot own a + /// snapshot. + pub fn derive(namespace: &str, scope: &str, secret: &str) -> Option { + let scope = scope.trim(); + let secret = secret.trim(); + if scope.is_empty() || secret.is_empty() { + return None; + } + let mut hasher = Sha256::new(); + for part in [OWNER_NAMESPACE, namespace, scope, secret] { + hasher.update(part.as_bytes()); + hasher.update([0u8]); + } + let digest = hasher.finalize(); + Some(Self( + digest.iter().map(|byte| format!("{byte:02x}")).collect(), + )) + } +} + +impl std::fmt::Debug for LastGoodOwner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("LastGoodOwner()") + } +} + +/// Which session a failed refresh came from, as far as the provider can prove. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub enum FailureOwnership { + /// The error carries no session identity; the provider's policy decides. + #[default] + Unchecked, + /// The failed request used exactly this session. + Owned(LastGoodOwner), + /// The failure has no attributable session (for example a resolution + /// deadline, or a session already rejected). Retention must fail closed. + Unowned, +} + +impl FailureOwnership { + /// Whether a cached snapshot owned by `cached` may be shown in place of + /// this failure. + pub fn allows_retention(&self, cached: Option<&LastGoodOwner>) -> bool { + match self { + Self::Unchecked => true, + Self::Owned(owner) => cached == Some(owner), + Self::Unowned => false, + } + } +} + +impl ProviderError { + /// Wrap a transport failure with the session that produced it. + /// + /// A failure that is not a transport failure is returned unchanged, because + /// only transport failures are eligible for owner-checked retention. + pub fn with_failure_owner(self, owner: Option) -> Self { + if self.is_transport_failure() { + Self::OwnedTransport { + owner, + source: Box::new(self), + } + } else { + self + } + } + + /// Session ownership of this failure for last-good retention. + pub fn failure_ownership(&self) -> FailureOwnership { + match self { + Self::OwnedTransport { + owner: Some(owner), .. + } => FailureOwnership::Owned(owner.clone()), + Self::OwnedTransport { owner: None, .. } => FailureOwnership::Unowned, + _ => FailureOwnership::Unchecked, + } + } + + /// The underlying error with any session-ownership wrapper removed. + pub fn without_failure_owner(&self) -> &Self { + match self { + Self::OwnedTransport { source, .. } => source.without_failure_owner(), + other => other, + } + } +} + +#[cfg(test)] +#[path = "last_good_owner_tests.rs"] +mod tests; diff --git a/rust/src/core/last_good_owner_tests.rs b/rust/src/core/last_good_owner_tests.rs new file mode 100644 index 0000000000..525899f018 --- /dev/null +++ b/rust/src/core/last_good_owner_tests.rs @@ -0,0 +1,86 @@ +use super::*; + +fn owner(scope: &str, secret: &str) -> LastGoodOwner { + LastGoodOwner::derive("test", scope, secret).expect("owner") +} + +#[test] +fn owner_requires_scope_and_secret() { + assert!(LastGoodOwner::derive("test", " ", "secret").is_none()); + assert!(LastGoodOwner::derive("test", "chrome:Default", " \n").is_none()); +} + +#[test] +fn owner_trims_and_separates_scope_from_secret() { + assert_eq!( + owner(" chrome:Default ", " token "), + owner("chrome:Default", "token") + ); + assert_ne!( + owner("chrome:Default", "token"), + owner("chrome:Profile 1", "token") + ); + assert_ne!( + owner("chrome:Default", "token"), + owner("chrome:Default", "token2") + ); + assert_ne!(owner("ab", "c"), owner("a", "bc")); +} + +#[test] +fn owner_debug_output_never_contains_digest_or_secret() { + let rendered = format!("{:?}", owner("chrome:Default", "super-secret-token")); + assert_eq!(rendered, "LastGoodOwner()"); +} + +#[test] +fn only_transport_failures_are_wrapped() { + let wrapped = ProviderError::Timeout.with_failure_owner(Some(owner("p", "t"))); + assert!(matches!(wrapped, ProviderError::OwnedTransport { .. })); + assert!(wrapped.is_transport_failure()); + + let parse = ProviderError::Parse("bad".into()).with_failure_owner(Some(owner("p", "t"))); + assert!(matches!(parse, ProviderError::Parse(_))); +} + +#[test] +fn ownership_of_wrapped_and_plain_errors() { + let a = owner("p", "t"); + assert_eq!( + ProviderError::Timeout + .with_failure_owner(Some(a.clone())) + .failure_ownership(), + FailureOwnership::Owned(a) + ); + assert_eq!( + ProviderError::Timeout + .with_failure_owner(None) + .failure_ownership(), + FailureOwnership::Unowned + ); + assert_eq!( + ProviderError::Timeout.failure_ownership(), + FailureOwnership::Unchecked + ); +} + +#[test] +fn retention_requires_matching_owner_unless_unchecked() { + let a = owner("p", "t"); + let b = owner("p", "other"); + assert!(FailureOwnership::Unchecked.allows_retention(None)); + assert!(FailureOwnership::Owned(a.clone()).allows_retention(Some(&a))); + assert!(!FailureOwnership::Owned(a.clone()).allows_retention(Some(&b))); + assert!(!FailureOwnership::Owned(a).allows_retention(None)); + assert!(!FailureOwnership::Unowned.allows_retention(Some(&b))); +} + +#[test] +fn wrapper_displays_the_underlying_message() { + let wrapped = ProviderError::Timeout.with_failure_owner(None); + assert_eq!(wrapped.to_string(), "Timeout"); + assert!(matches!( + wrapped.without_failure_owner(), + ProviderError::Timeout + )); +} diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index 436c7d5e0a..510501ad4b 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -14,6 +14,7 @@ mod hooks; mod http; mod http_proxy; mod jsonl_scanner; +mod last_good_owner; mod models_dev_pricing; mod openai_dashboard; mod provider; @@ -42,6 +43,7 @@ pub use hooks::*; pub use http::*; pub use http_proxy::*; pub use jsonl_scanner::*; +pub use last_good_owner::{FailureOwnership, LastGoodOwner}; pub use models_dev_pricing::*; pub use openai_dashboard::*; pub use provider::*; diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..f9b052f086 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use thiserror::Error; use super::ProviderFetchResult; +use super::last_good_owner::LastGoodOwner; use super::provider_state::ProviderStateKind; /// Unique identifier for a provider @@ -653,6 +654,15 @@ pub enum ProviderError { #[error("{0}")] Other(String), + + /// A transport failure tagged with the session that produced it, so the + /// shell can retain a cached snapshot only for the same session. Build it + /// with [`ProviderError::with_failure_owner`]. + #[error("{source}")] + OwnedTransport { + owner: Option, + source: Box, + }, } impl ProviderError { @@ -664,6 +674,7 @@ impl ProviderError { ReqwestFailureClass::Timeout | ReqwestFailureClass::Connect ), ProviderError::Timeout => true, + ProviderError::OwnedTransport { source, .. } => source.is_transport_failure(), _ => false, } } diff --git a/rust/src/core/provider_state.rs b/rust/src/core/provider_state.rs index a4087e1ca5..7c6cd97a15 100644 --- a/rust/src/core/provider_state.rs +++ b/rust/src/core/provider_state.rs @@ -72,6 +72,7 @@ impl ProviderError { | ProviderError::Parse(_) | ProviderError::UnsupportedSource(_) | ProviderError::Other(_) => ProviderStateKind::Unknown, + ProviderError::OwnedTransport { source, .. } => source.state_kind(), } } } diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index ebbdcb8c20..9612832ea6 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -3,6 +3,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use super::LastGoodOwner; use super::ProviderDisplayDetail; use super::RateWindow; /// Subscription dates explicitly reported by an authenticated provider @@ -643,6 +644,11 @@ pub struct ProviderFetchResult { /// actions. It never crosses the frontend bridge. #[serde(skip)] pub account_identity: Option, + + /// Live session that supplied this result, used only for owner-checked + /// last-good retention. It is in memory only and never crosses the bridge. + #[serde(skip)] + pub last_good_owner: Option, } fn default_pace_authoritative() -> bool { @@ -662,9 +668,16 @@ impl ProviderFetchResult { has_successful_claude_cli_quota: false, pace_authoritative: true, account_identity: None, + last_good_owner: None, } } + /// Record which live session supplied this result. + pub fn with_last_good_owner(mut self, owner: Option) -> Self { + self.last_good_owner = owner; + self + } + /// Attach the provider's stable account identity without exposing it to /// serialized UI payloads. pub fn with_account_identity(mut self, account_identity: impl Into) -> Self { diff --git a/rust/src/providers/deepseek/chrome_session.rs b/rust/src/providers/deepseek/chrome_session.rs new file mode 100644 index 0000000000..098ee91be7 --- /dev/null +++ b/rust/src/providers/deepseek/chrome_session.rs @@ -0,0 +1,160 @@ +//! Import DeepSeek Platform sessions from Chrome's local storage. +//! +//! The platform dashboard keeps the signed-in session as a `userToken` item in +//! `localStorage` for `https://platform.deepseek.com`. Each Chrome profile holds +//! its own copy, so a candidate is identified by `chrome:`. + +use std::fmt; + +use serde_json::Value; + +use crate::browser::detection::{BrowserDetector, BrowserType}; +use crate::browser::leveldb::local_storage::{local_storage_dir, read_local_storage_entries}; + +const PLATFORM_ORIGIN: &str = "https://platform.deepseek.com"; +const TOKEN_KEY: &str = "userToken"; +const TOKEN_FIELDS: [&str; 5] = ["value", "token", "access_token", "accessToken", "userToken"]; +const MIN_TOKEN_CHARS: usize = 20; + +/// Environment variable that pins the Chrome profile to read when several +/// profiles hold a DeepSeek session. It accepts `chrome:` or a +/// profile directory path. +pub(super) const PROFILE_ID_ENV: &str = "CODEXBAR_DEEPSEEK_PROFILE_ID"; + +/// A session token found in one Chrome profile. +#[derive(Clone, PartialEq, Eq)] +pub(super) struct TokenInfo { + pub(super) id: String, + pub(super) token: String, + pub(super) label: String, +} + +impl fmt::Debug for TokenInfo { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TokenInfo") + .field("id", &self.id) + .field("token", &"") + .field("label", &self.label) + .finish() + } +} + +/// Which profile the user asked for, if any. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(super) struct ProfileSelection { + pub(super) profile_id: Option, +} + +impl ProfileSelection { + pub(super) fn from_env() -> Self { + Self::from_value(std::env::var(PROFILE_ID_ENV).ok().as_deref()) + } + + pub(super) fn from_value(value: Option<&str>) -> Self { + Self { + profile_id: value + .map(canonical_profile_id) + .filter(|profile| !profile.is_empty()), + } + } +} + +/// Normalize a profile setting: a directory path becomes `chrome:`. +pub(super) fn canonical_profile_id(raw: &str) -> String { + let value = raw.trim(); + let is_path = value.starts_with('/') || value.contains('\\') || value.contains(":/"); + if !is_path { + return value.to_string(); + } + let leaf = value + .trim_end_matches(['/', '\\']) + .rsplit(['/', '\\']) + .next() + .unwrap_or_default(); + if leaf.is_empty() { + value.to_string() + } else { + format!("chrome:{leaf}") + } +} + +/// Read the DeepSeek session token of every Chrome profile that has one. +/// +/// Blocking: reads LevelDB files from disk. Never logs token values. +pub(super) fn import_tokens() -> Vec { + let Some(chrome) = BrowserDetector::detect(BrowserType::Chrome) else { + tracing::debug!("deepseek chrome session: Chrome profiles not found"); + return Vec::new(); + }; + + let mut tokens = Vec::new(); + for profile in &chrome.profiles { + let id = format!("chrome:{}", profile.name); + let dir = local_storage_dir(&profile.path); + let entries = match read_local_storage_entries(&dir, PLATFORM_ORIGIN) { + Ok(entries) => entries, + Err(error) => { + tracing::debug!(profile = %id, %error, "deepseek chrome session: local storage unreadable"); + continue; + } + }; + let Some(token) = entries + .iter() + .find(|entry| entry.key == TOKEN_KEY) + .and_then(|entry| extract_user_token(&entry.value)) + else { + tracing::debug!(profile = %id, "deepseek chrome session: no userToken"); + continue; + }; + tracing::debug!(profile = %id, "deepseek chrome session: found userToken"); + tokens.push(TokenInfo { + label: format!("Google Chrome {}", profile.name), + id, + token, + }); + } + tokens.sort_by(|a, b| a.id.cmp(&b.id)); + tokens +} + +/// Extract the token from a `userToken` value: a JSON object carrying the +/// token in one of the known fields, or a bare (optionally quoted) token. +pub(super) fn extract_user_token(raw: &str) -> Option { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return None; + } + + if let Ok(value @ (Value::Object(_) | Value::Array(_))) = serde_json::from_str::(trimmed) + { + return token_from_json(&value); + } + + let unquoted = if (trimmed.starts_with('"') && trimmed.ends_with('"')) + || (trimmed.starts_with('\'') && trimmed.ends_with('\'')) + { + trimmed + .get(1..trimmed.len().saturating_sub(1)) + .unwrap_or_default() + } else { + trimmed + }; + is_plausible_token(unquoted).then(|| unquoted.to_string()) +} + +fn token_from_json(value: &Value) -> Option { + let object = value.as_object()?; + TOKEN_FIELDS.iter().find_map(|field| { + let token = object.get(*field)?.as_str()?; + is_plausible_token(token).then(|| token.to_string()) + }) +} + +fn is_plausible_token(value: &str) -> bool { + let trimmed = value.trim(); + trimmed.chars().count() >= MIN_TOKEN_CHARS && !trimmed.chars().any(char::is_whitespace) +} + +#[cfg(test)] +#[path = "chrome_session_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/chrome_session_tests.rs b/rust/src/providers/deepseek/chrome_session_tests.rs new file mode 100644 index 0000000000..98f4b4ba9b --- /dev/null +++ b/rust/src/providers/deepseek/chrome_session_tests.rs @@ -0,0 +1,78 @@ +use super::*; + +const TOKEN: &str = "abcdefghijklmnopqrstuvwxyz0123456789"; + +#[test] +fn extracts_token_from_json_object_fields() { + for field in ["value", "token", "access_token", "accessToken", "userToken"] { + let raw = format!(r#"{{"{field}":"{TOKEN}","__version":"0"}}"#); + assert_eq!(extract_user_token(&raw).as_deref(), Some(TOKEN), "{field}"); + } +} + +#[test] +fn json_object_skips_implausible_fields_and_prefers_field_order() { + let raw = format!(r#"{{"value":"short","token":"{TOKEN}"}}"#); + assert_eq!(extract_user_token(&raw).as_deref(), Some(TOKEN)); + assert_eq!(extract_user_token(r#"{"value":"short"}"#), None); + assert_eq!(extract_user_token(r#"{"value":42}"#), None); +} + +#[test] +fn extracts_bare_and_quoted_tokens() { + assert_eq!(extract_user_token(TOKEN).as_deref(), Some(TOKEN)); + assert_eq!( + extract_user_token(&format!(" \"{TOKEN}\"\n")).as_deref(), + Some(TOKEN) + ); + assert_eq!( + extract_user_token(&format!("'{TOKEN}'")).as_deref(), + Some(TOKEN) + ); +} + +#[test] +fn rejects_empty_short_and_spaced_values() { + assert_eq!(extract_user_token(""), None); + assert_eq!(extract_user_token(" "), None); + assert_eq!(extract_user_token("short-token"), None); + assert_eq!(extract_user_token("abcdefghij klmnopqrstuvwxyz"), None); + assert_eq!(extract_user_token("[\"abcdefghijklmnopqrstuvwxyz\"]"), None); +} + +#[test] +fn canonical_profile_id_normalizes_paths_only() { + assert_eq!(canonical_profile_id(" chrome:Default "), "chrome:Default"); + assert_eq!( + canonical_profile_id("C:\\Users\\me\\AppData\\Local\\Google\\Chrome\\User Data\\Profile 1"), + "chrome:Profile 1" + ); + assert_eq!( + canonical_profile_id("/home/me/.config/google-chrome/Default/"), + "chrome:Default" + ); +} + +#[test] +fn profile_selection_ignores_blank_values() { + assert_eq!(ProfileSelection::from_value(None).profile_id, None); + assert_eq!(ProfileSelection::from_value(Some(" ")).profile_id, None); + assert_eq!( + ProfileSelection::from_value(Some("chrome:Profile 2")) + .profile_id + .as_deref(), + Some("chrome:Profile 2") + ); +} + +#[test] +fn token_info_debug_redacts_the_token() { + let info = TokenInfo { + id: "chrome:Default".into(), + token: TOKEN.into(), + label: "Google Chrome Default".into(), + }; + let rendered = format!("{info:?}"); + assert!(!rendered.contains(TOKEN)); + assert!(rendered.contains("chrome:Default")); +} diff --git a/rust/src/providers/deepseek/mod.rs b/rust/src/providers/deepseek/mod.rs index 4461944d77..1a684a2071 100644 --- a/rust/src/providers/deepseek/mod.rs +++ b/rust/src/providers/deepseek/mod.rs @@ -1,6 +1,9 @@ //! DeepSeek provider implementation. //! -//! Fetches API account balance from DeepSeek's `/user/balance` endpoint. +//! Fetches API account balance from DeepSeek's `/user/balance` endpoint. When +//! no API key is configured, Auto mode reads the Platform balance through a +//! signed-in Chrome session and keeps that balance visible through temporary +//! connection failures of the same session. use async_trait::async_trait; use reqwest::Client; @@ -8,16 +11,31 @@ use std::collections::{HashMap, HashSet}; use serde::Deserialize; +mod chrome_session; +mod platform_balance; pub mod pricing; +mod session_resolver; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + CostSnapshot, FetchContext, LastGoodFailurePolicy, Provider, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; const DEEPSEEK_API_BASE: &str = "https://api.deepseek.com"; const DEEPSEEK_CREDENTIAL_TARGET: &str = "codexbar-deepseek"; +/// Upper bound for importing Chrome sessions and validating them. Staying +/// under the shell's fetch timeout keeps a slow resolution attributable: it +/// fails with no session owner instead of a shell-level timeout. +const CHROME_SESSION_DEADLINE: std::time::Duration = std::time::Duration::from_secs(28); + +/// Validation results outlive a single provider instance, which the shell +/// recreates for every refresh. +static CHROME_SESSION_CACHE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + session_resolver::ValidationCache::new(session_resolver::VALIDITY_TTL) + }); + #[derive(Debug, Deserialize)] struct BalanceResponse { #[serde(default)] @@ -203,7 +221,57 @@ impl DeepSeekProvider { ctx: &FetchContext, ) -> Result { let api_key = Self::get_api_key(ctx.api_key.as_deref())?; + self.fetch_usage_with_api_key(&api_key).await + } + /// Balance from a signed-in Chrome session on platform.deepseek.com, used + /// when no API key is configured. Reads the balance only. + async fn fetch_usage_chrome_session( + &self, + missing_key: ProviderError, + ) -> Result { + let selection = chrome_session::ProfileSelection::from_env(); + let client = self.client.clone(); + let resolution = tokio::time::timeout(CHROME_SESSION_DEADLINE, async { + let candidates = tokio::task::spawn_blocking(chrome_session::import_tokens) + .await + .unwrap_or_default(); + session_resolver::resolve(&candidates, &selection, &CHROME_SESSION_CACHE, |token| { + let client = client.clone(); + async move { platform_balance::fetch_platform_balance(&client, &token).await } + }) + .await + }) + .await; + + match resolution { + // A deadline has no attributable session, so it never retains a balance. + Err(_) => Err(ProviderError::Timeout.with_failure_owner(None)), + Ok(session_resolver::Resolution::Balance { balance, owner }) => Ok( + ProviderFetchResult::new(Self::snapshot_from_balance(*balance), "web") + .with_last_good_owner(owner), + ), + Ok(session_resolver::Resolution::SessionRequired) => Err(match missing_key { + ProviderError::NotInstalled(message) => ProviderError::NotInstalled(format!( + "{message} Or sign in to platform.deepseek.com in Chrome." + )), + other => other, + }), + Ok(session_resolver::Resolution::SelectionRequired(profiles)) => { + Err(ProviderError::Other(format!( + "DeepSeek is signed in on several Chrome profiles ({}). Set {} to choose one.", + profiles.join(", "), + chrome_session::PROFILE_ID_ENV, + ))) + } + Ok(session_resolver::Resolution::Failed(error)) => Err(error), + } + } + + async fn fetch_usage_with_api_key( + &self, + api_key: &str, + ) -> Result { let resp = self .client .get(format!("{DEEPSEEK_API_BASE}/user/balance")) @@ -228,7 +296,7 @@ impl DeepSeekProvider { let mut usage = Self::snapshot_from_balance(balance); let mut result = ProviderFetchResult::new(usage.clone(), "api"); - if let Ok(summary) = self.fetch_usage_summary(&api_key).await { + if let Ok(summary) = self.fetch_usage_summary(api_key).await { usage = Self::apply_usage_summary(usage, &summary); result = ProviderFetchResult::new(usage, "api"); if summary.month_cost > 0.0 || !summary.model_costs.is_empty() { @@ -501,7 +569,11 @@ impl Provider for DeepSeekProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => self.fetch_usage_api(ctx).await, + SourceMode::Auto => match Self::get_api_key(ctx.api_key.as_deref()) { + Ok(api_key) => self.fetch_usage_with_api_key(&api_key).await, + Err(missing_key) => self.fetch_usage_chrome_session(missing_key).await, + }, + SourceMode::OAuth => self.fetch_usage_api(ctx).await, SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) } @@ -511,6 +583,17 @@ impl Provider for DeepSeekProvider { fn available_sources(&self) -> Vec { vec![SourceMode::Auto, SourceMode::OAuth] } + + /// Only a transport failure attributed to a Chrome session keeps the last + /// balance; the shell then checks that the session matches the cached one. + /// API-key balances and every other failure surface normally. + fn last_good_failure_policy_for_error(&self, error: &ProviderError) -> LastGoodFailurePolicy { + if matches!(error, ProviderError::OwnedTransport { .. }) { + LastGoodFailurePolicy::Preserve + } else { + LastGoodFailurePolicy::Replace + } + } } fn parse_money(value: &str) -> f64 { diff --git a/rust/src/providers/deepseek/platform_balance.rs b/rust/src/providers/deepseek/platform_balance.rs new file mode 100644 index 0000000000..27e0ca0ec7 --- /dev/null +++ b/rust/src/providers/deepseek/platform_balance.rs @@ -0,0 +1,160 @@ +//! DeepSeek Platform balance read through a signed-in web session. +//! +//! The platform dashboard exposes the account wallets at +//! `GET /api/v0/users/get_user_summary`, authorized with the session's +//! `userToken`. This module only reads the balance; it never reads token-level +//! usage or cost. + +use std::collections::BTreeMap; + +use reqwest::{Client, StatusCode}; +use serde::Deserialize; + +use super::{BalanceInfo, BalanceResponse, FlexibleF64, FlexibleI64, select_balance_info}; +use crate::core::ProviderError; + +const PLATFORM_USER_SUMMARY_URL: &str = + "https://platform.deepseek.com/api/v0/users/get_user_summary"; + +/// Envelope error codes the platform uses for a missing or expired session. +const AUTH_ERROR_CODES: [i64; 2] = [40002, 40003]; + +/// Top-level envelope. Error envelopes are not schema-stable, so `data` stays +/// raw until the envelope code has been checked. +#[derive(Debug, Deserialize)] +struct UserSummaryResponse { + code: Option, + data: Option, +} + +#[derive(Debug, Deserialize)] +struct UserSummaryData { + biz_code: Option, + biz_data: Option, +} + +#[derive(Debug, Deserialize)] +struct UserSummary { + #[serde(default)] + normal_wallets: Vec, + #[serde(default)] + bonus_wallets: Vec, +} + +#[derive(Debug, Deserialize)] +struct Wallet { + currency: String, + balance: FlexibleF64, +} + +/// Fetch the wallet balance for one platform session token. +/// +/// `ProviderError::AuthRequired` means the platform rejected the session +/// (HTTP 401/403 or an auth envelope code). Transport failures keep their +/// reqwest classification so callers can tell an outage from a rejection. +pub(super) async fn fetch_platform_balance( + client: &Client, + token: &str, +) -> Result { + let token = token.trim(); + if token.is_empty() { + return Err(ProviderError::AuthRequired); + } + let response = client + .get(PLATFORM_USER_SUMMARY_URL) + .header("Authorization", format!("Bearer {token}")) + .header("Accept", "application/json") + .header("x-client-platform", "web") + .send() + .await?; + let status = response.status(); + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + return Err(ProviderError::AuthRequired); + } + if status != StatusCode::OK { + return Err(ProviderError::Other(format!( + "DeepSeek Platform returned status {status}" + ))); + } + let body = response.bytes().await?; + parse_platform_balance(&body) +} + +pub(super) fn parse_platform_balance(body: &[u8]) -> Result { + let parse_error = |error: serde_json::Error| { + ProviderError::Parse(format!( + "Failed to parse DeepSeek Platform balance: {error}" + )) + }; + let envelope: UserSummaryResponse = serde_json::from_slice(body).map_err(parse_error)?; + check_code(envelope.code, "user summary code")?; + + let data: UserSummaryData = envelope + .data + .map(serde_json::from_value) + .transpose() + .map_err(parse_error)? + .ok_or_else(|| ProviderError::Parse("Missing DeepSeek Platform user summary".into()))?; + check_code(data.biz_code, "user summary biz_code")?; + + let summary: UserSummary = data + .biz_data + .map(serde_json::from_value) + .transpose() + .map_err(parse_error)? + .ok_or_else(|| ProviderError::Parse("Missing DeepSeek Platform biz_data".into()))?; + Ok(balance_response(&summary)) +} + +fn check_code(code: Option, label: &str) -> Result<(), ProviderError> { + match code.map(|code| code.0) { + None | Some(0) => Ok(()), + Some(code) if AUTH_ERROR_CODES.contains(&code) => Err(ProviderError::AuthRequired), + Some(code) => Err(ProviderError::Other(format!( + "DeepSeek Platform {label} {code}" + ))), + } +} + +/// Sum wallets per currency and express them as the API-key balance shape, so +/// both lanes render through the same snapshot code. +fn balance_response(summary: &UserSummary) -> BalanceResponse { + let mut totals: BTreeMap<&str, (f64, f64)> = BTreeMap::new(); + for wallet in &summary.normal_wallets { + totals.entry(wallet.currency.as_str()).or_default().0 += wallet.balance.0; + } + for wallet in &summary.bonus_wallets { + totals.entry(wallet.currency.as_str()).or_default().1 += wallet.balance.0; + } + + let balance_infos: Vec = totals + .into_iter() + .map(|(currency, (topped_up, granted))| BalanceInfo { + currency: currency.to_string(), + total_balance: (topped_up + granted).to_string(), + granted_balance: granted.to_string(), + topped_up_balance: topped_up.to_string(), + }) + .collect(); + + let Some(selected) = select_balance_info(&balance_infos) else { + return BalanceResponse { + is_available: false, + balance_infos: vec![BalanceInfo { + currency: "USD".into(), + total_balance: "0".into(), + granted_balance: "0".into(), + topped_up_balance: "0".into(), + }], + }; + }; + let is_available = selected.total_balance.parse::().unwrap_or(0.0) > 0.0; + BalanceResponse { + is_available, + balance_infos, + } +} + +#[cfg(test)] +#[path = "platform_balance_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/platform_balance_tests.rs b/rust/src/providers/deepseek/platform_balance_tests.rs new file mode 100644 index 0000000000..eee4c68c4f --- /dev/null +++ b/rust/src/providers/deepseek/platform_balance_tests.rs @@ -0,0 +1,93 @@ +use super::*; + +fn parse(body: &str) -> Result { + parse_platform_balance(body.as_bytes()) +} + +fn summary_body(normal: &str, bonus: &str) -> String { + format!( + r#"{{"code":0,"msg":"","data":{{"biz_code":0,"biz_msg":"","biz_data":{{"normal_wallets":{normal},"bonus_wallets":{bonus}}}}}}}"# + ) +} + +#[test] +fn sums_wallets_per_currency_and_prefers_funded_usd() { + let body = summary_body( + r#"[{"currency":"USD","balance":"2.50"},{"currency":"CNY","balance":"9"}]"#, + r#"[{"currency":"USD","balance":1.25}]"#, + ); + let response = parse(&body).unwrap(); + assert!(response.is_available); + + let snapshot = crate::providers::deepseek::DeepSeekProvider::snapshot_from_balance(response); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("$3.75 (Paid: $2.50 / Granted: $1.25)") + ); +} + +#[test] +fn falls_back_to_funded_cny_when_usd_is_empty() { + let body = summary_body( + r#"[{"currency":"USD","balance":"0"},{"currency":"CNY","balance":"40"}]"#, + r#"[{"currency":"CNY","balance":"2.25"}]"#, + ); + let snapshot = + crate::providers::deepseek::DeepSeekProvider::snapshot_from_balance(parse(&body).unwrap()); + assert_eq!( + snapshot.login_method.as_deref(), + Some("CNY balance: ¥42.25") + ); +} + +#[test] +fn empty_wallets_report_an_unavailable_usd_balance() { + let response = parse(&summary_body("[]", "[]")).unwrap(); + assert!(!response.is_available); + assert_eq!(response.balance_infos.len(), 1); + assert_eq!(response.balance_infos[0].currency, "USD"); +} + +#[test] +fn zero_balance_is_not_available() { + let response = parse(&summary_body(r#"[{"currency":"USD","balance":"0"}]"#, "[]")).unwrap(); + assert!(!response.is_available); +} + +#[test] +fn auth_envelope_codes_are_session_rejections() { + for code in [40002, 40003] { + let body = format!(r#"{{"code":{code},"msg":"auth","data":{{"unexpected":true}}}}"#); + assert!(matches!(parse(&body), Err(ProviderError::AuthRequired))); + + let body = format!( + r#"{{"code":0,"data":{{"biz_code":{code},"biz_msg":"auth","biz_data":"not a summary"}}}}"# + ); + assert!(matches!(parse(&body), Err(ProviderError::AuthRequired))); + } +} + +#[test] +fn other_envelope_codes_are_not_session_rejections() { + let error = parse(r#"{"code":50000,"msg":"busy","data":null}"#).unwrap_err(); + assert!(matches!(error, ProviderError::Other(_)), "{error:?}"); + assert!(!error.is_transport_failure()); + + let error = parse(r#"{"code":0,"data":{"biz_code":1,"biz_data":null}}"#).unwrap_err(); + assert!(matches!(error, ProviderError::Other(_)), "{error:?}"); +} + +#[test] +fn malformed_or_missing_payloads_are_parse_errors() { + for body in [ + "not json", + r#"{"code":0}"#, + r#"{"code":0,"data":{"biz_code":0}}"#, + r#"{"code":0,"data":{"biz_code":0,"biz_data":{"normal_wallets":[{"currency":"USD","balance":"abc"}]}}}"#, + ] { + assert!( + matches!(parse(body), Err(ProviderError::Parse(_))), + "{body}" + ); + } +} diff --git a/rust/src/providers/deepseek/session_resolver.rs b/rust/src/providers/deepseek/session_resolver.rs new file mode 100644 index 0000000000..f38ffc2d34 --- /dev/null +++ b/rust/src/providers/deepseek/session_resolver.rs @@ -0,0 +1,311 @@ +//! Choose the Chrome session that supplies the DeepSeek Platform balance, and +//! decide which session a failed request belongs to. +//! +//! Every candidate is validated against the platform (a rejected session is +//! remembered for [`VALIDITY_TTL`]). The selected session's balance is always +//! fetched live. A transport failure keeps the identity of the session it came +//! from, so the shell can keep showing the last balance only when that same +//! session (same profile, same token) produced it. A session the platform has +//! rejected never keeps that authority. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use futures::future::join_all; + +use super::BalanceResponse; +use super::chrome_session::{ProfileSelection, TokenInfo}; +use crate::core::{LastGoodOwner, ProviderError}; + +/// How long a validation result is trusted before the platform is asked again. +pub(super) const VALIDITY_TTL: Duration = Duration::from_secs(30 * 60); + +const OWNER_NAMESPACE: &str = "deepseek-platform-balance"; + +/// Owner of a balance fetched with `token` from Chrome profile `profile_id`. +pub(super) fn balance_owner(profile_id: &str, token: &str) -> Option { + LastGoodOwner::derive(OWNER_NAMESPACE, profile_id, token) +} + +/// Outcome of asking the platform about one candidate. +enum Validation { + Valid(Box), + Rejected, + Unavailable(Option), +} + +struct Outcome { + index: usize, + validation: Validation, +} + +/// Result of resolving the selected Chrome session. +#[derive(Debug)] +pub(super) enum Resolution { + Balance { + balance: Box, + owner: Option, + }, + /// No profile holds a usable DeepSeek session. + SessionRequired, + /// Several profiles hold a session and none was selected. + SelectionRequired(Vec), + Failed(ProviderError), +} + +#[derive(Clone, Copy)] +struct Lookup { + fresh: Option, + last_known: Option, +} + +struct CacheEntry { + proof: Option, + status: bool, + checked_at: Instant, +} + +/// Remembers which sessions the platform accepted or rejected. Holds a digest +/// of each token, never the token. +pub(super) struct ValidationCache { + ttl: Duration, + entries: Mutex>, +} + +impl ValidationCache { + pub(super) fn new(ttl: Duration) -> Self { + Self { + ttl, + entries: Mutex::new(HashMap::new()), + } + } + + fn lookup(&self, candidate: &TokenInfo, now: Instant) -> Lookup { + let proof = balance_owner(&candidate.id, &candidate.token); + let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner()); + match entries.get(&candidate.id) { + Some(entry) if entry.proof == proof => Lookup { + fresh: (now.saturating_duration_since(entry.checked_at) < self.ttl) + .then_some(entry.status), + last_known: Some(entry.status), + }, + _ => Lookup { + fresh: None, + last_known: None, + }, + } + } + + fn record(&self, candidate: &TokenInfo, status: bool, now: Instant) { + let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner()); + entries.insert( + candidate.id.clone(), + CacheEntry { + proof: balance_owner(&candidate.id, &candidate.token), + status, + checked_at: now, + }, + ); + } +} + +/// Resolve the session to use and fetch its balance through `validate`. +pub(super) async fn resolve( + candidates: &[TokenInfo], + selection: &ProfileSelection, + cache: &ValidationCache, + validate: F, +) -> Resolution +where + F: Fn(String) -> Fut, + Fut: Future>, +{ + if candidates.is_empty() { + return Resolution::SessionRequired; + } + + let now = Instant::now(); + let lookups: Vec = candidates.iter().map(|c| cache.lookup(c, now)).collect(); + let selected_id = selection.profile_id.as_deref(); + let to_validate: Vec = (0..candidates.len()) + .filter(|&i| lookups[i].fresh.is_none() || Some(candidates[i].id.as_str()) == selected_id) + .collect(); + + let mut outcomes = run_validation(candidates, &to_validate, &validate).await; + record(cache, candidates, &outcomes, now); + let mut statuses = resolved_statuses(&lookups, &outcomes); + let mut valid = valid_indices(&statuses); + + // A status that came only from the cache carries no balance. + if let Some(index) = selected_candidate(candidates, &valid, selection) + && !has_balance(&outcomes, index) + { + outcomes.extend(run_validation(candidates, &[index], &validate).await); + record(cache, candidates, &outcomes, now); + statuses = resolved_statuses(&lookups, &outcomes); + valid = valid_indices(&statuses); + } + + let all: Vec = (0..candidates.len()).collect(); + let unresolved = selected_candidate(candidates, &all, selection); + + if valid.is_empty() { + let failure = + unresolved.and_then(|index| take_failure(candidates, &mut outcomes, &statuses, index)); + if let Some(error) = failure { + return Resolution::Failed(error); + } + if outcomes + .iter() + .any(|o| matches!(o.validation, Validation::Unavailable(_))) + { + return Resolution::Failed(validation_unavailable()); + } + return Resolution::SessionRequired; + } + + let Some(selected) = selected_candidate(candidates, &valid, selection) else { + let rejected = unresolved.is_some_and(|index| statuses[index] == Some(false)); + let failure = if rejected { + None + } else { + unresolved.and_then(|index| take_failure(candidates, &mut outcomes, &statuses, index)) + }; + return match failure { + Some(error) => Resolution::Failed(error), + None => Resolution::SelectionRequired( + valid.iter().map(|&i| candidates[i].id.clone()).collect(), + ), + }; + }; + + if let Some(balance) = take_balance(&mut outcomes, selected) { + let candidate = &candidates[selected]; + return Resolution::Balance { + balance, + owner: balance_owner(&candidate.id, &candidate.token), + }; + } + Resolution::Failed( + take_failure(candidates, &mut outcomes, &statuses, selected) + .unwrap_or_else(validation_unavailable), + ) +} + +fn validation_unavailable() -> ProviderError { + ProviderError::Other("Chrome DeepSeek session could not be validated right now".to_string()) +} + +async fn run_validation( + candidates: &[TokenInfo], + indices: &[usize], + validate: &F, +) -> Vec +where + F: Fn(String) -> Fut, + Fut: Future>, +{ + join_all(indices.iter().map(|&index| async move { + let validation = match validate(candidates[index].token.clone()).await { + Ok(balance) => Validation::Valid(Box::new(balance)), + Err(ProviderError::AuthRequired) => Validation::Rejected, + Err(error) => Validation::Unavailable(Some(error)), + }; + Outcome { index, validation } + })) + .await +} + +fn record(cache: &ValidationCache, candidates: &[TokenInfo], outcomes: &[Outcome], now: Instant) { + for outcome in outcomes { + match outcome.validation { + Validation::Valid(_) => cache.record(&candidates[outcome.index], true, now), + Validation::Rejected => cache.record(&candidates[outcome.index], false, now), + Validation::Unavailable(_) => {} + } + } +} + +/// Known status per candidate: a fresh cache entry, overridden by this round's +/// outcomes. An unavailable check falls back to the last known status. +fn resolved_statuses(lookups: &[Lookup], outcomes: &[Outcome]) -> Vec> { + let mut statuses: Vec> = lookups.iter().map(|l| l.fresh).collect(); + for outcome in outcomes { + match outcome.validation { + Validation::Valid(_) => statuses[outcome.index] = Some(true), + Validation::Rejected => statuses[outcome.index] = Some(false), + Validation::Unavailable(_) => { + if let Some(known) = lookups[outcome.index].last_known { + statuses[outcome.index] = Some(known); + } + } + } + } + statuses +} + +fn valid_indices(statuses: &[Option]) -> Vec { + statuses + .iter() + .enumerate() + .filter_map(|(index, status)| (*status == Some(true)).then_some(index)) + .collect() +} + +/// The candidate the user's selection points at among `pool`. Without an +/// explicit selection only a single candidate is unambiguous. +fn selected_candidate( + candidates: &[TokenInfo], + pool: &[usize], + selection: &ProfileSelection, +) -> Option { + match selection.profile_id.as_deref() { + Some(id) => pool.iter().copied().find(|&i| candidates[i].id == id), + None => (pool.len() == 1).then(|| pool[0]), + } +} + +fn has_balance(outcomes: &[Outcome], index: usize) -> bool { + outcomes + .iter() + .any(|o| o.index == index && matches!(o.validation, Validation::Valid(_))) +} + +fn take_balance(outcomes: &mut [Outcome], index: usize) -> Option> { + let outcome = outcomes + .iter_mut() + .rev() + .find(|o| o.index == index && matches!(o.validation, Validation::Valid(_)))?; + match std::mem::replace(&mut outcome.validation, Validation::Unavailable(None)) { + Validation::Valid(balance) => Some(balance), + _ => None, + } +} + +/// The latest failed check of `index` as an error. A transport failure carries +/// the session that produced it, except when the platform already rejected +/// that session: an outage after a rejection must not revive the old balance. +fn take_failure( + candidates: &[TokenInfo], + outcomes: &mut [Outcome], + statuses: &[Option], + index: usize, +) -> Option { + let latest = outcomes.iter_mut().rev().find(|o| o.index == index)?; + let Validation::Unavailable(error) = &mut latest.validation else { + return None; + }; + let error = error.take()?; + let owner = if statuses[index] == Some(false) { + None + } else { + balance_owner(&candidates[index].id, &candidates[index].token) + }; + Some(error.with_failure_owner(owner)) +} + +#[cfg(test)] +#[path = "session_resolver_tests.rs"] +mod tests; diff --git a/rust/src/providers/deepseek/session_resolver_tests.rs b/rust/src/providers/deepseek/session_resolver_tests.rs new file mode 100644 index 0000000000..d0da44ebfe --- /dev/null +++ b/rust/src/providers/deepseek/session_resolver_tests.rs @@ -0,0 +1,366 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use super::*; +use crate::core::FailureOwnership; +use crate::providers::deepseek::BalanceInfo; + +fn candidate(id: &str, token: &str) -> TokenInfo { + TokenInfo { + id: id.to_string(), + token: token.to_string(), + label: format!("Google Chrome {id}"), + } +} + +fn select(id: &str) -> ProfileSelection { + ProfileSelection::from_value(Some(id)) +} + +fn balance(total: &str) -> BalanceResponse { + BalanceResponse { + is_available: true, + balance_infos: vec![BalanceInfo { + currency: "USD".into(), + total_balance: total.into(), + granted_balance: "0".into(), + topped_up_balance: total.into(), + }], + } +} + +fn fresh_cache() -> ValidationCache { + ValidationCache::new(VALIDITY_TTL) +} + +fn owner_of(id: &str, token: &str) -> LastGoodOwner { + balance_owner(id, token).expect("owner") +} + +/// Token-keyed outcomes: `ok-*` succeed, `rejected-*` are refused by the +/// platform, `down-*` hit a transport failure, `broken-*` a non-transport one. +async fn answer(token: String) -> Result { + if token.starts_with("rejected") { + Err(ProviderError::AuthRequired) + } else if token.starts_with("down") { + Err(ProviderError::Timeout) + } else if token.starts_with("broken") { + Err(ProviderError::Parse("bad body".into())) + } else { + Ok(balance("5")) + } +} + +fn failed(resolution: Resolution) -> ProviderError { + match resolution { + Resolution::Failed(error) => error, + other => panic!("expected failure, got {other:?}"), + } +} + +#[tokio::test] +async fn no_candidates_requires_a_session() { + let resolution = resolve(&[], &ProfileSelection::default(), &fresh_cache(), answer).await; + assert!(matches!(resolution, Resolution::SessionRequired)); +} + +#[tokio::test] +async fn single_valid_session_supplies_balance_and_owner() { + let candidates = [candidate("chrome:Default", "ok-token")]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + let Resolution::Balance { balance, owner } = resolution else { + panic!("expected balance"); + }; + assert!(balance.is_available); + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-token"))); +} + +#[tokio::test] +async fn owner_follows_profile_and_token() { + assert_ne!( + owner_of("chrome:Default", "ok-one"), + owner_of("chrome:Default", "ok-two") + ); + assert_ne!( + owner_of("chrome:Default", "ok-one"), + owner_of("chrome:Profile 1", "ok-one") + ); + assert_eq!( + owner_of("chrome:Default", "ok-one"), + owner_of(" chrome:Default", "ok-one ") + ); +} + +#[tokio::test] +async fn rejected_session_requires_a_new_sign_in() { + let candidates = [candidate("chrome:Default", "rejected-token")]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + assert!(matches!(resolution, Resolution::SessionRequired)); +} + +#[tokio::test] +async fn transport_failure_keeps_the_session_that_failed() { + let candidates = [candidate("chrome:Default", "down-token")]; + let error = failed( + resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await, + ); + assert!(error.is_transport_failure()); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Default", "down-token")) + ); +} + +#[tokio::test] +async fn non_transport_failure_is_not_owned() { + let candidates = [candidate("chrome:Default", "broken-token")]; + let error = failed( + resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await, + ); + assert!(matches!(error, ProviderError::Parse(_))); + assert_eq!(error.failure_ownership(), FailureOwnership::Unchecked); +} + +#[tokio::test] +async fn outage_after_rejection_cannot_revive_the_old_balance() { + let cache = fresh_cache(); + let candidates = [candidate("chrome:Default", "token-flaky")]; + let calls = Arc::new(AtomicUsize::new(0)); + let flaky = |token: String| { + let calls = Arc::clone(&calls); + async move { + let _ = token; + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => Err(ProviderError::AuthRequired), + _ => Err(ProviderError::Timeout), + } + } + }; + let selection = select("chrome:Default"); + + assert!(matches!( + resolve(&candidates, &selection, &cache, &flaky).await, + Resolution::SessionRequired + )); + let error = failed(resolve(&candidates, &selection, &cache, &flaky).await); + assert!(error.is_transport_failure()); + assert_eq!(error.failure_ownership(), FailureOwnership::Unowned); +} + +#[tokio::test] +async fn several_sessions_without_a_selection_ask_for_one() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + let Resolution::SelectionRequired(ids) = resolution else { + panic!("expected selection"); + }; + assert_eq!(ids, vec!["chrome:Default", "chrome:Profile 1"]); +} + +#[tokio::test] +async fn explicit_selection_supplies_that_profiles_balance() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await; + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Profile 1", "ok-b"))); +} + +#[tokio::test] +async fn selected_profile_keeps_its_own_outage_beside_a_working_profile() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "down-b"), + ]; + let error = failed( + resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await, + ); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Profile 1", "down-b")) + ); +} + +#[tokio::test] +async fn ambiguous_outage_does_not_borrow_another_profiles_failure() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "down-b"), + ]; + let resolution = resolve( + &candidates, + &ProfileSelection::default(), + &fresh_cache(), + answer, + ) + .await; + // Only Default is known valid; with no explicit selection the single valid + // profile is selected and supplies its own balance. + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-a"))); +} + +#[tokio::test] +async fn rejected_selection_beside_a_valid_profile_asks_for_a_choice() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "rejected-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Profile 1"), + &fresh_cache(), + answer, + ) + .await; + let Resolution::SelectionRequired(ids) = resolution else { + panic!("expected selection, got {resolution:?}"); + }; + assert_eq!(ids, vec!["chrome:Default"]); +} + +#[tokio::test] +async fn unknown_selection_asks_for_a_choice() { + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let resolution = resolve( + &candidates, + &select("chrome:Missing"), + &fresh_cache(), + answer, + ) + .await; + assert!(matches!(resolution, Resolution::SelectionRequired(_))); +} + +#[tokio::test] +async fn fresh_cache_skips_revalidating_unselected_profiles() { + let cache = fresh_cache(); + let calls = Arc::new(AtomicUsize::new(0)); + let counting = |token: String| { + let calls = Arc::clone(&calls); + async move { + calls.fetch_add(1, Ordering::SeqCst); + answer(token).await + } + }; + let candidates = [ + candidate("chrome:Default", "ok-a"), + candidate("chrome:Profile 1", "ok-b"), + ]; + let selection = select("chrome:Default"); + + resolve(&candidates, &selection, &cache, &counting).await; + assert_eq!(calls.load(Ordering::SeqCst), 2); + resolve(&candidates, &selection, &cache, &counting).await; + // Only the selected profile is asked again: its balance is always live. + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn changed_token_in_the_same_profile_is_validated_again() { + let cache = fresh_cache(); + let selection = select("chrome:Default"); + let first = [candidate("chrome:Default", "rejected-old")]; + assert!(matches!( + resolve(&first, &selection, &cache, answer).await, + Resolution::SessionRequired + )); + + let second = [candidate("chrome:Default", "ok-new")]; + let resolution = resolve(&second, &selection, &cache, answer).await; + let Resolution::Balance { owner, .. } = resolution else { + panic!("expected balance"); + }; + assert_eq!(owner, Some(owner_of("chrome:Default", "ok-new"))); +} + +#[tokio::test] +async fn cached_status_survives_a_later_validation_outage() { + let cache = ValidationCache::new(Duration::ZERO); + let candidates = [candidate("chrome:Default", "token-x")]; + let calls = Arc::new(AtomicUsize::new(0)); + let script = |_token: String| { + let calls = Arc::clone(&calls); + async move { + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => Ok(balance("5")), + _ => Err(ProviderError::Timeout), + } + } + }; + let selection = ProfileSelection::default(); + + assert!(matches!( + resolve(&candidates, &selection, &cache, &script).await, + Resolution::Balance { .. } + )); + let error = failed(resolve(&candidates, &selection, &cache, &script).await); + assert_eq!( + error.failure_ownership(), + FailureOwnership::Owned(owner_of("chrome:Default", "token-x")) + ); +} + +#[test] +fn validation_cache_stores_no_plain_token() { + let cache = fresh_cache(); + let candidate = candidate("chrome:Default", "super-secret-token-value"); + cache.record(&candidate, true, Instant::now()); + let entries = cache.entries.lock().unwrap(); + let rendered = format!("{:?}", entries.get("chrome:Default").unwrap().proof); + assert!(!rendered.contains("super-secret")); +} diff --git a/rust/src/providers/deepseek/tests.rs b/rust/src/providers/deepseek/tests.rs index badaaf05e1..4acac37f8d 100644 --- a/rust/src/providers/deepseek/tests.rs +++ b/rust/src/providers/deepseek/tests.rs @@ -239,3 +239,34 @@ fn applies_deepseek_summary_as_extra_windows() { Some("¥0.0000 · Current month") ); } + +#[test] +fn only_owned_transport_failures_retain_the_last_balance() { + let provider = DeepSeekProvider::new(); + let owner = session_resolver::balance_owner("chrome:Default", "token-value"); + + assert_eq!( + provider + .last_good_failure_policy_for_error(&ProviderError::Timeout.with_failure_owner(owner)), + LastGoodFailurePolicy::Preserve + ); + // A deadline has no owner but is still an owned failure; the shell then + // fails closed because no session can match. + assert_eq!( + provider + .last_good_failure_policy_for_error(&ProviderError::Timeout.with_failure_owner(None)), + LastGoodFailurePolicy::Preserve + ); + for error in [ + ProviderError::Timeout, + ProviderError::AuthRequired, + ProviderError::Other("HTTP 500".into()), + ProviderError::Parse("bad".into()), + ] { + assert_eq!( + provider.last_good_failure_policy_for_error(&error), + LastGoodFailurePolicy::Replace, + "{error:?}" + ); + } +} From b883c478713dbb74de21be163174b910c0c36997 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:41:47 +0700 Subject: [PATCH 122/259] Port upstream 0.63.0: support StepFun credit plans with Credit label and no invented reset --- rust/src/providers/stepfun/mod.rs | 266 +++++++++++++++++++++++++++++- 1 file changed, 258 insertions(+), 8 deletions(-) diff --git a/rust/src/providers/stepfun/mod.rs b/rust/src/providers/stepfun/mod.rs index 4f9a875cb8..b35f184016 100644 --- a/rust/src/providers/stepfun/mod.rs +++ b/rust/src/providers/stepfun/mod.rs @@ -24,8 +24,10 @@ const STEPFUN_REFRESH_TOKEN_URL: &str = const STEPFUN_CREDENTIAL_TARGET: &str = "codexbar-stepfun"; const STEPFUN_WEB_ID: &str = "734152690100432"; const STEPFUN_APP_ID: &str = "111003695"; +const CREDIT_LABEL: &str = "Credit"; +const NO_CREDIT_BALANCE: &str = "No credit balance reported"; -#[derive(Debug, Deserialize)] +#[derive(Debug, Default, Deserialize)] struct StepFunRateLimitResponse { status: Option, code: Option, @@ -35,6 +37,111 @@ struct StepFunRateLimitResponse { weekly_usage_left_rate: Option, five_hour_usage_reset_time: Option, weekly_usage_reset_time: Option, + plan_family: Option, + plan_credit_rate_limit: Option, +} + +impl StepFunRateLimitResponse { + /// StepFun serves two Step Plan billing models. The Coding Plan meters + /// rolling 5-hour / weekly windows; the Token Plan meters a monthly Credit + /// pool through `plan_credit_rate_limit` and reports its rolling windows as + /// 0 with a `"0"` reset time ("no window configured", not "used up"). + /// + /// Classify by the payload shape: a live rolling window means Coding Plan, + /// no window plus a credit pool means Token Plan. `plan_family == 2` only + /// breaks the tie for an ambiguous payload, so a future family-id change + /// cannot flip a windowed plan onto the credit renderer or vice versa. + fn is_credit_plan(&self) -> bool { + let has_live_window = [ + &self.five_hour_usage_reset_time, + &self.weekly_usage_reset_time, + ] + .into_iter() + .any(|reset| reset.as_ref().is_some_and(|ts| ts.0 > 0)); + if has_live_window { + return false; + } + if self + .plan_credit_rate_limit + .as_ref() + .is_some_and(StepFunPlanCreditRateLimit::has_credit_pool) + { + return true; + } + self.plan_family + .as_ref() + .is_some_and(|family| family.0 == 2.0) + } +} + +/// The `plan_credit_rate_limit` object returned for credit-based plans. +#[derive(Debug, Default, Deserialize)] +struct StepFunPlanCreditRateLimit { + subscription_credit_left_rate: Option, + subscription_credit_reset_time: Option, + topup_credit_left_rate: Option, + credit_buckets: Option>, +} + +#[derive(Debug, Default, Deserialize)] +struct StepFunCreditBucket { + credit_total: Option, + credit_residual: Option, +} + +impl StepFunPlanCreditRateLimit { + fn has_credit_pool(&self) -> bool { + self.subscription_credit_left_rate.is_some() + || self.topup_credit_left_rate.is_some() + || self + .credit_buckets + .as_ref() + .is_some_and(|buckets| !buckets.is_empty()) + } + + /// Remaining fraction of the credit pool, or `None` when no balance is reported. + /// + /// Subscription and top-up rates are independent fractions, so adding them + /// does not give a combined rate; absolute bucket balances are preferred. + /// Without usable bucket sizes the subscription rate is the plan allowance + /// and the top-up rate is used only when no subscription rate is present. + fn left_rate(&self) -> Option { + let buckets = self.credit_buckets.as_deref().unwrap_or_default(); + let balances: Vec<(f64, f64)> = buckets + .iter() + .filter_map(StepFunCreditBucket::balance) + .collect(); + if !buckets.is_empty() && balances.len() == buckets.len() { + let total: f64 = balances.iter().map(|(total, _)| total).sum(); + let residual: f64 = balances.iter().map(|(_, residual)| residual).sum(); + return Some(residual / total); + } + self.subscription_credit_left_rate + .as_ref() + .or(self.topup_credit_left_rate.as_ref()) + .map(|rate| rate.0) + } + + /// A real monthly reset; a missing or zero timestamp stays unknown. + fn reset_at(&self) -> Option> { + self.subscription_credit_reset_time + .as_ref() + .filter(|ts| ts.0 > 0) + .and_then(|ts| Utc.timestamp_opt(ts.0, 0).single()) + } +} + +impl StepFunCreditBucket { + /// `(total, residual)` when the bucket carries a sound balance. + fn balance(&self) -> Option<(f64, f64)> { + let total = self.credit_total.as_ref()?.0; + let residual = self.credit_residual.as_ref()?.0; + (total.is_finite() + && residual.is_finite() + && total > 0.0 + && (0.0..=total).contains(&residual)) + .then_some((total, residual)) + } } #[derive(Debug, Deserialize)] @@ -225,6 +332,14 @@ fn snapshot_from_response( return Err(ProviderError::Other(format!("StepFun API error: {msg}"))); } + let login_method = plan_name + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| "Oasis-Token".into()); + + if response.is_credit_plan() { + return Ok(credit_snapshot(response).with_login_method(login_method)); + } + let five_left = response .five_hour_usage_left_rate .as_ref() @@ -257,13 +372,28 @@ fn snapshot_from_response( weekly_reset.map(reset_description), ); - let mut snapshot = UsageSnapshot::new(primary).with_secondary(secondary); - if let Some(plan_name) = plan_name.filter(|value| !value.trim().is_empty()) { - snapshot = snapshot.with_login_method(plan_name); - } else { - snapshot = snapshot.with_login_method("Oasis-Token"); - } - Ok(snapshot) + Ok(UsageSnapshot::new(primary) + .with_secondary(secondary) + .with_login_method(login_method)) +} + +/// Credit plans populate only the primary lane, including balances without a +/// reset timestamp, and never invent a reset date. +fn credit_snapshot(response: &StepFunRateLimitResponse) -> UsageSnapshot { + let credit = response.plan_credit_rate_limit.as_ref(); + let primary = match credit.and_then(StepFunPlanCreditRateLimit::left_rate) { + Some(left_rate) => { + let reset = credit.and_then(StepFunPlanCreditRateLimit::reset_at); + RateWindow::with_details( + (1.0 - left_rate).clamp(0.0, 1.0) * 100.0, + RateWindow::monthly_window_minutes(reset), + reset, + reset.map(reset_description), + ) + } + None => RateWindow::informational(NO_CREDIT_BALANCE), + }; + UsageSnapshot::new(primary).with_primary_label(CREDIT_LABEL) } struct StepFunTokenParts { @@ -455,12 +585,132 @@ mod tests { weekly_usage_left_rate: Some(FlexibleNumber(0.75)), five_hour_usage_reset_time: Some(FlexibleTimestamp(1_800_000_000)), weekly_usage_reset_time: Some(FlexibleTimestamp(1_800_000_000)), + ..Default::default() }; let snapshot = snapshot_from_response(&response, Some("Step Plan".into())).unwrap(); assert_eq!(snapshot.primary.used_percent, 75.0); + assert_eq!(snapshot.primary_label, None); assert_eq!(snapshot.secondary.unwrap().used_percent, 25.0); } + fn snapshot_from_json(json: &str) -> Result { + let response: StepFunRateLimitResponse = serde_json::from_str(json).unwrap(); + snapshot_from_response(&response, None) + } + + #[test] + fn stepfun_coding_plan_payload_keeps_window_labels_and_lanes() { + let snapshot = snapshot_from_json( + r#"{"status":1,"five_hour_usage_left_rate":0.99781543,"weekly_usage_left_rate":1, + "five_hour_usage_reset_time":"1777528800","weekly_usage_reset_time":"1777852800", + "plan_family":1}"#, + ) + .unwrap(); + assert_eq!(snapshot.primary_label, None); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!(snapshot.secondary.unwrap().window_minutes, Some(10080)); + } + + #[test] + fn stepfun_coding_plan_payload_still_requires_window_fields() { + let error = snapshot_from_json( + r#"{"status":1,"five_hour_usage_reset_time":"1777528800","plan_family":1}"#, + ) + .unwrap_err(); + assert!(error.to_string().contains("five-hour")); + } + + #[test] + fn stepfun_credit_plan_with_reset_uses_monthly_window() { + let snapshot = snapshot_from_json( + r#"{"status":1,"five_hour_usage_left_rate":0,"weekly_usage_left_rate":0, + "five_hour_usage_reset_time":"0","weekly_usage_reset_time":"0","plan_family":2, + "plan_credit_rate_limit":{"subscription_credit_left_rate":0.75, + "subscription_credit_reset_time":"1777528800","topup_credit_left_rate":0}}"#, + ) + .unwrap(); + assert_eq!(snapshot.primary_label.as_deref(), Some("Credit")); + assert!(snapshot.secondary.is_none()); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!( + snapshot.primary.resets_at.map(|t| t.timestamp()), + Some(1_777_528_800) + ); + assert!(snapshot.primary.reset_description.is_some()); + let minutes = snapshot.primary.window_minutes.unwrap(); + assert!((28 * 1440..=31 * 1440).contains(&minutes)); + } + + #[test] + fn stepfun_credit_plan_without_or_with_zero_reset_invents_no_reset() { + for reset in ["", r#","subscription_credit_reset_time":"0""#] { + let json = format!( + r#"{{"status":1,"plan_credit_rate_limit":{{"subscription_credit_left_rate":0.4{reset}}}}}"# + ); + let snapshot = snapshot_from_json(&json).unwrap(); + assert_eq!(snapshot.primary_label.as_deref(), Some("Credit")); + assert_eq!(snapshot.primary.used_percent, 60.0); + assert_eq!(snapshot.primary.resets_at, None); + assert_eq!(snapshot.primary.reset_description, None); + assert_eq!(snapshot.primary.window_minutes, None); + assert!(snapshot.secondary.is_none()); + } + } + + #[test] + fn stepfun_credit_plan_weights_buckets_by_balance() { + let snapshot = snapshot_from_json( + r#"{"status":1,"plan_credit_rate_limit":{"subscription_credit_left_rate":0.9, + "topup_credit_left_rate":0.9,"credit_buckets":[ + {"credit_total":"400000000","credit_residual":"100000000"}, + {"credit_total":100000000,"credit_residual":100000000}]}}"#, + ) + .unwrap(); + // residual 200M of total 500M => 60% used, not derived from the rates. + assert!((snapshot.primary.used_percent - 60.0).abs() < 1e-9); + } + + #[test] + fn stepfun_credit_plan_falls_back_to_rates_when_buckets_are_unsound() { + let snapshot = snapshot_from_json( + r#"{"status":1,"plan_credit_rate_limit":{"topup_credit_left_rate":0.5, + "credit_buckets":[{"credit_total":10,"credit_residual":20}]}}"#, + ) + .unwrap(); + assert_eq!(snapshot.primary.used_percent, 50.0); + let snapshot = snapshot_from_json( + r#"{"status":1,"plan_credit_rate_limit":{"subscription_credit_left_rate":0.8, + "topup_credit_left_rate":0.1}}"#, + ) + .unwrap(); + assert!((snapshot.primary.used_percent - 20.0).abs() < 1e-9); + } + + #[test] + fn stepfun_credit_family_without_balance_has_no_quota_window() { + let snapshot = snapshot_from_json( + r#"{"status":1,"plan_family":2,"five_hour_usage_reset_time":"0", + "weekly_usage_reset_time":"0"}"#, + ) + .unwrap(); + assert!(snapshot.primary.is_informational); + assert_eq!(snapshot.primary_label.as_deref(), Some("Credit")); + assert!(snapshot.secondary.is_none()); + } + + #[test] + fn stepfun_live_window_wins_over_credit_family() { + let snapshot = snapshot_from_json( + r#"{"status":1,"plan_family":2,"five_hour_usage_left_rate":0.5, + "weekly_usage_left_rate":0.5,"five_hour_usage_reset_time":"1777528800", + "weekly_usage_reset_time":"1777852800", + "plan_credit_rate_limit":{"subscription_credit_left_rate":0.1}}"#, + ) + .unwrap(); + assert_eq!(snapshot.primary_label, None); + assert!(snapshot.secondary.is_some()); + } + #[test] fn stepfun_token_parts_extract_cookie_and_refresh_token() { let parts = token_parts("Cookie: Oasis-Token=access...refresh; Oasis-Webid=abc"); From 093031ce9884219e99aa43c962f48bf2a0b06b5a Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:19:15 +0700 Subject: [PATCH 123/259] Address thermo review --- rust/src/providers/antigravity/mod.rs | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 0d79dd4000..498a846b15 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -784,13 +784,9 @@ impl AntigravityProvider { } } - match failure { - Some(error) => Self::resolve_probe_failure(error, offline), - None => Self::resolve_probe_failure( - ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.into()), - offline, - ), - } + let error = + failure.unwrap_or_else(|| ProviderError::NotInstalled(AGY_NOT_FOUND_MESSAGE.into())); + Self::resolve_probe_failure(error, offline) } fn locate_agy_binary() -> Option { From bd005516fe3dc203a5cfc90f62bdb20dd19378b3 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:43:46 +0700 Subject: [PATCH 124/259] Address thermo review Zed fails closed on empty manual cookie in browser mode; localize Zed usage-source picker. --- apps/desktop-tauri/src/i18n/keys.ts | 4 +++ .../providers/sections/UsageSourceSection.tsx | 8 ++++-- .../providers/sections/usageSourcePolicy.ts | 28 +++++++++++++++++-- rust/src/locale.rs | 4 +++ rust/src/locale/en-US.ftl | 4 +++ rust/src/locale/es-MX.ftl | 4 +++ rust/src/locale/ja-JP.ftl | 4 +++ rust/src/locale/ko-KR.ftl | 4 +++ rust/src/locale/ru-RU.ftl | 4 +++ rust/src/locale/tr-TR.ftl | 4 +++ rust/src/locale/zh-CN.ftl | 4 +++ rust/src/locale/zh-TW.ftl | 4 +++ rust/src/providers/zed/mod.rs | 8 ++++-- 13 files changed, 77 insertions(+), 7 deletions(-) diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 1c07040562..1d3055053b 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -285,6 +285,10 @@ export const ALL_LOCALE_KEYS = [ "MenuBarMetric", "MenuBarMetricHelper", "UsageSource", + "ProviderZedBrowserSession", + "ProviderZedUsageSourceAutoHelp", + "ProviderZedUsageSourceApiHelp", + "ProviderZedUsageSourceWebHelp", "ProviderNoCodexAccountsDetected", "ProviderCodexAutoImportHelp", "ProviderCodexHistoryHelp", diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.tsx index d8acbb4a28..eccd7bf0a2 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.tsx @@ -57,12 +57,16 @@ export function UsageSourceSection({ className={`provider-detail-segmented__option${isActive ? " is-active" : ""}`} onClick={() => void handleSelect(option.value)} > - {option.label} + {option.labelKey ? t(option.labelKey) : option.label} ); })}
-

{selectedOption.description}

+

+ {selectedOption.descriptionKey + ? t(selectedOption.descriptionKey) + : selectedOption.description} +

{error &&

{error}

} ); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts b/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts index 25e76fddf3..f5cc72dc62 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts @@ -1,7 +1,11 @@ +import type { LocaleKey } from "../../../../i18n/keys"; + export interface UsageSourceOption { value: string; label: string; + labelKey?: LocaleKey; description: string; + descriptionKey?: LocaleKey; } export interface UsageSourcePolicy { @@ -55,9 +59,27 @@ const POLICIES: Readonly> = { }, zed: { options: [ - { value: "auto", label: "Auto", description: "Uses the Zed editor credential; the browser session is used only when Browser session is selected." }, - { value: "oauth", label: "API", description: "Uses the Zed editor credential only." }, - { value: "web", label: "Browser session", description: "Reads token spend from the zed.dev browser session or manual cookie header only, with no editor-credential fallback." }, + { + value: "auto", + label: "Auto", + labelKey: "ProviderSourceAutoShort", + description: "Uses the Zed editor credential; the browser session is used only when Browser session is selected.", + descriptionKey: "ProviderZedUsageSourceAutoHelp", + }, + { + value: "oauth", + label: "API", + labelKey: "ProviderSourceApiShort", + description: "Uses the Zed editor credential only.", + descriptionKey: "ProviderZedUsageSourceApiHelp", + }, + { + value: "web", + label: "Browser session", + labelKey: "ProviderZedBrowserSession", + description: "Reads token spend from the zed.dev browser session or manual cookie header only, with no editor-credential fallback.", + descriptionKey: "ProviderZedUsageSourceWebHelp", + }, ], }, gitkraken: { diff --git a/rust/src/locale.rs b/rust/src/locale.rs index 276a31160e..c509bf91d0 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -527,6 +527,10 @@ locale_keys! { MenuBarMetric, MenuBarMetricHelper, UsageSource, + ProviderZedBrowserSession, + ProviderZedUsageSourceAutoHelp, + ProviderZedUsageSourceApiHelp, + ProviderZedUsageSourceWebHelp, ProviderNoCodexAccountsDetected, ProviderCodexAutoImportHelp, ProviderCodexHistoryHelp, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 53b4453f6d..1a47f00bc8 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -272,6 +272,10 @@ ProviderOptionsTitle = Options MenuBarMetric = Menu and floating bar metric MenuBarMetricHelper = Choose which window drives the menu bar and floating bar percentages. UsageSource = Usage source +ProviderZedBrowserSession = Browser session +ProviderZedUsageSourceAutoHelp = Uses the Zed editor credential; browser billing is used only when Browser session is selected. +ProviderZedUsageSourceApiHelp = Uses the Zed editor credential only. +ProviderZedUsageSourceWebHelp = Reads token spend only from the selected zed.dev browser session or saved Cookie header; it never falls back to the editor credential. ProviderNoCodexAccountsDetected = No Codex accounts detected yet. ProviderCodexAutoImportHelp = Automatic imports browser cookies for dashboard extras. ProviderCodexHistoryHelp = Stores local Codex usage history (8 weeks) to personalize Pace predictions. diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index bae2c814c3..0d200615d5 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -263,6 +263,10 @@ ProviderOptionsTitle = Opciones MenuBarMetric = Métrica de las barras de menú y flotante MenuBarMetricHelper = Elige qué ventana controla los porcentajes de las barras de menú y flotante. UsageSource = Fuente de uso +ProviderZedBrowserSession = Sesión del navegador +ProviderZedUsageSourceAutoHelp = Usa la credencial del editor de Zed; la facturación del navegador solo se usa al seleccionar «Sesión del navegador». +ProviderZedUsageSourceApiHelp = Usa solo la credencial del editor de Zed. +ProviderZedUsageSourceWebHelp = Lee el gasto de tokens solo de la sesión seleccionada de zed.dev o del encabezado Cookie guardado; nunca recurre a la credencial del editor. ProviderNoCodexAccountsDetected = Aún no se detectaron cuentas de Codex. ProviderCodexAutoImportHelp = Importa automáticamente cookies del navegador para extras del panel. ProviderCodexHistoryHelp = Almacena el historial local de uso de Codex (8 semanas) para personalizar predicciones de ritmo. diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 92618f72f7..45ed6e9aff 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -263,6 +263,10 @@ ProviderOptionsTitle = オプション MenuBarMetric = メニューバーとフローティングバーのメトリクス MenuBarMetricHelper = メニューバーとフローティングバーのパーセント表示に使用する時間枠を選択します。 UsageSource = 使用量ソース +ProviderZedBrowserSession = ブラウザーセッション +ProviderZedUsageSourceAutoHelp = Zed エディターの認証情報を使用します。ブラウザー課金を使うのは「ブラウザーセッション」を選択した場合のみです。 +ProviderZedUsageSourceApiHelp = Zed エディターの認証情報のみを使用します。 +ProviderZedUsageSourceWebHelp = 選択した zed.dev のブラウザーセッションまたは保存済み Cookie ヘッダーからのみトークン使用額を読み取り、エディターの認証情報にはフォールバックしません。 ProviderNoCodexAccountsDetected = Codex アカウントはまだ検出されていません。 ProviderCodexAutoImportHelp = ダッシュボードの追加情報のためブラウザ Cookie を自動インポートします。 ProviderCodexHistoryHelp = Pace 予測をパーソナライズするため、ローカルの Codex 使用履歴(8週間)を保存します。 diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 7f59c6fc29..08028728b4 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -263,6 +263,10 @@ ProviderOptionsTitle = 옵션 MenuBarMetric = 메뉴 바 및 플로팅 바 메트릭 MenuBarMetricHelper = 메뉴 바와 플로팅 바의 백분율을 결정할 시간 창을 선택합니다. UsageSource = 사용량 소스 +ProviderZedBrowserSession = 브라우저 세션 +ProviderZedUsageSourceAutoHelp = Zed 편집기 자격 증명을 사용하며, 브라우저 결제 정보는 브라우저 세션을 선택한 경우에만 사용합니다. +ProviderZedUsageSourceApiHelp = Zed 편집기 자격 증명만 사용합니다. +ProviderZedUsageSourceWebHelp = 선택한 zed.dev 브라우저 세션 또는 저장된 Cookie 헤더에서만 토큰 사용액을 읽으며 편집기 자격 증명으로 대체하지 않습니다. ProviderNoCodexAccountsDetected = 감지된 Codex 계정이 아직 없습니다. ProviderCodexAutoImportHelp = 대시보드 추가 기능을 위해 브라우저 쿠키를 자동으로 가져옵니다. ProviderCodexHistoryHelp = 사용 속도 예측을 개인화하기 위해 로컬 Codex 사용 기록(8주)을 저장합니다. diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index 855f7a9eb6..4ec3e58d7a 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -247,6 +247,10 @@ ProviderOptionsTitle = Опции MenuBarMetric = Метрика строки меню и плавающей панели MenuBarMetricHelper = Выберите временное окно для процентов в строке меню и на плавающей панели. UsageSource = Источник использования +ProviderZedBrowserSession = Сеанс браузера +ProviderZedUsageSourceAutoHelp = Используются учетные данные редактора Zed; данные браузерного биллинга применяются только при выборе сеанса браузера. +ProviderZedUsageSourceApiHelp = Используются только учетные данные редактора Zed. +ProviderZedUsageSourceWebHelp = Расход токенов считывается только из выбранного сеанса zed.dev или сохраненного заголовка Cookie; перехода на учетные данные редактора нет. ProviderNoCodexAccountsDetected = Учетные записи Codex пока не обнаружены. ProviderCodexAutoImportHelp = Автоматический импорт файлов cookie браузера для дополнительных функций панели управления. ProviderCodexHistoryHelp = Сохраняет локальную историю использования Codex (8 недель) для персонализации прогнозов темпа. diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index ca4fabe791..af6247168e 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -269,6 +269,10 @@ ProviderOptionsTitle = Seçenekler MenuBarMetric = Menü çubuğu ölçümü MenuBarMetricHelper = Menü çubuğundaki yüzdeyi hangi pencerenin belirleyeceğini seçin. UsageSource = Kullanım kaynağı +ProviderZedBrowserSession = Tarayıcı oturumu +ProviderZedUsageSourceAutoHelp = Zed düzenleyici kimlik bilgisi kullanılır; tarayıcı faturalandırması yalnızca Tarayıcı oturumu seçildiğinde kullanılır. +ProviderZedUsageSourceApiHelp = Yalnızca Zed düzenleyici kimlik bilgisi kullanılır. +ProviderZedUsageSourceWebHelp = Token harcaması yalnızca seçilen zed.dev tarayıcı oturumundan veya kaydedilmiş Cookie başlığından okunur; düzenleyici kimlik bilgisine geri dönülmez. ProviderNoCodexAccountsDetected = Henüz Codex hesabı algılanmadı. ProviderCodexAutoImportHelp = Otomatik olarak panel ekleri için tarayıcı çerezlerini içe aktarır. ProviderCodexHistoryHelp = Tempo tahminlerini kişiselleştirmek için yerel Codex kullanım geçmişini (8 hafta) depolar. diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 98dd870db8..f4d16b22ac 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -263,6 +263,10 @@ ProviderOptionsTitle = 选项 MenuBarMetric = 菜单栏和浮动栏指标 MenuBarMetricHelper = 选择用于菜单栏和浮动栏百分比的时间窗口。 UsageSource = 用量来源 +ProviderZedBrowserSession = 浏览器会话 +ProviderZedUsageSourceAutoHelp = 使用 Zed 编辑器凭据；仅在选择“浏览器会话”时使用浏览器账单。 +ProviderZedUsageSourceApiHelp = 仅使用 Zed 编辑器凭据。 +ProviderZedUsageSourceWebHelp = 仅从所选的 zed.dev 浏览器会话或已保存的 Cookie 标头读取令牌支出；不会回退到编辑器凭据。 ProviderNoCodexAccountsDetected = 尚未检测到 Codex 账号。 ProviderCodexAutoImportHelp = 自动导入浏览器 Cookie 以补充仪表盘信息。 ProviderCodexHistoryHelp = 在本地保存 Codex 用量历史（8 周），用于个性化 Pace 预测。 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index f893d5b423..9156e0a7e7 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -263,6 +263,10 @@ ProviderOptionsTitle = 選項 MenuBarMetric = 選單欄和浮動列指標 MenuBarMetricHelper = 選擇用於選單欄和浮動列百分比的時間範圍。 UsageSource = 用量來源 +ProviderZedBrowserSession = 瀏覽器工作階段 +ProviderZedUsageSourceAutoHelp = 使用 Zed 編輯器憑證；只有選擇「瀏覽器工作階段」時才使用瀏覽器帳單。 +ProviderZedUsageSourceApiHelp = 僅使用 Zed 編輯器憑證。 +ProviderZedUsageSourceWebHelp = 僅從所選的 zed.dev 瀏覽器工作階段或已儲存的 Cookie 標頭讀取 Token 支出；不會回退到編輯器憑證。 ProviderNoCodexAccountsDetected = 尚未檢測到 Codex 賬號。 ProviderCodexAutoImportHelp = 自動匯入瀏覽器 Cookie 以補充儀表板資訊。 ProviderCodexHistoryHelp = 在本機儲存 Codex 用量歷史（8 周），用於個性化 Pace 預測。 diff --git a/rust/src/providers/zed/mod.rs b/rust/src/providers/zed/mod.rs index 75e64808ed..9416a7eacc 100644 --- a/rust/src/providers/zed/mod.rs +++ b/rust/src/providers/zed/mod.rs @@ -6,8 +6,8 @@ use async_trait::async_trait; use reqwest::{Client, StatusCode}; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - ProviderStateKind, SourceMode, + FetchContext, ManualEmptyCookiePolicy, Provider, ProviderError, ProviderFetchResult, + ProviderId, ProviderMetadata, ProviderStateKind, SourceMode, }; use crate::providers::{BoundedBodyError, read_bounded_response}; @@ -128,6 +128,10 @@ impl Provider for ZedProvider { true } + fn manual_empty_cookie_policy(&self) -> ManualEmptyCookiePolicy { + ManualEmptyCookiePolicy::FailClosedWeb + } + fn error_state_kind(&self, error: &ProviderError) -> ProviderStateKind { match error { ProviderError::Other(message) if message == SESSION_EXPIRED => { From fc9a0c176fdd1956d91427617e051f610dc187b9 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:46:15 +0700 Subject: [PATCH 125/259] Address thermo review --- .../src/lib/currencyCatalog.generated.ts | 2 +- rust/src/currency.rs | 29 +++++++------------ 2 files changed, 12 insertions(+), 19 deletions(-) diff --git a/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts b/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts index e1ec5b1f2a..8ca3e7089f 100644 --- a/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts +++ b/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts @@ -1,5 +1,5 @@ // Generated from rust/src/currency.rs (CURRENCIES). Do not edit by hand. -// Regenerate: UPDATE_CURRENCY_CATALOG=1 cargo test -p codexbar currency_catalog +// Regenerate (PowerShell): $env:UPDATE_CURRENCY_CATALOG = '1'; cargo test -p codexbar currency_catalog export const CURRENCY_CATALOG = [ { code: "USD", symbol: "$", fallbackRate: 1 }, { code: "GBP", symbol: "£", fallbackRate: 0.79 }, diff --git a/rust/src/currency.rs b/rust/src/currency.rs index 608dc9978a..2670016b90 100644 --- a/rust/src/currency.rs +++ b/rust/src/currency.rs @@ -22,7 +22,7 @@ const fn currency(code: &'static str, symbol: &'static str, fallback_rate: f64) /// Single source of truth for picker order, picker symbols, and offline rates. /// The frontend copy (`apps/desktop-tauri/src/lib/currencyCatalog.generated.ts`) -/// is generated from this table by `render_typescript_catalog`. +/// is checked against this table by the catalog sync test. pub const CURRENCIES: &[CurrencyInfo] = &[ currency("USD", "$", 1.0), currency("GBP", "£", 0.79), @@ -59,26 +59,22 @@ pub fn is_supported_currency(code: &str) -> bool { CURRENCIES.iter().any(|currency| currency.code == code) } -/// Picker label in the catalog's `CODE (symbol)` form. -pub fn picker_label(code: &str) -> Option { - CURRENCIES - .iter() - .find(|currency| currency.code == code) - .map(|currency| format!("{} ({})", currency.code, currency.symbol)) -} - /// Renders the catalog as the TypeScript module the settings picker and -/// offline converter import, so both sides share this table. -pub fn render_typescript_catalog() -> String { +/// offline converter import, so both sides share this table. This is only +/// needed by the sync test and the opt-in regeneration path. +#[cfg(test)] +fn render_typescript_catalog() -> String { let mut out = String::from( "// Generated from rust/src/currency.rs (CURRENCIES). Do not edit by hand.\n\ - // Regenerate: UPDATE_CURRENCY_CATALOG=1 cargo test -p codexbar currency_catalog\n\ + // Regenerate (PowerShell): $env:UPDATE_CURRENCY_CATALOG = '1'; cargo test -p codexbar currency_catalog\n\ export const CURRENCY_CATALOG = [\n", ); for currency in CURRENCIES { + let code = serde_json::to_string(currency.code).expect("currency codes serialize"); + let symbol = serde_json::to_string(currency.symbol).expect("currency symbols serialize"); out.push_str(&format!( - " {{ code: \"{}\", symbol: \"{}\", fallbackRate: {} }},\n", - currency.code, currency.symbol, currency.fallback_rate + " {{ code: {code}, symbol: {symbol}, fallbackRate: {} }},\n", + currency.fallback_rate )); } out.push_str("] as const;\n"); @@ -268,9 +264,6 @@ mod tests { (*code, *symbol, *rate) ); } - assert_eq!(picker_label("VND").as_deref(), Some("VND (₫)")); - assert_eq!(picker_label("TRY").as_deref(), Some("TRY (₺)")); - assert_eq!(picker_label("BTC"), None); } #[test] @@ -312,7 +305,7 @@ mod tests { assert_eq!( actual.replace("\r\n", "\n"), expected, - "regenerate with UPDATE_CURRENCY_CATALOG=1 cargo test -p codexbar currency_catalog" + "set UPDATE_CURRENCY_CATALOG=1, then rerun `cargo test -p codexbar currency_catalog`" ); } } From 27ff74ec1343403c82c47c6b17eda491300e6233 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:48:14 +0700 Subject: [PATCH 126/259] Address thermo review --- rust/src/providers/openrouter/mod.rs | 14 +++++--------- rust/src/providers/openrouter/tests.rs | 14 -------------- rust/src/settings/tests.rs | 14 ++++++++++++++ 3 files changed, 19 insertions(+), 23 deletions(-) diff --git a/rust/src/providers/openrouter/mod.rs b/rust/src/providers/openrouter/mod.rs index a9b3f0880f..b9e244db0e 100755 --- a/rust/src/providers/openrouter/mod.rs +++ b/rust/src/providers/openrouter/mod.rs @@ -172,15 +172,11 @@ impl OpenRouterProvider { return Ok(key.to_string()); } - match keyring::Entry::new(OPENROUTER_CREDENTIAL_TARGET, "api_token") { - Ok(entry) => match entry.get_password() { - Ok(token) => Ok(token), - Err(_) => std::env::var("OPENROUTER_API_KEY") - .map_err(|_| ProviderError::NotInstalled(MISSING_API_KEY_MESSAGE.to_string())), - }, - Err(_) => std::env::var("OPENROUTER_API_KEY") - .map_err(|_| ProviderError::NotInstalled(MISSING_API_KEY_MESSAGE.to_string())), - } + keyring::Entry::new(OPENROUTER_CREDENTIAL_TARGET, "api_token") + .ok() + .and_then(|entry| entry.get_password().ok()) + .or_else(|| std::env::var("OPENROUTER_API_KEY").ok()) + .ok_or_else(|| ProviderError::NotInstalled(MISSING_API_KEY_MESSAGE.to_string())) } fn configured_management_key() -> Option { diff --git a/rust/src/providers/openrouter/tests.rs b/rust/src/providers/openrouter/tests.rs index 3ebc28e53c..3a1db06431 100644 --- a/rust/src/providers/openrouter/tests.rs +++ b/rust/src/providers/openrouter/tests.rs @@ -397,17 +397,3 @@ fn missing_key_message_explains_primary_and_management_fields() { "Enter a regular API key or a Management API key in the API key field, or set OPENROUTER_API_KEY. In Settings, the optional Management API key field does not replace it." ); } - -#[test] -fn api_key_field_help_says_management_keys_are_accepted() { - let info = crate::settings::get_api_key_providers() - .into_iter() - .find(|info| info.id == ProviderId::OpenRouter) - .expect("OpenRouter api key metadata"); - assert_eq!( - info.api_key_help, - Some( - "Required. Enter a regular API key or a Management API key here. Management keys also enable account Activity on the official OpenRouter API." - ) - ); -} diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..d78b52b36f 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -548,6 +548,20 @@ fn test_api_key_provider_catalog_includes_token_providers() { } } +#[test] +fn openrouter_api_key_help_explains_management_keys() { + let info = get_api_key_providers() + .into_iter() + .find(|info| info.id == ProviderId::OpenRouter) + .expect("OpenRouter api key metadata"); + assert_eq!( + info.api_key_help, + Some( + "Required. Enter a regular API key or a Management API key here. Management keys also enable account Activity on the official OpenRouter API." + ) + ); +} + #[test] fn test_t3_chat_is_cookie_configured_not_api_key_configured() { let providers = get_api_key_providers(); From 29b34cc1f4d98d918c52724655982fc5c546a433 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:58:10 +0700 Subject: [PATCH 127/259] Address thermo review --- .../src-tauri/src/commands/chart.rs | 19 ++++++- .../src-tauri/src/commands/chart/tests.rs | 30 +++++++++-- .../src-tauri/src/commands/spend_contract.rs | 13 ++--- .../src-tauri/src/commands/usage_spend.rs | 50 ++++++++++--------- .../src/commands/usage_spend/tests.rs | 30 +++++++++++ .../src/components/MenuCard.test.tsx | 6 ++- .../src/components/MenuCardDetails.tsx | 11 ++-- .../src/floatbar/FloatBar.test.tsx | 7 ++- apps/desktop-tauri/src/floatbar/FloatBar.tsx | 23 +++------ apps/desktop-tauri/src/i18n/keys.ts | 6 --- apps/desktop-tauri/src/lib/costPeriod.test.ts | 16 ------ apps/desktop-tauri/src/lib/costPeriod.ts | 7 --- .../src/lib/usageSpendSharing.test.ts | 16 ++++-- .../src/lib/usageSpendSharing.ts | 27 +++------- .../src/surfaces/TrayPanel.test.tsx | 12 ++++- apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 26 +++------- .../settings/tabs/CostPeriodControl.tsx | 28 +++++------ .../settings/tabs/UsageSpendTab.test.tsx | 29 +++++++++-- .../surfaces/settings/tabs/UsageSpendTab.tsx | 10 +--- apps/desktop-tauri/src/types/bridge.ts | 16 +++--- rust/src/cost_reporting_period.rs | 38 ++++---------- rust/src/locale.rs | 6 --- rust/src/locale/en-US.ftl | 6 --- 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/tests.rs | 8 --- rust/src/locale/tr-TR.ftl | 3 -- rust/src/locale/zh-CN.ftl | 3 -- rust/src/locale/zh-TW.ftl | 3 -- rust/src/providers/bedrock/daily.rs | 2 +- 32 files changed, 223 insertions(+), 240 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/chart.rs b/apps/desktop-tauri/src-tauri/src/commands/chart.rs index fa49349028..9cdc0e7dfe 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/chart.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/chart.rs @@ -555,6 +555,14 @@ fn local_usage_period_identity() -> String { current_reporting_period().identity(Utc::now(), CostTimeZone::Local) } +/// Cache a completed scan under the reporting window it actually used. The +/// summary timestamp preserves the month for month-to-date identities. +fn local_usage_summary_period_identity(summary: &ProviderLocalUsageSummary) -> Option { + let period = CostReportingPeriod::parse(&summary.reporting_period)?; + let scanned_at = DateTime::::from_timestamp_millis(summary.token_cost_updated_at_ms)?; + Some(period.identity(scanned_at, CostTimeZone::Local)) +} + fn local_usage_cache() -> &'static Mutex> { static CACHE: OnceLock>> = OnceLock::new(); CACHE.get_or_init(|| Mutex::new(HashMap::new())) @@ -662,7 +670,16 @@ fn load_local_usage_summary_cached( } fn store_local_usage_summary(provider_id: &str, summary: Option) { - let period_identity = local_usage_period_identity(); + let period_identity = match summary.as_ref() { + Some(summary) => { + let Some(period_identity) = local_usage_summary_period_identity(summary) else { + tracing::warn!("Skipping local usage cache entry with invalid period metadata"); + return; + }; + period_identity + } + None => local_usage_period_identity(), + }; if let Ok(mut guard) = local_usage_cache().lock() { guard.insert( provider_id.to_string(), diff --git a/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs index a9a49dbdf4..f0d19e9adf 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs @@ -1,10 +1,10 @@ use super::{ CostFetchFailure, ProviderLocalUsageSummary, cost_fetch_failure_allows_early_retry, - local_usage_summary_from_cost_summary, localized_estimate_note, muse_local_usage_summary, - token_cost_cache_is_fresh, + local_usage_summary_from_cost_summary, local_usage_summary_period_identity, + localized_estimate_note, muse_local_usage_summary, token_cost_cache_is_fresh, }; use crate::commands::is_provider_cache_fresh; -use codexbar::cost_reporting_period::CostReportingPeriod; +use codexbar::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; use codexbar::cost_scanner::CostSummary; use codexbar::providers::muse::local_usage::{DailyUsage, Report}; use codexbar::settings::Language; @@ -59,6 +59,30 @@ fn local_usage_summary_serializes_token_cost_timestamp() { ); } +#[test] +fn local_usage_cache_identity_uses_the_summary_period_and_timestamp() { + let scanned_at = chrono::DateTime::parse_from_rfc3339("2026-09-30T23:00:00Z") + .expect("valid timestamp") + .with_timezone(&chrono::Utc); + let summary = ProviderLocalUsageSummary { + today_cost: None, + thirty_day_cost: None, + thirty_day_tokens: None, + period_cost: None, + period_tokens: None, + reporting_period: "month-to-date".to_string(), + latest_tokens: None, + top_model: None, + estimate_note: String::new(), + token_cost_updated_at_ms: scanned_at.timestamp_millis(), + }; + + assert_eq!( + local_usage_summary_period_identity(&summary), + Some(CostReportingPeriod::MonthToDate.identity(scanned_at, CostTimeZone::Local)) + ); +} + #[test] fn muse_local_usage_summary_exposes_complete_tokens_without_cost() { let report = Report { 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 ee926ab736..02ae85c413 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs @@ -5,13 +5,11 @@ use codexbar::cost_scanner::CostScanner; use codexbar::settings::Settings; use codexbar::spend_contract::{SpendContract, build_contract_from_period_summary}; -/// `period` is a raw reporting period (`rolling:N`, `month-to-date`, `all`). -/// `history_days` is the legacy rolling count (`0` meant All). With neither, -/// the saved `cost_reporting_period` applies. +/// `period` is a raw reporting period (`rolling:N`, `month-to-date`, `all`); +/// without a valid one the saved `cost_reporting_period` applies. #[tauri::command] pub async fn get_spend_contract( provider_id: String, - history_days: Option, period: Option, include_open_codex: Option, ) -> Result { @@ -24,11 +22,8 @@ pub async fn get_spend_contract( let include_import = include_open_codex.unwrap_or(false) && provider == "codex"; tauri::async_runtime::spawn_blocking(move || { let settings = Settings::load(); - let period = CostReportingPeriod::resolve_request( - period.as_deref(), - history_days, - settings.cost_reporting_period, - ); + let period = + CostReportingPeriod::resolve_request(period.as_deref(), settings.cost_reporting_period); let scanner = CostScanner::for_period(period); let summary = match provider.as_str() { "codex" => scanner.scan_codex(), 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 3e9e2ec118..3b4dcebdc3 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -193,7 +193,6 @@ struct BuiltUsageSpendSummary { #[tauri::command] pub async fn get_usage_spend_summary( state: State<'_, Mutex>, - history_days: Option, period: Option, force_refresh: Option, ) -> Result { @@ -202,11 +201,9 @@ pub async fn get_usage_spend_summary( guard.provider_cache.clone() }; - // An explicit `period` wins, then the legacy `history_days` (0 = All), - // then the saved History window. + // An explicit `period` wins; otherwise use the saved History window. let period = CostReportingPeriod::resolve_request( period.as_deref(), - history_days, codexbar::settings::Settings::load().cost_reporting_period, ); let force_refresh = force_refresh.unwrap_or(false); @@ -398,8 +395,7 @@ fn build_usage_spend_summary( }; let mut codex_scan_options = codex_scan_options; codex_scan_options.include_pi_sessions = include_pi_in_native; - // The selected period reuses the fixed 7d/30d scan when it is the same - // rolling window; any other period costs one extra scan per provider. + // Any period other than 7d/30d costs one extra scan per provider. let ( (codex_7_summary, codex_30_summary, codex_period_summary), (claude_7_summary, claude_30_summary, claude_period_summary), @@ -412,13 +408,11 @@ fn build_usage_spend_summary( let thirty = CostScanner::new(30) .with_options(codex_scan_options) .scan_codex(); - let selected = match period { - CostReportingPeriod::Rolling(7) => seven.clone(), - CostReportingPeriod::Rolling(30) => thirty.clone(), - _ => CostScanner::for_period(period) + let selected = selected_period_scan(period, &seven, &thirty, || { + CostScanner::for_period(period) .with_options(codex_scan_options) - .scan_codex(), - }; + .scan_codex() + }); (seven, thirty, selected) }); let claude = scope.spawn(|| { @@ -426,22 +420,18 @@ fn build_usage_spend_summary( .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native); let thirty = CostScanner::new(30) .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native); - let selected = match period { - CostReportingPeriod::Rolling(7) => seven.clone(), - CostReportingPeriod::Rolling(30) => thirty.clone(), - _ => CostScanner::for_period(period) - .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native), - }; + let selected = selected_period_scan(period, &seven, &thirty, || { + CostScanner::for_period(period) + .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native) + }); (seven, thirty, selected) }); let pi = scope.spawn(|| { let seven = CostScanner::new(7).scan_pi(); let thirty = CostScanner::new(30).scan_pi(); - let selected = match period { - CostReportingPeriod::Rolling(7) => seven.clone(), - CostReportingPeriod::Rolling(30) => thirty.clone(), - _ => CostScanner::for_period(period).scan_pi(), - }; + let selected = selected_period_scan(period, &seven, &thirty, || { + CostScanner::for_period(period).scan_pi() + }); (seven, thirty, selected) }); ( @@ -731,6 +721,20 @@ fn build_usage_spend_summary( } } +/// Reuse the fixed scan when it already covers the selected period. +fn selected_period_scan( + period: CostReportingPeriod, + seven: &T, + thirty: &T, + scan: impl FnOnce() -> T, +) -> T { + match period { + CostReportingPeriod::Rolling(7) => seven.clone(), + CostReportingPeriod::Rolling(30) => thirty.clone(), + _ => scan(), + } +} + /// 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. diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs index b42400eeb3..7187b98187 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs @@ -1,4 +1,5 @@ use super::*; +use codexbar::cost_scanner::CostSummary; #[test] fn invalidated_owner_clears_orphaned_indexing_activity() { @@ -49,6 +50,35 @@ fn privacy_mode_is_part_of_usage_spend_cache_identity() { assert_ne!(public, private); } +#[test] +fn selected_period_scan_reuses_fixed_windows_and_scans_other_periods() { + let seven = CostSummary { + total_cost_usd: 7.0, + ..CostSummary::default() + }; + let thirty = CostSummary { + total_cost_usd: 30.0, + ..CostSummary::default() + }; + let mut scanned = false; + let selected = selected_period_scan(CostReportingPeriod::Rolling(7), &seven, &thirty, || { + scanned = true; + CostSummary::default() + }); + assert_eq!(selected.total_cost_usd, 7.0); + assert!(!scanned); + + let selected = selected_period_scan(CostReportingPeriod::Rolling(90), &seven, &thirty, || { + scanned = true; + CostSummary { + total_cost_usd: 90.0, + ..CostSummary::default() + } + }); + assert_eq!(selected.total_cost_usd, 90.0); + assert!(scanned); +} + fn identity(period: CostReportingPeriod, now: &str) -> String { let now = chrono::DateTime::parse_from_rfc3339(now) .expect("valid instant") diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index ffd2f94677..75addfb86e 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -127,11 +127,10 @@ describe("MenuCard", () => { PanelFiveHours: "5h", PanelOnPaceBudget: "On-pace budget", PanelReserveSuffix: "in reserve", - PanelThirtyDayCost: "30d cost", - PanelThirtyDayTokens: "30d tokens", PanelPeriodCost: "{} cost", PanelPeriodTokens: "{} tokens", CostPeriodShortMonthToDate: "MTD", + CostPeriodShortDays: "{}d", PanelTodayBudget: "today", PanelUsedSuffix: "used", ResetsInHoursMinutes: "Resets in {}h {}m", @@ -169,6 +168,9 @@ describe("MenuCard", () => { todayCost: null, thirtyDayCost: 1.23, thirtyDayTokens: 584_000, + periodCost: 1.23, + periodTokens: 584_000, + reportingPeriod: "rolling:30", latestTokens: null, topModel: "glim-4.6", estimateNote: "Estimated from local logs", diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 930a04a0f9..7ba3428353 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -116,10 +116,7 @@ function LocalUsageBlock({ const isCodex = providerId === "codex"; const isMuse = providerId === "muse"; // The selected History window; the histogram below stays a fixed 30 days. - // Payloads without a period (older cache entries) fall back to 30 days. - const period = summary.reportingPeriod || null; - const periodCost = period ? summary.periodCost : summary.thirtyDayCost; - const periodTokens = period ? summary.periodTokens : summary.thirtyDayTokens; + const { reportingPeriod, periodCost, periodTokens } = summary; const visibleHistory = costHistory.slice(-30); const maxCost = Math.max( ...visibleHistory.flatMap((point) => (point.value == null ? [] : [point.value])), @@ -144,7 +141,7 @@ function LocalUsageBlock({ {!isMuse && (
- {period ? periodCostLabel(period, t) : t("PanelThirtyDayCost")} + {periodCostLabel(reportingPeriod, t)} {periodCost != null ? formatCurrency(periodCost, "USD") : "—"} @@ -153,9 +150,9 @@ function LocalUsageBlock({ )}
- {period ? periodTokensLabel(period, t) : t("PanelThirtyDayTokens")} + {periodTokensLabel(reportingPeriod, t)} - {formatCompactCount(periodTokens ?? null)} + {formatCompactCount(periodTokens)}
{!isMuse && (
diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index d7b07da417..03d8addd1b 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -230,7 +230,6 @@ describe("FloatBar", () => { ProviderIssueSessionExpired: "Session expired", ProviderIssueLocalRuntimeOffline: "Local runtime offline", ProviderIssueUnknown: "Usage unavailable", - FloatBarThirtyDayShort: "30d", CostPeriodShortMonthToDate: "MTD", CostPeriodShortDays: "{}d", FloatBarNoProviders: "No providers", @@ -390,6 +389,9 @@ describe("FloatBar", () => { todayCost: 1.25, thirtyDayCost: 12.5, thirtyDayTokens: 1000, + periodCost: 12.5, + periodTokens: 1000, + reportingPeriod: "rolling:30", latestTokens: 200, topModel: "gpt-5", estimateNote: "Estimated from local logs", @@ -411,6 +413,9 @@ describe("FloatBar", () => { todayCost: 1.25, thirtyDayCost: 12.5, thirtyDayTokens: 1000, + periodCost: 12.5, + periodTokens: 1000, + reportingPeriod: "rolling:30", latestTokens: 200, topModel: "gpt-5", estimateNote: "Estimated from local logs", diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.tsx index d067f49cb4..f38f1eb3b2 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.tsx @@ -93,10 +93,10 @@ type FloatBarCostSummary = { providerId: string; displayName: string; todayCost: number | null; - /** Cost over the selected History window (30 days when the backend sends none). */ + /** Cost over the selected History window. */ periodCost: number | null; - /** Raw History window, or null for payloads that only carry 30 days. */ - period: string | null; + /** Raw History window `periodCost` covers. */ + period: string; }; type FloatBarCostTarget = { @@ -109,13 +109,8 @@ function providerCostKey(provider: ProviderUsageSnapshot): string { return `${provider.providerId}:${provider.accountEmail ?? ""}`; } -/** Cost for the selected History window; older payloads only carry 30 days. */ -function windowCost(summary: ProviderLocalUsageSummary): number | null { - return (summary.reportingPeriod ? summary.periodCost : summary.thirtyDayCost) ?? null; -} - function hasLocalCost(summary: ProviderLocalUsageSummary | null): summary is ProviderLocalUsageSummary { - return summary != null && (summary.todayCost != null || windowCost(summary) != null); + return summary?.todayCost != null || summary?.periodCost != null; } function formatUsd(value: number | null): string | null { @@ -391,8 +386,8 @@ export default function FloatBar({ state }: { state: BootstrapState }) { providerId: target.providerId, displayName: target.displayName, todayCost: localUsage.todayCost, - periodCost: windowCost(localUsage), - period: localUsage.reportingPeriod || null, + periodCost: localUsage.periodCost, + period: localUsage.reportingPeriod, } satisfies FloatBarCostSummary; }), ) @@ -533,11 +528,7 @@ export default function FloatBar({ state }: { state: BootstrapState }) { summary={summary} scale={scale} todayLabel={t("PanelToday")} - periodLabel={ - summary.period - ? costPeriodShortLabel(summary.period, t) - : t("FloatBarThirtyDayShort") - } + periodLabel={costPeriodShortLabel(summary.period, t)} estimateLabel={t("OverviewSpendEstimate")} /> ))} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 1b6feb649a..6c07f03044 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -396,7 +396,6 @@ export const ALL_LOCALE_KEYS = [ "UsageSpendTitle", "UsageSpendCaption", "UsageSpendModels", - "UsageSpendAllTime", "UsageSpendOpenCodexImport", "UsageSpendHideNativeCodex", "UsageSpendSpend", @@ -416,9 +415,7 @@ export const ALL_LOCALE_KEYS = [ "UsageSpendHourlyActivity", "UsageSpendRequests", "UsageSpendTokens", - "UsageSpendAllTimeHistory", "UsageSpendCustomPricing", - "OverviewSpendTitle", "OverviewSpendProviderCoverage", "OverviewSpendEstimate", "UsageSpendProjects", @@ -641,8 +638,6 @@ export const ALL_LOCALE_KEYS = [ "PanelMenu", "PanelCopied", "PanelToday", - "PanelThirtyDayCost", - "PanelThirtyDayTokens", "PanelLatestTokens", "PanelThirtyDayCostHistogram", "PanelTopModelPrefix", @@ -662,7 +657,6 @@ export const ALL_LOCALE_KEYS = [ "PanelReserveLastsUntilReset", "PanelReserveRunsOutInDaysHours", "PanelReserveRunsOutInHours", - "FloatBarThirtyDayShort", "FloatBarNoProviders", "FloatBarRemainingSuffix", "FloatBarShowCost", diff --git a/apps/desktop-tauri/src/lib/costPeriod.test.ts b/apps/desktop-tauri/src/lib/costPeriod.test.ts index 2064aeb10a..fd8547865f 100644 --- a/apps/desktop-tauri/src/lib/costPeriod.test.ts +++ b/apps/desktop-tauri/src/lib/costPeriod.test.ts @@ -10,7 +10,6 @@ import { normalizeCostPeriod, parseCostPeriod, periodCostLabel, - periodFromLegacyDays, periodTokensLabel, rollingDays, } from "./costPeriod"; @@ -56,21 +55,6 @@ describe("normalizeCostPeriod", () => { }); }); -describe("periodFromLegacyDays", () => { - it("migrates the old day selection, where 0 meant all", () => { - expect(periodFromLegacyDays(0)).toBe("all"); - expect(periodFromLegacyDays(7)).toBe("rolling:7"); - expect(periodFromLegacyDays(30)).toBe("rolling:30"); - expect(periodFromLegacyDays(1000)).toBe("rolling:365"); - }); - - it("uses the default for missing or negative counts", () => { - expect(periodFromLegacyDays(undefined)).toBe("rolling:30"); - expect(periodFromLegacyDays(null)).toBe("rolling:30"); - expect(periodFromLegacyDays(-1)).toBe("rolling:30"); - }); -}); - describe("customPeriodRaw", () => { it("accepts integers in 1..=365 only", () => { expect(customPeriodRaw("1")).toBe("rolling:1"); diff --git a/apps/desktop-tauri/src/lib/costPeriod.ts b/apps/desktop-tauri/src/lib/costPeriod.ts index c79b7a970e..f03f1d8b64 100644 --- a/apps/desktop-tauri/src/lib/costPeriod.ts +++ b/apps/desktop-tauri/src/lib/costPeriod.ts @@ -57,13 +57,6 @@ export function normalizeCostPeriod(raw: string | null | undefined): string { return parsed ? toRaw(parsed) : DEFAULT_COST_PERIOD; } -/** Migrate the pre-0.67.0 day selection (0 = all, otherwise a rolling day count). */ -export function periodFromLegacyDays(days: number | null | undefined): string { - if (days == null || !Number.isFinite(days) || days < 0) return DEFAULT_COST_PERIOD; - if (days === 0) return "all"; - return normalizeCostPeriod(`rolling:${Math.floor(days)}`); -} - /** Persisted form for a typed custom day count; null unless it is an integer in 1..=365. */ export function customPeriodRaw(input: string): string | null { const text = input.trim(); diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index 9f838100c2..16f400a597 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -35,11 +35,14 @@ describe("usage spend sharing", () => { expect( usageSpendShareFooter({ contract: {} as SpendContract, + reportingPeriod: "rolling:30", rows: [{ providerId: "codex", displayName: "Codex", sevenDay: null, thirtyDay: null, + periodCost: null, + periodTokens: null, currency: "USD", source: "local", includedInOverview: true, @@ -53,6 +56,7 @@ describe("usage spend sharing", () => { it("keeps hidden sources out of the Overview share summary", () => { const summary: UsageSpendSummary = { contract: {} as SpendContract, + reportingPeriod: "rolling:30", reportingDay: "2026-09-19", dashboardTimezone: "UTC", rows: [ @@ -61,6 +65,8 @@ describe("usage spend sharing", () => { displayName: "Codex", sevenDay: 1, thirtyDay: 2, + periodCost: 1, + periodTokens: null, currency: "USD", source: "local", includedInOverview: true, @@ -70,6 +76,8 @@ describe("usage spend sharing", () => { displayName: "Claude", sevenDay: 3, thirtyDay: 4, + periodCost: 3, + periodTokens: null, currency: "USD", source: "hidden", includedInOverview: false, @@ -101,6 +109,8 @@ describe("usage spend sharing", () => { displayName: "Codex", sevenDay: 1, thirtyDay: 2, + periodCost: 2, + periodTokens: null, currency: "USD", source: "local", includedInOverview: true, @@ -110,7 +120,7 @@ describe("usage spend sharing", () => { ); // Nothing outside the drawn cells may be read by the renderer, and no // UsageSpendRow field may carry account-identity data (no email/org). - expect(Object.keys(row).filter((key) => /email|org|token|account/i.test(key))).toEqual([]); + expect(Object.keys(row).filter((key) => /email|org|account|secret|(access|auth|api|refresh)token/i.test(key))).toEqual([]); expect(unsafeKeys).toEqual(["providerId", "includedInOverview"]); }); @@ -126,10 +136,6 @@ describe("usage spend sharing", () => { source: "local", includedInOverview: true, }; - const plain = usageSpendShareColumns(); - expect(plain.headers).toEqual(["Provider", "7 days", "30 days", "Currency", "Source"]); - expect(plain.cellsFor(row)).toHaveLength(plain.headers.length); - const withPeriod = usageSpendShareColumns("Month to date"); expect(withPeriod.headers).toEqual([ "Provider", diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 5d80e3e4d9..7f4957b0a8 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -140,31 +140,20 @@ export function formatSpendMetric( } /** - * Share-card columns. The fixed 7 and 30 day columns stay for compatibility; - * when `periodLabel` is given, a column for the selected History window - * (read from `periodCost` / `periodTokens`) follows the provider name. + * Share-card columns: the selected History window (`periodCost` / `periodTokens`) + * follows the provider name, then the fixed 7 and 30 day compatibility columns. */ -export function usageSpendShareColumns(periodLabel?: string): { +export function usageSpendShareColumns(periodLabel: string): { headers: string[]; colW: number[]; cellsFor: (row: UsageSpendRow) => string[]; } { - const withPeriod = Boolean(periodLabel); return { - headers: [ - "Provider", - ...(withPeriod ? [periodLabel as string] : []), - "7 days", - "30 days", - "Currency", - "Source", - ], - colW: withPeriod ? [150, 130, 100, 100, 70, 130] : [160, 100, 100, 80, 160], + headers: ["Provider", periodLabel, "7 days", "30 days", "Currency", "Source"], + colW: [150, 130, 100, 100, 70, 130], cellsFor: (row) => [ row.displayName, - ...(withPeriod - ? [formatSpendMetric(row.periodCost, row.periodTokens, row.currency, "tokens")] - : []), + formatSpendMetric(row.periodCost, row.periodTokens, row.currency, "tokens"), formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, "tokens"), formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, "tokens"), row.currency || "USD", @@ -185,7 +174,7 @@ export function usageSpendShareColumns(periodLabel?: string): { export function renderUsageSpendSharePng( summary: UsageSpendSummary, title: string, - periodLabel?: string, + periodLabel: string, ): string { const rows = summary.rows; const pad = 24; @@ -281,7 +270,7 @@ export function shareUsageSpendPng( summary: UsageSpendSummary | null, title: string, filename: string, - periodLabel?: string, + periodLabel: string, ): string | null { if (!summary) return "UsageSpendShareEmpty"; try { diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 8e156a34bf..318cc43b3c 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -320,6 +320,7 @@ describe("TrayPanel provider grid", () => { it("offers an Overview share snapshot using only included spend rows", async () => { tauriMocks.getUsageSpendSummary.mockResolvedValue({ contract: {}, + reportingPeriod: "rolling:30", reportingDay: "2026-09-19", dashboardTimezone: "UTC", rows: [ @@ -328,6 +329,8 @@ describe("TrayPanel provider grid", () => { displayName: "Codex", sevenDay: 1, thirtyDay: 2, + periodCost: 2, + periodTokens: null, currency: "USD", source: "local", includedInOverview: true, @@ -337,6 +340,8 @@ describe("TrayPanel provider grid", () => { displayName: "Claude", sevenDay: 3, thirtyDay: 4, + periodCost: 4, + periodTokens: null, currency: "USD", source: "hidden", includedInOverview: false, @@ -428,7 +433,9 @@ describe("TrayPanel provider grid", () => { PanelAllProviders: "すべてのプロバイダー", PanelAllProvidersShort: "すべて", PanelLatestTokens: "最新トークン", - PanelThirtyDayCost: "30日間のコスト", + CostPeriodShortDays: "{}日", + PanelPeriodCost: "{}間のコスト", + PanelPeriodTokens: "{}間のトークン", PanelTopModelPrefix: "トップモデル", PanelEstimatedFromLocalLogs: "ローカルログから推定", PanelZoom: "ズーム", @@ -446,6 +453,9 @@ describe("TrayPanel provider grid", () => { todayCost: null, thirtyDayCost: 1.23, thirtyDayTokens: 584_000, + periodCost: 1.23, + periodTokens: 584_000, + reportingPeriod: "rolling:30", latestTokens: 1200, topModel: "gpt-5.5", estimateNote: "Estimated from local logs", diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 920f5db80b..5f3f8c9e91 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -1,11 +1,6 @@ import { Fragment, useEffect, useState, type CSSProperties } from "react"; import { getCurrentWindow } from "@tauri-apps/api/window"; -import type { - BootstrapState, - ProviderUsageSnapshot, - UsageSpendRow, - UsageSpendSummary, -} from "../types/bridge"; +import type { BootstrapState, ProviderUsageSnapshot, UsageSpendSummary } from "../types/bridge"; import type { LocaleKey } from "../i18n/keys"; import { costPeriodLabel, costPeriodShortLabel } from "../lib/costPeriod"; import { @@ -373,32 +368,25 @@ function OverviewSpendSummary({ const overviewSummary = summary ? filterUsageSpendSummaryForOverview(summary) : null; if (!overviewSummary) return null; - // The summary names the History window its period columns cover. Payloads - // without one only carry the fixed 30-day column. - const summaryPeriod = overviewSummary.reportingPeriod || null; - const spendOf = (row: UsageSpendRow) => (summaryPeriod ? row.periodCost : row.thirtyDay); - const title = summaryPeriod - ? t("OverviewSpendPeriodTitle").replace("{}", costPeriodShortLabel(summaryPeriod, t)) - : t("OverviewSpendTitle"); + // The summary names the History window its period columns cover. + const summaryPeriod = overviewSummary.reportingPeriod; + const title = t("OverviewSpendPeriodTitle").replace("{}", costPeriodShortLabel(summaryPeriod, t)); const onShare = () => { setShareError(null); const error = shareUsageSpendPng( overviewSummary, title, `codexbar-overview-usage-${overviewSummary.reportingDay}.png`, - summaryPeriod ? costPeriodLabel(summaryPeriod, t) : undefined, + costPeriodLabel(summaryPeriod, t), ); if (error) setShareError(t(error as LocaleKey)); }; const rows = overviewSummary.rows; const summable = rows.filter((row) => (row.currency || "USD") === "USD"); - const known = summable.filter((row) => { - const spend = spendOf(row); - return spend != null && Number.isFinite(spend); - }); + const known = summable.filter((row) => row.periodCost != null && Number.isFinite(row.periodCost)); if (known.length === 0) return null; - const total = known.reduce((sum, row) => sum + (spendOf(row) ?? 0), 0); + const total = known.reduce((sum, row) => sum + (row.periodCost ?? 0), 0); const partial = known.length < rows.length; const formatter = new Intl.NumberFormat(undefined, { style: "currency", currency: "USD", maximumFractionDigits: 2 }); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx index 6229e9ede0..5e0c517a2d 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx @@ -1,4 +1,4 @@ -import { useEffect, useRef, useState } from "react"; +import { useEffect, useState } from "react"; import type { LocaleKey } from "../../../i18n/keys"; import { MAX_ROLLING_DAYS, @@ -14,7 +14,8 @@ const CUSTOM_OPTION = "custom"; /** * History window picker: month to date, all, 1/7/30/90/365 days, or a custom * day count in 1..=365. `value` is the persisted form (`rolling:N`, - * `month-to-date`, `all`). A custom count is only reported once it is valid. + * `month-to-date`, `all`). A custom count is committed on blur or Enter, once + * it is valid, so typing a multi-digit value triggers one rescan. */ export default function CostPeriodControl({ value, @@ -29,13 +30,8 @@ export default function CostPeriodControl({ }) { const [customMode, setCustomMode] = useState(() => !isPresetCostPeriod(value)); const [customText, setCustomText] = useState(() => String(rollingDays(value) ?? "")); - const typedRaw = useRef(null); - // Follow external changes (settings loaded after mount, another window). - // Values this control just emitted while typing a custom count are skipped so - // typing "1" on the way to "14" does not collapse the custom input. useEffect(() => { - if (value === typedRaw.current) return; if (isPresetCostPeriod(value)) { setCustomMode(false); } else { @@ -45,6 +41,10 @@ export default function CostPeriodControl({ }, [value]); const customInvalid = customMode && customText !== "" && customPeriodRaw(customText) === null; + const commitCustom = () => { + const raw = customPeriodRaw(customText); + if (raw && raw !== value) onChange(raw); + }; return (
@@ -66,7 +66,6 @@ export default function CostPeriodControl({ setCustomText(String(rollingDays(value) ?? "")); return; } - typedRaw.current = null; setCustomMode(false); onChange(next); }} @@ -90,13 +89,12 @@ export default function CostPeriodControl({ disabled={disabled} aria-label={t("CostPeriodCustomDays")} aria-invalid={customInvalid} - onChange={(event) => { - const text = event.target.value; - setCustomText(text); - const raw = customPeriodRaw(text); - if (raw && raw !== value) { - typedRaw.current = raw; - onChange(raw); + onChange={(event) => setCustomText(event.target.value)} + onBlur={commitCustom} + onKeyDown={(event) => { + if (event.key === "Enter") { + event.preventDefault(); + commitCustom(); } }} /> diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx index f4221cf77f..7ed5879627 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx @@ -150,7 +150,7 @@ describe("UsageSpendTab History window", () => { expect(await screen.findByRole("columnheader", { name: "Last 90 days" })).toBeTruthy(); }); - it("offers a custom day count and only persists values in 1..=365", async () => { + it("offers a custom day count and persists one valid value on blur", async () => { render(); await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledTimes(1)); @@ -166,22 +166,41 @@ describe("UsageSpendTab History window", () => { fireEvent.change(input, { target: { value: "0" } }); expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); - // Typing "1" on the way to "14" keeps the custom input open. + // Typing "1" on the way to "14" saves nothing and keeps the input open. fireEvent.change(input, { target: { value: "1" } }); - await waitFor(() => - expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:1" }), - ); + expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); expect(screen.getByLabelText("CostPeriodCustomDays")).toBe(input); fireEvent.change(input, { target: { value: "14" } }); + expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); + fireEvent.blur(input); await waitFor(() => expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:14" }), ); + expect(tauriMocks.updateSettings).toHaveBeenCalledTimes(1); await waitFor(() => expect(lastScanPeriod()).toBe("rolling:14")); expect(input.getAttribute("aria-invalid")).toBe("false"); expect(await screen.findByRole("columnheader", { name: "Last 14 days" })).toBeTruthy(); }); + it("commits a valid custom day count on Enter", async () => { + render(); + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledTimes(1)); + + fireEvent.change(screen.getByLabelText("CostPeriodHistoryWindow"), { + target: { value: "custom" }, + }); + const input = screen.getByLabelText("CostPeriodCustomDays") as HTMLInputElement; + fireEvent.change(input, { target: { value: "45" } }); + expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); + + fireEvent.keyDown(input, { key: "Enter" }); + await waitFor(() => + expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:45" }), + ); + expect(tauriMocks.updateSettings).toHaveBeenCalledTimes(1); + }); + it("starts in custom mode for a saved count that is not a preset", async () => { tauriMocks.getSettingsSnapshot.mockResolvedValue({ costReportingPeriod: "rolling:14", diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index 7a66a469bf..9433e0439c 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -19,7 +19,6 @@ import { DEFAULT_COST_PERIOD, costPeriodLabel, normalizeCostPeriod, - periodFromLegacyDays, } from "../../../lib/costPeriod"; import type { TabProps } from "../settingsTabs"; import CostPeriodControl from "./CostPeriodControl"; @@ -311,11 +310,6 @@ export default function UsageSpendTab(_props: TabProps) { -/** Window label for a contract; older payloads only carry a day count (0 = all). */ -function contractPeriodLabel(contract: SpendContract, t: (key: LocaleKey) => string): string { - return costPeriodLabel(contract.reportingPeriod ?? periodFromLegacyDays(contract.historyDays), t); -} - function SpendContractOverview({ contract, t }: { contract: SpendContract; t: (key: LocaleKey) => string }) { const coverage = contract.priceCoverageRatio == null ? t("UsageSpendUnknown") @@ -400,7 +394,7 @@ function ContractModelsPanel({ contract, showAll, onToggleAll, t }: { contract:

{t("UsageSpendModels")}

- {t("UsageSpendModelsPeriodCaption").replace("{}", contractPeriodLabel(contract, t))} + {t("UsageSpendModelsPeriodCaption").replace("{}", costPeriodLabel(contract.reportingPeriod, t))}

{contract.models.length > 8 && ( @@ -454,7 +448,7 @@ function ProjectsPanel({

{t("UsageSpendProjects")}

- {t("UsageSpendProjectsPeriodCaption").replace("{}", contractPeriodLabel(contract, t))} + {t("UsageSpendProjectsPeriodCaption").replace("{}", costPeriodLabel(contract.reportingPeriod, t))} {partial ? ` · ${t("UsageSpendPartialHistory")}` : ""}.

diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 68b214c1bf..237e1ac159 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -389,8 +389,8 @@ export interface UsageSpendRow { sevenDayTokens?: number | null; thirtyDayTokens?: number | null; /** Cost over the selected History window (`UsageSpendSummary.reportingPeriod`). */ - periodCost?: number | null; - periodTokens?: number | null; + periodCost: number | null; + periodTokens: number | null; currency: string; source: string; includedInOverview: boolean; @@ -405,7 +405,7 @@ export interface UsageSpendSummary { rows: UsageSpendRow[]; contract: SpendContract; /** Raw History window the `period*` columns were built for. */ - reportingPeriod?: string; + reportingPeriod: string; reportingDay: string; dashboardTimezone: string; } @@ -530,7 +530,7 @@ export interface SpendContract { providerId: string; historyDays: number; /** Raw History window this contract was built for. */ - reportingPeriod?: string; + reportingPeriod: string; knownCostUsd: number | null; knownZero: boolean; provenance: CostProvenance; @@ -836,10 +836,10 @@ export interface ProviderLocalUsageSummary { /** Always the trailing 30 days. */ thirtyDayCost: number | null; thirtyDayTokens: number | null; - /** Selected History window totals; absent from older backends. */ - periodCost?: number | null; - periodTokens?: number | null; - reportingPeriod?: string; + /** Selected History window totals. */ + periodCost: number | null; + periodTokens: number | null; + reportingPeriod: string; latestTokens: number | null; topModel: string | null; estimateNote: string; diff --git a/rust/src/cost_reporting_period.rs b/rust/src/cost_reporting_period.rs index 56b6071025..8f1fb87cd5 100644 --- a/rust/src/cost_reporting_period.rs +++ b/rust/src/cost_reporting_period.rs @@ -241,17 +241,10 @@ impl CostReportingPeriod { self.scan_days(now).clamp(1, MAX_ROLLING_DAYS) } - /// Pick the period a request runs with: an explicit raw `period` wins, - /// then a legacy day count (`0` meant All in the old UI), then `saved`. - pub fn resolve_request(period: Option<&str>, legacy_days: Option, saved: Self) -> Self { - if let Some(parsed) = period.and_then(Self::parse) { - return parsed; - } - match legacy_days { - Some(0) => Self::AllAvailable, - Some(days) => Self::rolling(days), - None => saved, - } + /// Pick the period a request runs with: a valid raw `period` wins, + /// otherwise `saved`. + pub fn resolve_request(period: Option<&str>, saved: Self) -> Self { + period.and_then(Self::parse).unwrap_or(saved) } } @@ -281,7 +274,7 @@ mod tests { use chrono_tz::{America::Los_Angeles, Asia::Tokyo}; const LA: CostTimeZone = CostTimeZone::Named(Los_Angeles); - const UTC: CostTimeZone = CostTimeZone::Named(chrono_tz::UTC); + const UTC: CostTimeZone = CostTimeZone::UTC; fn at(value: &str) -> DateTime { DateTime::parse_from_rfc3339(value) @@ -439,25 +432,16 @@ mod tests { } #[test] - fn request_resolution_prefers_explicit_then_legacy_then_saved() { + fn request_resolution_prefers_explicit_then_saved() { let saved = CostReportingPeriod::MonthToDate; - let resolve = |period, legacy| CostReportingPeriod::resolve_request(period, legacy, saved); + let resolve = |period| CostReportingPeriod::resolve_request(period, saved); + assert_eq!(resolve(Some("all")), CostReportingPeriod::AllAvailable); assert_eq!( - resolve(Some("all"), Some(7)), - CostReportingPeriod::AllAvailable - ); - assert_eq!( - resolve(Some("rolling:90"), None), + resolve(Some("rolling:90")), CostReportingPeriod::Rolling(90) ); - assert_eq!(resolve(None, Some(0)), CostReportingPeriod::AllAvailable); - assert_eq!(resolve(None, Some(7)), CostReportingPeriod::Rolling(7)); - assert_eq!( - resolve(Some("bogus"), Some(30)), - CostReportingPeriod::Rolling(30) - ); - assert_eq!(resolve(None, None), saved); - assert_eq!(resolve(Some("bogus"), None), saved); + assert_eq!(resolve(None), saved); + assert_eq!(resolve(Some("bogus")), saved); } #[test] diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e4b9821d21..2ffcaa5df4 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -640,7 +640,6 @@ locale_keys! { UsageSpendTitle, UsageSpendCaption, UsageSpendModels, - UsageSpendAllTime, UsageSpendOpenCodexImport, UsageSpendHideNativeCodex, UsageSpendSpend, @@ -660,9 +659,7 @@ locale_keys! { UsageSpendHourlyActivity, UsageSpendRequests, UsageSpendTokens, - UsageSpendAllTimeHistory, UsageSpendCustomPricing, - OverviewSpendTitle, OverviewSpendProviderCoverage, OverviewSpendEstimate, UsageSpendProjects, @@ -915,8 +912,6 @@ locale_keys! { PanelMenu, PanelCopied, PanelToday, - PanelThirtyDayCost, - PanelThirtyDayTokens, PanelLatestTokens, PanelThirtyDayCostHistogram, PanelTopModelPrefix, @@ -936,7 +931,6 @@ locale_keys! { PanelReserveLastsUntilReset, PanelReserveRunsOutInDaysHours, PanelReserveRunsOutInHours, - FloatBarThirtyDayShort, FloatBarNoProviders, FloatBarRemainingSuffix, FloatBarShowCost, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 72289872bb..22f50ed1d1 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -336,7 +336,6 @@ HooksConfigPathHint = Config path: %APPDATA%\CodexBar\hooks.json (same folder as UsageSpendTitle = Usage & Spend UsageSpendCaption = Local estimated cost history for Codex and Claude (JSONL logs), plus period cost snapshots from other providers when available. UsageSpendModels = Models -UsageSpendAllTime = All time UsageSpendOpenCodexImport = OpenCodex import UsageSpendHideNativeCodex = Hide native Codex when OpenCodex is present UsageSpendSpend = Spend @@ -354,9 +353,7 @@ UsageSpendDefaultPricing = Default pricing UsageSpendHourlyActivity = Hourly activity UsageSpendRequests = requests UsageSpendTokens = tokens -UsageSpendAllTimeHistory = All-time is backed by the latest UsageSpendCustomPricing = custom pricing -OverviewSpendTitle = Usage & Spend · 30d OverviewSpendProviderCoverage = providers have spend OverviewSpendEstimate = local/provider estimates UsageSpendProjects = Projects @@ -578,8 +575,6 @@ PanelZoom = Zoom PanelMenu = Menu PanelCopied = Copied PanelToday = Today -PanelThirtyDayCost = 30d cost -PanelThirtyDayTokens = 30d tokens PanelLatestTokens = Latest tokens PanelThirtyDayCostHistogram = 30 day cost histogram PanelTopModelPrefix = Top model @@ -599,7 +594,6 @@ PanelReserveSuffix = in reserve PanelReserveLastsUntilReset = Lasts until reset PanelReserveRunsOutInDaysHours = Runs out in { "{}" }d { "{}" }h PanelReserveRunsOutInHours = Runs out in { "{}" }h -FloatBarThirtyDayShort = 30d FloatBarNoProviders = No providers FloatBarRemainingSuffix = remaining BannerCheckingForUpdates = Checking for updates… diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index 38dc0b05ae..a5519e0c8b 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -534,8 +534,6 @@ PanelZoom = Zoom PanelMenu = Menú PanelCopied = Copiado PanelToday = Hoy -PanelThirtyDayCost = Costo 30d -PanelThirtyDayTokens = Tokens 30d PanelLatestTokens = Últimos tokens PanelThirtyDayCostHistogram = Histograma de costos 30 días PanelTopModelPrefix = Modelo principal @@ -555,7 +553,6 @@ PanelReserveSuffix = en reserva PanelReserveLastsUntilReset = Dura hasta el reinicio PanelReserveRunsOutInDaysHours = Se agota en { "{}" }d { "{}" }h PanelReserveRunsOutInHours = Se agota en { "{}" }h -FloatBarThirtyDayShort = 30d FloatBarNoProviders = Sin proveedores FloatBarRemainingSuffix = restante BannerCheckingForUpdates = Buscando actualizaciones… diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index b6c2dc0fc8..43b331827b 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -516,8 +516,6 @@ PanelZoom = ズーム PanelMenu = メニュー PanelCopied = コピーしました PanelToday = 今日 -PanelThirtyDayCost = 30日間のコスト -PanelThirtyDayTokens = 30日間のトークン PanelLatestTokens = 最新トークン PanelThirtyDayCostHistogram = 30日間コストのヒストグラム PanelTopModelPrefix = トップモデル @@ -537,7 +535,6 @@ PanelReserveSuffix = 予備 PanelReserveLastsUntilReset = リセットまで持ちます PanelReserveRunsOutInDaysHours = 残り { "{}" }日 { "{}" }時間 PanelReserveRunsOutInHours = 残り { "{}" }時間 -FloatBarThirtyDayShort = 30日 FloatBarNoProviders = プロバイダーなし FloatBarRemainingSuffix = 残り BannerCheckingForUpdates = 更新を確認中… diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 299dcf1a54..110e1124da 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -521,8 +521,6 @@ PanelZoom = 확대/축소 PanelMenu = 메뉴 PanelCopied = 복사됨 PanelToday = 오늘 -PanelThirtyDayCost = 30일 비용 -PanelThirtyDayTokens = 30일 토큰 PanelLatestTokens = 최신 토큰 PanelThirtyDayCostHistogram = 30일 비용 히스토그램 PanelTopModelPrefix = 상위 모델 @@ -542,7 +540,6 @@ PanelReserveSuffix = 예비 PanelReserveLastsUntilReset = 리셋까지 지속 PanelReserveRunsOutInDaysHours = { "{}" }일 { "{}" }시간 후 소진 PanelReserveRunsOutInHours = { "{}" }시간 후 소진 -FloatBarThirtyDayShort = 30일 FloatBarNoProviders = 공급자 없음 FloatBarRemainingSuffix = 남음 BannerCheckingForUpdates = 업데이트 확인 중… diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index f978d540e7..d6e2885101 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -500,8 +500,6 @@ PanelZoom = Увеличить PanelMenu = Меню PanelCopied = Скопировано PanelToday = Сегодня -PanelThirtyDayCost = стоимость 30 пенсов -PanelThirtyDayTokens = жетоны 30 пенсов PanelLatestTokens = Последние токены PanelThirtyDayCostHistogram = Гистограмма затрат за 30 дней PanelTopModelPrefix = Топ-модель @@ -521,7 +519,6 @@ PanelReserveSuffix = в резерве PanelReserveLastsUntilReset = Длится до сброса PanelReserveRunsOutInDaysHours = Заканчивается через { "{}" }d { "{}" }h PanelReserveRunsOutInHours = Заканчивается через { "{}" }h -FloatBarThirtyDayShort = 30 дней FloatBarNoProviders = Нет провайдеров FloatBarRemainingSuffix = оставшееся BannerCheckingForUpdates = Проверка обновлений… diff --git a/rust/src/locale/tests.rs b/rust/src/locale/tests.rs index 9daf93d809..a2a57fd598 100644 --- a/rust/src/locale/tests.rs +++ b/rust/src/locale/tests.rs @@ -146,8 +146,6 @@ fn test_japanese_tray_panel_locale_values_are_translated() { (LocaleKey::PanelZoom, "ズーム"), (LocaleKey::PanelMenu, "メニュー"), (LocaleKey::PanelToday, "今日"), - (LocaleKey::PanelThirtyDayCost, "30日間のコスト"), - (LocaleKey::PanelThirtyDayTokens, "30日間のトークン"), (LocaleKey::PanelLatestTokens, "最新トークン"), (LocaleKey::PanelTopModelPrefix, "トップモデル"), (LocaleKey::PanelUsedSuffix, "使用済み"), @@ -184,8 +182,6 @@ fn test_chinese_tray_panel_locale_values_are_translated() { (LocaleKey::PanelZoom, "缩放"), (LocaleKey::PanelMenu, "菜单"), (LocaleKey::PanelToday, "今日"), - (LocaleKey::PanelThirtyDayCost, "30天成本"), - (LocaleKey::PanelThirtyDayTokens, "30天令牌"), (LocaleKey::PanelLatestTokens, "最新令牌"), (LocaleKey::PanelTopModelPrefix, "热门模型"), (LocaleKey::PanelUsedSuffix, "已使用"), @@ -220,8 +216,6 @@ fn test_korean_tray_panel_locale_values_are_translated() { (LocaleKey::PanelZoom, "확대/축소"), (LocaleKey::PanelMenu, "메뉴"), (LocaleKey::PanelToday, "오늘"), - (LocaleKey::PanelThirtyDayCost, "30일 비용"), - (LocaleKey::PanelThirtyDayTokens, "30일 토큰"), (LocaleKey::PanelLatestTokens, "최신 토큰"), (LocaleKey::PanelTopModelPrefix, "상위 모델"), (LocaleKey::PanelUsedSuffix, "사용됨"), @@ -256,8 +250,6 @@ fn test_spanish_tray_panel_locale_values_are_translated() { (LocaleKey::PanelZoom, "Zoom"), (LocaleKey::PanelMenu, "Menú"), (LocaleKey::PanelToday, "Hoy"), - (LocaleKey::PanelThirtyDayCost, "Costo 30d"), - (LocaleKey::PanelThirtyDayTokens, "Tokens 30d"), (LocaleKey::PanelLatestTokens, "Últimos tokens"), (LocaleKey::PanelTopModelPrefix, "Modelo principal"), (LocaleKey::PanelUsedSuffix, "usado"), diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index 08de23e2e3..b3233c057d 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -535,8 +535,6 @@ PanelZoom = Yakınlaştırma PanelMenu = Menü PanelCopied = Kopyalandı PanelToday = Bugün -PanelThirtyDayCost = 30 günlük maliyet -PanelThirtyDayTokens = 30 günlük tokenlar PanelLatestTokens = Son tokenlar PanelThirtyDayCostHistogram = 30 günlük maliyet histogramı PanelTopModelPrefix = En çok kullanılan model @@ -556,7 +554,6 @@ PanelReserveSuffix = yedekte PanelReserveLastsUntilReset = Sıfırlamaya kadar sürer PanelReserveRunsOutInDaysHours = { "{}" } g { "{}" } sa içinde tükenir PanelReserveRunsOutInHours = { "{}" } sa içinde tükenir -FloatBarThirtyDayShort = 30 gün FloatBarNoProviders = Sağlayıcı yok FloatBarRemainingSuffix = kaldı BannerCheckingForUpdates = Güncellemeler denetleniyor… diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index daf8e3c9af..d406be1559 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -515,8 +515,6 @@ PanelZoom = 缩放 PanelMenu = 菜单 PanelCopied = 已复制 PanelToday = 今日 -PanelThirtyDayCost = 30天成本 -PanelThirtyDayTokens = 30天令牌 PanelLatestTokens = 最新令牌 PanelThirtyDayCostHistogram = 30天成本直方图 PanelTopModelPrefix = 热门模型 @@ -536,7 +534,6 @@ PanelReserveSuffix = 储备 PanelReserveLastsUntilReset = 持续到重置 PanelReserveRunsOutInDaysHours = { "{}" }天{ "{}" }小时后耗尽 PanelReserveRunsOutInHours = { "{}" }小时后耗尽 -FloatBarThirtyDayShort = 30天 FloatBarNoProviders = 无提供商 FloatBarRemainingSuffix = 剩余 BannerCheckingForUpdates = 正在检查更新… diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 31bb3bcf02..99a3588e27 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -515,8 +515,6 @@ PanelZoom = 縮放 PanelMenu = 選單 PanelCopied = 已複製 PanelToday = 今日 -PanelThirtyDayCost = 30天成本 -PanelThirtyDayTokens = 30天權杖 PanelLatestTokens = 最新權杖 PanelThirtyDayCostHistogram = 30天成本直方圖 PanelTopModelPrefix = 熱門模型 @@ -536,7 +534,6 @@ PanelReserveSuffix = 儲備 PanelReserveLastsUntilReset = 持續到重置 PanelReserveRunsOutInDaysHours = { "{}" }天{ "{}" }小時後耗盡 PanelReserveRunsOutInHours = { "{}" }小時後耗盡 -FloatBarThirtyDayShort = 30天 FloatBarNoProviders = 無提供商 FloatBarRemainingSuffix = 剩餘 BannerCheckingForUpdates = 正在檢查更新… diff --git a/rust/src/providers/bedrock/daily.rs b/rust/src/providers/bedrock/daily.rs index 56b388fe41..cfb975eeee 100644 --- a/rust/src/providers/bedrock/daily.rs +++ b/rust/src/providers/bedrock/daily.rs @@ -30,7 +30,7 @@ pub(super) fn daily_range(since: NaiveDate, now: DateTime) -> (String, Stri fn utc_month_start(now: DateTime) -> NaiveDate { CostReportingPeriod::MonthToDate - .bounds(now, CostTimeZone::Named(chrono_tz::UTC), None) + .bounds(now, CostTimeZone::UTC, None) .start } From 5d993548c1411e60d2579a167d02e9ffec24b52f Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:02:15 +0700 Subject: [PATCH 128/259] Address thermo review --- rust/src/core/usage_snapshot.rs | 24 +++- rust/src/providers/openaiapi/history.rs | 109 +++++++++++++----- rust/src/providers/openaiapi/history_tests.rs | 18 +++ rust/src/providers/openaiapi/mod.rs | 8 +- 4 files changed, 124 insertions(+), 35 deletions(-) diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index e0eb197a4d..a8db374133 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -666,8 +666,9 @@ pub struct ProviderFetchResult { pub wayfinder_usage: Option, /// Per-day OpenAI Admin API history for the daily usage chart. Set only by the - /// Admin usage path; the balance fallback has no per-day data. - #[serde(default, skip_serializing_if = "Option::is_none")] + /// Admin usage path; the balance fallback has no per-day data. It is projected + /// through the frontend bridge and intentionally excluded from core serialization. + #[serde(skip)] pub open_ai_api_usage: Option, /// Transient non-quota inventory for provider-specific display. @@ -803,6 +804,25 @@ mod tests { assert!(decoded.inventory.is_empty()); } + #[test] + fn openai_api_history_is_transient_and_not_serialized() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let result = ProviderFetchResult::new(usage, "admin-api").with_open_ai_api_usage( + OpenAiApiUsageHistory { + history_days: 30, + project_id: Some("proj_abc".to_string()), + daily: Vec::new(), + }, + ); + + assert!(result.open_ai_api_usage.is_some()); + let encoded = serde_json::to_value(&result).unwrap(); + assert!(encoded.get("open_ai_api_usage").is_none()); + + let decoded: ProviderFetchResult = serde_json::from_value(encoded).unwrap(); + assert!(decoded.open_ai_api_usage.is_none()); + } + #[test] fn display_details_reject_invalid_shapes_and_duplicate_ids() { let usage = UsageSnapshot::new(RateWindow::new(25.0)); diff --git a/rust/src/providers/openaiapi/history.rs b/rust/src/providers/openaiapi/history.rs index d09dc9358b..b6d28a5910 100644 --- a/rust/src/providers/openaiapi/history.rs +++ b/rust/src/providers/openaiapi/history.rs @@ -49,7 +49,7 @@ impl DayAccumulator { } } - fn finish(self) -> OpenAiApiDailyUsage { + fn finish(self) -> Result { let mut line_items: Vec<_> = self .lines .into_iter() @@ -63,37 +63,46 @@ impl DayAccumulator { let mut models: Vec<_> = self .models .into_iter() - .map(|(name, totals)| OpenAiApiModelUsage { - name, - requests: totals.requests, - input_tokens: totals.input, - cached_input_tokens: totals.cached, - output_tokens: totals.output, - total_tokens: totals.input.saturating_add(totals.output), + .map(|(name, totals)| { + Ok(OpenAiApiModelUsage { + name, + requests: totals.requests, + input_tokens: totals.input, + cached_input_tokens: totals.cached, + output_tokens: totals.output, + total_tokens: checked_count_sum( + totals.input, + totals.output, + "model total_tokens", + )?, + }) }) - .collect(); + .collect::, ProviderError>>()?; models.sort_by(|a, b| { b.total_tokens .cmp(&a.total_tokens) .then_with(|| a.name.cmp(&b.name)) }); - let sum = |field: fn(&OpenAiApiModelUsage) -> u64| { - models - .iter() - .fold(0_u64, |sum, model| sum.saturating_add(field(model))) + let sum = |field: fn(&OpenAiApiModelUsage) -> u64, name: &str| { + models.iter().try_fold(0_u64, |sum, model| { + checked_count_sum(sum, field(model), name) + }) }; - OpenAiApiDailyUsage { + Ok(OpenAiApiDailyUsage { start_time: self.start, end_time: self.end, cost_usd: self.cost, - requests: sum(|model| model.requests), - input_tokens: sum(|model| model.input_tokens), - cached_input_tokens: sum(|model| model.cached_input_tokens), - output_tokens: sum(|model| model.output_tokens), - total_tokens: sum(|model| model.total_tokens), + requests: sum(|model| model.requests, "daily requests")?, + input_tokens: sum(|model| model.input_tokens, "daily input_tokens")?, + cached_input_tokens: sum( + |model| model.cached_input_tokens, + "daily cached_input_tokens", + )?, + output_tokens: sum(|model| model.output_tokens, "daily output_tokens")?, + total_tokens: sum(|model| model.total_tokens, "daily total_tokens")?, line_items, models, - } + }) } } @@ -135,14 +144,18 @@ pub(super) fn daily_usage( .models .entry(display_name(result.model.as_deref(), DEFAULT_MODEL)) .or_default(); - model.requests = model.requests.saturating_add(requests); - model.input = model - .input - .saturating_add(input.saturating_add(audio_input)); - model.cached = model.cached.saturating_add(cached); - model.output = model - .output - .saturating_add(output.saturating_add(audio_output)); + model.requests = checked_count_sum(model.requests, requests, "model requests")?; + model.input = checked_count_sum( + model.input, + checked_count_sum(input, audio_input, "input_tokens")?, + "model input_tokens", + )?; + model.cached = checked_count_sum(model.cached, cached, "model cached_input_tokens")?; + model.output = checked_count_sum( + model.output, + checked_count_sum(output, audio_output, "output_tokens")?, + "model output_tokens", + )?; } } @@ -151,7 +164,7 @@ pub(super) fn daily_usage( .into_values() .filter(|day| day.start <= now) .map(DayAccumulator::finish) - .collect(); + .collect::>()?; let excess = daily.len().saturating_sub(history_days as usize); daily.drain(..excess); Ok(daily) @@ -170,6 +183,12 @@ pub(super) fn usage_history( tracing::warn!("Dropping OpenAI API daily history: a bucket does not end after it starts"); return None; } + if daily.iter().any(|day| !counts_fit_js_number(day)) { + tracing::warn!( + "Dropping OpenAI API daily history: an aggregate count exceeds the JavaScript safe-integer range" + ); + return None; + } let entries: usize = daily .iter() .map(|day| day.line_items.len() + day.models.len()) @@ -185,6 +204,38 @@ pub(super) fn usage_history( }) } +fn counts_fit_js_number(day: &OpenAiApiDailyUsage) -> bool { + let safe = |count| count <= MAX_SAFE_COUNT; + [ + day.requests, + day.input_tokens, + day.cached_input_tokens, + day.output_tokens, + day.total_tokens, + ] + .into_iter() + .all(safe) + && day.models.iter().all(|model| { + [ + model.requests, + model.input_tokens, + model.cached_input_tokens, + model.output_tokens, + model.total_tokens, + ] + .into_iter() + .all(safe) + }) +} + +fn checked_count_sum(left: u64, right: u64, field: &str) -> Result { + left.checked_add(right).ok_or_else(|| { + ProviderError::Parse(format!( + "OpenAI API completions {field} total exceeds the supported integer range" + )) + }) +} + /// Upstream `name()`: a trimmed non-empty string, else the fallback. fn display_name(raw: Option<&str>, fallback: &str) -> String { raw.map(str::trim) diff --git a/rust/src/providers/openaiapi/history_tests.rs b/rust/src/providers/openaiapi/history_tests.rs index b17a895f7a..d037d252c7 100644 --- a/rust/src/providers/openaiapi/history_tests.rs +++ b/rust/src/providers/openaiapi/history_tests.rs @@ -277,6 +277,24 @@ fn history_is_dropped_but_the_summary_kept_for_a_bucket_that_does_not_end_after_ assert_eq!(result.cost.unwrap().used, 3.0); } +#[test] +fn history_is_dropped_when_aggregated_counts_exceed_javascript_integer_precision() { + let max_safe_count = 9_007_199_254_740_991; + let costs = [cost_bucket_at(NOW - DAY, 3.0, None)]; + let completions = [completions_bucket( + NOW - DAY, + vec![ + completion(Some("gpt-5.2"), max_safe_count, 0, 0), + completion(Some("gpt-5.2"), max_safe_count, 0, 0), + ], + )]; + + let result = result_from_admin_usage(&costs, &completions, now(), None).unwrap(); + + assert!(result.open_ai_api_usage.is_none()); + assert_eq!(result.cost.unwrap().used, 3.0); +} + #[test] fn history_is_dropped_beyond_ten_thousand_breakdown_rows() { let models = |count: usize| -> Vec { diff --git a/rust/src/providers/openaiapi/mod.rs b/rust/src/providers/openaiapi/mod.rs index dda101c00b..4d73b4dada 100644 --- a/rust/src/providers/openaiapi/mod.rs +++ b/rust/src/providers/openaiapi/mod.rs @@ -484,13 +484,13 @@ fn result_from_admin_usage( let daily = history::daily_usage(costs, completions, now, HISTORY_DAYS)?; let cost_total: f64 = daily.iter().map(|day| day.cost_usd).sum(); - let request_total: u64 = daily.iter().map(|day| day.requests).sum(); - let token_total: u64 = daily.iter().map(|day| day.total_tokens).sum(); - let mut model_tokens: HashMap<&str, u64> = HashMap::new(); + let request_total: u128 = daily.iter().map(|day| u128::from(day.requests)).sum(); + let token_total: u128 = daily.iter().map(|day| u128::from(day.total_tokens)).sum(); + let mut model_tokens: HashMap<&str, u128> = HashMap::new(); let mut line_item_costs: HashMap<&str, f64> = HashMap::new(); for day in &daily { for model in &day.models { - *model_tokens.entry(&model.name).or_default() += model.total_tokens; + *model_tokens.entry(&model.name).or_default() += u128::from(model.total_tokens); } for item in &day.line_items { *line_item_costs.entry(&item.name).or_default() += item.cost_usd; From 813458d0f9a99b359051794fd3fbf0ad6eb7275b Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:06:35 +0700 Subject: [PATCH 129/259] Address thermo review --- .../desktop-tauri/src/components/MenuCard.tsx | 13 +++-- .../components/MenuCardOpenAIUsage.test.tsx | 17 ++++++- .../components/OpenAIApiUsageChart.test.tsx | 14 ++++-- .../src/components/OpenAIApiUsageChart.tsx | 24 ++++----- .../src/components/charts/BarChart.tsx | 50 ++++++++++++------- .../components/charts/chartGeometry.test.ts | 11 ++++ .../src/components/charts/chartGeometry.ts | 26 ++++++---- .../charts/useChartAnimation.test.ts | 36 +++++++++++++ .../components/charts/useChartAnimation.ts | 11 ++-- .../src/lib/openAiApiIdentity.test.ts | 19 +++++++ .../src/lib/openAiApiIdentity.ts | 8 +++ apps/desktop-tauri/src/styles.css | 15 ++++++ .../settings/providers/ProviderDetailPane.tsx | 7 ++- .../sections/IdentitySection.test.tsx | 44 ++++++++++++++++ .../providers/sections/IdentitySection.tsx | 25 ++++++++-- 15 files changed, 265 insertions(+), 55 deletions(-) create mode 100644 apps/desktop-tauri/src/components/charts/useChartAnimation.test.ts create mode 100644 apps/desktop-tauri/src/lib/openAiApiIdentity.test.ts create mode 100644 apps/desktop-tauri/src/lib/openAiApiIdentity.ts create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.test.tsx diff --git a/apps/desktop-tauri/src/components/MenuCard.tsx b/apps/desktop-tauri/src/components/MenuCard.tsx index 493ff08e98..383ac5331a 100644 --- a/apps/desktop-tauri/src/components/MenuCard.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.tsx @@ -9,6 +9,7 @@ import { useLocale } from "../hooks/useLocale"; import { formatRelativeUpdated } from "../lib/relativeTime"; import type { LocaleKey } from "../i18n/keys"; import { providerSupportsChartData } from "../lib/providerCharts"; +import { hideOpenAiApiProjectId } from "../lib/openAiApiIdentity"; import MenuCardDetails, { describeCard, type MetricEntry } from "./MenuCardDetails"; import CodexAccountsMenu from "./CodexAccountsMenu"; import ClaudeAccountsMenu from "./ClaudeAccountsMenu"; @@ -194,7 +195,13 @@ export default function MenuCard({ ? maskEmail(provider.accountEmail) : provider.accountEmail : null; - const planName = !isWayfinder ? displayPlanName(provider.planName, t) : null; + const planName = !isWayfinder + ? displayPlanName(provider.planName, t) + : null; + const displayedPlanName = hideOpenAiApiProjectId( + planName, + provider.providerId === "openaiapi" && hideEmail, + ); const metrics: MetricEntry[] = [ ...(isWayfinder @@ -285,8 +292,8 @@ export default function MenuCard({ ? provider.updatedAt : formatRelativeUpdated(Date.parse(provider.updatedAt), t)} - {planName && ( - {planName} + {displayedPlanName && ( + {displayedPlanName} )}
)} diff --git a/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx b/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx index 4dc02025fc..fa72167b1c 100644 --- a/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCardOpenAIUsage.test.tsx @@ -91,12 +91,16 @@ function snapshot( }; } -function renderCard(snap: ProviderUsageSnapshot, compactOverview = false) { +function renderCard( + snap: ProviderUsageSnapshot, + compactOverview = false, + hidePersonalInfo = false, +) { return render( , ); @@ -139,6 +143,15 @@ describe("MenuCard OpenAI daily usage section", () => { expect(container.querySelector('[role="option"]')).toBeNull(); }); + it("masks the Admin project id in the card identity when privacy is enabled", async () => { + const provider = snapshot("openaiapi", usage()); + provider.planName = "Admin API: proj-private"; + const { container } = renderCard(provider, false, true); + await waitFor(() => expect(container.querySelector(".menu-card__plan-badge")).not.toBeNull()); + expect(container.querySelector(".menu-card__plan-badge")).toHaveTextContent("Admin API: ••••"); + expect(container).not.toHaveTextContent("proj-private"); + }); + it("hides the section on error, in compact overview and with an empty window", async () => { const errored = renderCard(snapshot("openaiapi", usage(), "boom")); await waitFor(() => expect(tauriMocks.getLocaleStrings).toHaveBeenCalled()); diff --git a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx index de38174925..f3d157fecf 100644 --- a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx +++ b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.test.tsx @@ -203,11 +203,19 @@ describe("OpenAIApiUsageChart", () => { expect(within(detail).queryByText("OpenAIChartModels")).toBeNull(); }); - it("draws only the latest 60 days and relabels the window to match", () => { + it("keeps the full history in a horizontally scrollable chart", () => { const days = Array.from({ length: 90 }, (_, i) => day(i)); renderChart(usage(days, 90)); - expect(bars()).toHaveLength(60); - expect(screen.getByText("Last 60 days")).toBeInTheDocument(); + expect(bars()).toHaveLength(90); + expect(screen.getByText("Last 90 days")).toBeInTheDocument(); + const chart = document.querySelector(".chart--scrollable"); + expect(chart).toBeInTheDocument(); + expect(Number(chart?.querySelector("svg")?.getAttribute("width"))).toBeGreaterThan(280); + expect(bars()[0]).toHaveAttribute("aria-label", "2026-09-01: $1.00"); + expect(bars()[89]).toHaveAttribute("aria-label", "2026-11-29: $90.00"); + expect(screen.getByText("Last 90 days").nextElementSibling).toHaveTextContent( + "$4,095.00", + ); }); it("shows the empty message when the window has no days", () => { diff --git a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx index 9177f10e5d..a24c640c0a 100644 --- a/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx +++ b/apps/desktop-tauri/src/components/OpenAIApiUsageChart.tsx @@ -17,14 +17,9 @@ type Metric = "cost" | "tokens"; type T = (key: LocaleKey) => string; const METRICS: readonly Metric[] = ["cost", "tokens"]; -/** Bars stay at least 1px wide inside the shared 280px chart geometry. */ -const MAX_CHART_DAYS = 60; const MAX_DETAIL_ROWS = 5; -const usdFormat = new Intl.NumberFormat("en-US", { - style: "currency", - currency: "USD", -}); +const usdFormat = new Intl.NumberFormat("en-US", { style: "currency", currency: "USD" }); const compactFormat = new Intl.NumberFormat("en-US", { notation: "compact", maximumFractionDigits: 1, @@ -107,7 +102,9 @@ function DayDetail({ day, t }: { day: OpenAiApiDailyUsageSnapshot; t: T }) { const models = day.models.map((model, index) => ({ key: `${model.name}-${index}`, name: model.name, - value: `${compactFormat.format(model.totalTokens)} · ${countFormat.format(model.requests)} ${t("UsageSpendRequests")}`, + value: `${compactFormat.format(model.totalTokens)} · ${countFormat.format(model.requests)} ${t( + "UsageSpendRequests", + )}`, })); return (
@@ -145,9 +142,13 @@ export function OpenAIApiUsageChart({ const [selectedStart, setSelectedStart] = useState(null); const baseId = useId(); const tabRefs = useRef>([]); - const days = useMemo(() => usage.daily.slice(-MAX_CHART_DAYS), [usage.daily]); + const days = usage.daily; const points = useMemo( - () => days.map((day) => ({ label: dayLabel(day), value: metricValue(day, metric) })), + () => + days.map((day) => ({ + label: dayLabel(day), + value: metricValue(day, metric), + })), [days, metric], ); @@ -170,8 +171,6 @@ export function OpenAIApiUsageChart({ const found = days.findIndex((day) => day.startTime === selectedStart); const selectedIndex = found >= 0 ? found : days.length - 1; const total = days.reduce((sum, day) => sum + metricValue(day, metric), 0); - const trimmed = days.length < usage.daily.length; - const windowDays = trimmed ? days.length : usage.historyDays; const tabLabel = (m: Metric) => t(m === "cost" ? "OpenAIChartMetricCost" : "OpenAIChartMetricTokens"); const tabId = (m: Metric) => `${baseId}-tab-${m}`; @@ -222,6 +221,7 @@ export function OpenAIApiUsageChart({ valueFormatter={(v) => formatMetric(v, metric)} animations={animations} emptyMessage={t("DetailChartEmpty")} + scrollable selection={{ index: selectedIndex, onSelect: (index) => setSelectedStart(days[index]?.startTime ?? null), @@ -230,7 +230,7 @@ export function OpenAIApiUsageChart({
{t("OpenAIChartHint")}
- {windowLabel(windowDays, t)} + {windowLabel(usage.historyDays, t)} {formatMetric(total, metric)}
diff --git a/apps/desktop-tauri/src/components/charts/BarChart.tsx b/apps/desktop-tauri/src/components/charts/BarChart.tsx index 65608c8395..a47c083c41 100644 --- a/apps/desktop-tauri/src/components/charts/BarChart.tsx +++ b/apps/desktop-tauri/src/components/charts/BarChart.tsx @@ -5,6 +5,7 @@ import { getBarCenter, getBarWidth, getBarX, + getScrollableChartWidth, shouldRenderCenterMax, } from "./chartGeometry"; @@ -46,6 +47,8 @@ export interface BarChartProps { /** Optional empty-state message rendered when `data.length === 0`. */ emptyMessage?: string; selection?: BarChartSelection; + /** Keep narrow bars readable by making long series horizontally scrollable. */ + scrollable?: boolean; } const DEFAULT_COLOR = "var(--chart-cost)"; @@ -60,6 +63,7 @@ export function BarChart({ animations = true, emptyMessage, selection, + scrollable = false, }: BarChartProps) { const fmt = valueFormatter ?? ((v: number) => v.toFixed(2)); const containerRef = useRef(null); @@ -93,7 +97,11 @@ export function BarChart({ ); } - const barWidth = getBarWidth(data.length); + const chartWidth = scrollable ? getScrollableChartWidth(data.length) : WIDTH; + const chartClass = scrollable + ? "chart chart--bar chart--scrollable" + : "chart chart--bar"; + const barWidth = getBarWidth(data.length, chartWidth); const plotHeight = Math.max(1, height - 4); const onMove = (e: React.MouseEvent, i: number) => { @@ -122,12 +130,13 @@ export function BarChart({ }; return ( -
+
CAP_HEIGHT; const bodyH = isPeak ? Math.max(0, barH - CAP_HEIGHT) : barH; const bodyY = isPeak ? y + CAP_HEIGHT : y; const isHovered = hover?.i === i; const isSelected = i === selectedIndex; + let opacity = 0.9; + if (p.value == null) opacity = 0; + else if (p.value === 0) opacity = 0.25; + else if (selection) opacity = isSelected ? 1 : 0.6; + else if (isHovered) opacity = 1; return ( @@ -153,15 +167,7 @@ export function BarChart({ width={barWidth} height={bodyH} fill={color} - opacity={ - p.value == null - ? 0 - : p.value === 0 - ? 0.25 - : selection - ? isSelected ? 1 : 0.6 - : isHovered ? 1 : 0.9 - } + opacity={opacity} rx={1} className="chart__bar" {...(selection @@ -205,14 +211,22 @@ export function BarChart({ ); })} -
- +
+ {data[0].label} {shouldRenderCenterMax(data.length) && ( - {fmt(max)} + + {fmt(max)} + )} - + {data[data.length - 1].label}
diff --git a/apps/desktop-tauri/src/components/charts/chartGeometry.test.ts b/apps/desktop-tauri/src/components/charts/chartGeometry.test.ts index b4c867b7cd..c55d2ae9ac 100644 --- a/apps/desktop-tauri/src/components/charts/chartGeometry.test.ts +++ b/apps/desktop-tauri/src/components/charts/chartGeometry.test.ts @@ -3,6 +3,8 @@ import { DATE_EDGE_PADDING, WIDTH, getBarCenter, + getBarWidth, + getScrollableChartWidth, shouldRenderCenterMax, } from "./chartGeometry"; @@ -22,4 +24,13 @@ describe("chart axis geometry", () => { it("renders the center max label for three or more points", () => { expect(shouldRenderCenterMax(3)).toBe(true); }); + + it("allocates enough scrollable width to keep long-series bars visible", () => { + const width = getScrollableChartWidth(90); + const lastBarRight = getBarCenter(89, 90, width) + getBarWidth(90, width) / 2; + + expect(width).toBeGreaterThan(WIDTH); + expect(getBarWidth(90, width)).toBeCloseTo(1.5); + expect(lastBarRight).toBeLessThanOrEqual(width - DATE_EDGE_PADDING); + }); }); diff --git a/apps/desktop-tauri/src/components/charts/chartGeometry.ts b/apps/desktop-tauri/src/components/charts/chartGeometry.ts index f3de6d951d..b428908d56 100644 --- a/apps/desktop-tauri/src/components/charts/chartGeometry.ts +++ b/apps/desktop-tauri/src/components/charts/chartGeometry.ts @@ -3,22 +3,30 @@ export const DATE_EDGE_PADDING = 36; export const BAR_GAP = 2; export const PLOT_WIDTH = WIDTH - DATE_EDGE_PADDING * 2; export const AXIS_MAX_X = WIDTH / 2; +const SCROLLABLE_BAR_WIDTH = 1.5; export function shouldRenderCenterMax(count: number): boolean { return count >= 3; } -export function getBarWidth(count: number): number { +export function getScrollableChartWidth(count: number): number { const barCount = Math.max(1, count); - return Math.max(1, (PLOT_WIDTH - (barCount - 1) * BAR_GAP) / barCount); + const plotWidth = barCount * SCROLLABLE_BAR_WIDTH + (barCount - 1) * BAR_GAP; + return Math.max(WIDTH, DATE_EDGE_PADDING * 2 + plotWidth); } -export function getBarX(index: number, count: number): number { - return DATE_EDGE_PADDING + index * (getBarWidth(count) + BAR_GAP); +export function getBarWidth(count: number, width = WIDTH): number { + const barCount = Math.max(1, count); + const plotWidth = width - DATE_EDGE_PADDING * 2; + return Math.max(1, (plotWidth - (barCount - 1) * BAR_GAP) / barCount); +} + +export function getBarX(index: number, count: number, width = WIDTH): number { + return DATE_EDGE_PADDING + index * (getBarWidth(count, width) + BAR_GAP); } -export function getBarCenter(index: number, count: number): number { - return getBarX(index, count) + getBarWidth(count) / 2; +export function getBarCenter(index: number, count: number, width = WIDTH): number { + return getBarX(index, count, width) + getBarWidth(count, width) / 2; } export interface BarGeometry { @@ -28,9 +36,9 @@ export interface BarGeometry { center: (index: number) => number; } -export function getBarGeometry(count: number): BarGeometry { - const barWidth = getBarWidth(count); - const x = (index: number) => getBarX(index, count); +export function getBarGeometry(count: number, width = WIDTH): BarGeometry { + const barWidth = getBarWidth(count, width); + const x = (index: number) => getBarX(index, count, width); return { barWidth, diff --git a/apps/desktop-tauri/src/components/charts/useChartAnimation.test.ts b/apps/desktop-tauri/src/components/charts/useChartAnimation.test.ts new file mode 100644 index 0000000000..2eeabef1f1 --- /dev/null +++ b/apps/desktop-tauri/src/components/charts/useChartAnimation.test.ts @@ -0,0 +1,36 @@ +import { act, renderHook } from "@testing-library/react"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { useChartAnimation } from "./useChartAnimation"; + +describe("useChartAnimation long series", () => { + afterEach(() => vi.unstubAllGlobals()); + + it("finishes long-series entrance animation after the bounded stagger", () => { + const frames = new Map(); + let nextFrame = 0; + vi.stubGlobal("requestAnimationFrame", (callback: FrameRequestCallback) => { + const id = ++nextFrame; + frames.set(id, callback); + return id; + }); + vi.stubGlobal("cancelAnimationFrame", (id: number) => frames.delete(id)); + + const { result, unmount } = renderHook(() => useChartAnimation(366, true)); + act(() => { + const first = frames.get(1); + if (!first) throw new Error("expected initial animation frame"); + frames.delete(1); + first(0); + }); + act(() => { + const next = frames.get(2); + if (!next) throw new Error("expected follow-up animation frame"); + frames.delete(2); + next(1_200); + }); + + expect(result.current.running).toBe(false); + expect(result.current.barProgress(365)).toBe(1); + unmount(); + }); +}); diff --git a/apps/desktop-tauri/src/components/charts/useChartAnimation.ts b/apps/desktop-tauri/src/components/charts/useChartAnimation.ts index f1e94ccecf..94417ffac9 100644 --- a/apps/desktop-tauri/src/components/charts/useChartAnimation.ts +++ b/apps/desktop-tauri/src/components/charts/useChartAnimation.ts @@ -19,6 +19,7 @@ import { useEffect, useRef, useState } from "react"; */ export const TOTAL_ANIMATION_MS = 600; export const STAGGER_PER_BAR_MS = 20; +const MAX_STAGGERED_ITEMS = 30; const easeOut = (t: number) => 1 - Math.pow(1 - t, 3); const clamp01 = (v: number) => (v < 0 ? 0 : v > 1 ? 1 : v); @@ -42,6 +43,9 @@ export function useChartAnimation( const rafRef = useRef(null); const prefersReduced = usePrefersReducedMotion(); const skip = !enabled || prefersReduced || count === 0; + // Keep long daily histories animated without making the entrance delay grow + // linearly with every bucket. Items after the first 30 start together. + const staggeredCount = Math.min(count, MAX_STAGGERED_ITEMS); useEffect(() => { if (skip) { @@ -52,7 +56,7 @@ export function useChartAnimation( startRef.current = null; setElapsed(0); - const totalMs = TOTAL_ANIMATION_MS + count * STAGGER_PER_BAR_MS; + const totalMs = TOTAL_ANIMATION_MS + staggeredCount * STAGGER_PER_BAR_MS; const tick = (now: number) => { if (startRef.current == null) startRef.current = now; @@ -74,10 +78,11 @@ export function useChartAnimation( const progress = skip ? 1 : clamp01(elapsed / TOTAL_ANIMATION_MS); const barProgress = (i: number) => { if (skip) return 1; - const barElapsed = Math.max(0, elapsed - i * STAGGER_PER_BAR_MS); + const stagger = Math.min(i, MAX_STAGGERED_ITEMS) * STAGGER_PER_BAR_MS; + const barElapsed = Math.max(0, elapsed - stagger); return easeOut(clamp01(barElapsed / TOTAL_ANIMATION_MS)); }; - const totalMs = TOTAL_ANIMATION_MS + count * STAGGER_PER_BAR_MS; + const totalMs = TOTAL_ANIMATION_MS + staggeredCount * STAGGER_PER_BAR_MS; const running = !skip && elapsed < totalMs; return { progress: easeOut(progress), barProgress, running }; diff --git a/apps/desktop-tauri/src/lib/openAiApiIdentity.test.ts b/apps/desktop-tauri/src/lib/openAiApiIdentity.test.ts new file mode 100644 index 0000000000..6160b5b8d6 --- /dev/null +++ b/apps/desktop-tauri/src/lib/openAiApiIdentity.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; +import { hideOpenAiApiProjectId } from "./openAiApiIdentity"; + +describe("hideOpenAiApiProjectId", () => { + it("masks only the project id in OpenAI identity labels", () => { + expect(hideOpenAiApiProjectId("Admin API: proj-private", true)).toBe("Admin API: ••••"); + expect(hideOpenAiApiProjectId("Project: proj-private", true)).toBe("Project: ••••"); + }); + + it("preserves visible identity text when privacy is off", () => { + expect(hideOpenAiApiProjectId("Admin API: proj-private", false)).toBe("Admin API: proj-private"); + }); + + it("does not alter unrelated labels", () => { + expect(hideOpenAiApiProjectId("Admin API", true)).toBe("Admin API"); + expect(hideOpenAiApiProjectId("gpt-5", true)).toBe("gpt-5"); + expect(hideOpenAiApiProjectId(null, true)).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/lib/openAiApiIdentity.ts b/apps/desktop-tauri/src/lib/openAiApiIdentity.ts new file mode 100644 index 0000000000..4e31c20c5b --- /dev/null +++ b/apps/desktop-tauri/src/lib/openAiApiIdentity.ts @@ -0,0 +1,8 @@ +/** Mask OpenAI Admin project ids in generic identity labels when privacy is on. */ +export function hideOpenAiApiProjectId( + value: string | null, + hidePersonalInfo: boolean, +): string | null { + if (!hidePersonalInfo || !value) return value; + return value.replace(/^(Admin API|Project):(\s*).+$/i, "$1:$2••••"); +} diff --git a/apps/desktop-tauri/src/styles.css b/apps/desktop-tauri/src/styles.css index 0da531f5c0..a558b15051 100644 --- a/apps/desktop-tauri/src/styles.css +++ b/apps/desktop-tauri/src/styles.css @@ -2704,6 +2704,21 @@ body:has(.tray-panel-reveal) { height: auto; } +.chart--scrollable { + overflow-x: auto; + overflow-y: hidden; +} + +.chart--scrollable .chart__svg { + flex: 0 0 auto; + width: auto; + max-width: none; +} + +.chart--scrollable .chart__axis { + flex: 0 0 auto; +} + .chart__axis { position: relative; font-size: 0.7rem; diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx index fc00d44dcd..12bfe8931a 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx @@ -288,7 +288,12 @@ export function ProviderDetailPane({ return (
- + {detail.id === "codex" && ( diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.test.tsx new file mode 100644 index 0000000000..5631cf95bb --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.test.tsx @@ -0,0 +1,44 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import type { LocaleKey } from "../../../../i18n/keys"; +import { IdentitySection } from "./IdentitySection"; + +type IdentityDetails = Parameters[0]["provider"]; + +function detail(overrides: Partial = {}): IdentityDetails { + return { + id: "openaiapi", + displayName: "OpenAI API", + organization: "Project: proj-private", + plan: "Admin API: proj-private", + email: null, + authType: null, + sourceLabel: null, + ...overrides, + }; +} + +const t = (key: LocaleKey) => key; + +describe("IdentitySection OpenAI project privacy", () => { + it("masks project ids in account and plan rows when privacy is enabled", () => { + render( + , + ); + + expect(screen.getByText("Project: ••••")).toBeInTheDocument(); + expect(screen.getByText("Admin API: ••••")).toBeInTheDocument(); + expect(screen.queryByText("proj-private")).toBeNull(); + }); + + it("keeps project ids visible when privacy is disabled", () => { + render(); + + expect(screen.getAllByText(/proj-private/)).toHaveLength(2); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.tsx index 624d5d1b07..7fb2b68dbe 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/IdentitySection.tsx @@ -1,11 +1,16 @@ import type { ProviderDetail } from "../../../../types/bridge"; import type { LocaleKey } from "../../../../i18n/keys"; import { ProviderIcon } from "../../../../components/providers/ProviderIcon"; +import { hideOpenAiApiProjectId } from "../../../../lib/openAiApiIdentity"; interface Props { - provider: ProviderDetail; + provider: Pick< + ProviderDetail, + "id" | "displayName" | "email" | "organization" | "plan" | "authType" | "sourceLabel" + >; subtitle: string; t: (key: LocaleKey) => string; + hidePersonalInfo?: boolean; } /** @@ -15,10 +20,22 @@ interface Props { * Port of the identity portion of * `rust/src/native_ui/preferences.rs::render_provider_detail_panel` (~4301). */ -export function IdentitySection({ provider, subtitle, t }: Props) { +export function IdentitySection({ + provider, + subtitle, + t, + hidePersonalInfo = false, +}: Props) { + const hideProjectId = provider.id === "openaiapi" && hidePersonalInfo; const rows: { label: string; value: string | null }[] = [ - { label: t("Account"), value: provider.email ?? provider.organization }, - { label: t("Plan"), value: displayIdentityValue(provider.plan, t) }, + { + label: t("Account"), + value: hideOpenAiApiProjectId(provider.email ?? provider.organization, hideProjectId), + }, + { + label: t("Plan"), + value: hideOpenAiApiProjectId(displayIdentityValue(provider.plan, t), hideProjectId), + }, { label: t("AuthType"), value: provider.authType }, { label: t("DataSource"), value: provider.sourceLabel }, ]; From d3e857c86cc0b3bb0b9964f67a31283674069508 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:13:18 +0700 Subject: [PATCH 130/259] Address thermo review --- rust/src/codex_costs/quota_windows.rs | 14 +-- rust/src/core/jsonl_scanner.rs | 3 +- rust/src/core/jsonl_scanner/codex/priority.rs | 19 ++-- .../core/jsonl_scanner/codex/source_rows.rs | 23 ++++ rust/src/cost_scanner/codex/cache_days.rs | 48 ++------- rust/src/cost_scanner/codex/priority_trace.rs | 102 ++++++++++++------ .../codex/priority_trace/tests.rs | 87 +++++++++++---- 7 files changed, 179 insertions(+), 117 deletions(-) diff --git a/rust/src/codex_costs/quota_windows.rs b/rust/src/codex_costs/quota_windows.rs index 46ff13a083..96286d35bd 100644 --- a/rust/src/codex_costs/quota_windows.rs +++ b/rust/src/codex_costs/quota_windows.rs @@ -11,7 +11,7 @@ use std::path::Path; use crate::core::{ CodexPriorityOverlay, CodexSourceRowCache, CodexSourceUsageRow, CostUsageCache, - CostUsagePricing, RateWindow, + CostUsagePricing, RateWindow, priced_model, }; const NOMINAL_WEEK_MINUTES: i64 = 7 * 24 * 60; @@ -370,17 +370,7 @@ fn slice_from_row(row: &CodexSourceUsageRow, overlay: Option<&CodexPriorityOverl let input = u64::try_from(row.input.max(0)).unwrap_or(0); let output = u64::try_from(row.output.max(0)).unwrap_or(0); let tokens = Some(input.saturating_add(output)); - let cost_usd = row.pricing.pricing_model.as_deref().and_then(|model| { - let model = if row.pricing.pricing_mode.as_deref() == Some("priority") - && !model.ends_with("-priority") - { - format!("{model}-priority") - } else { - model.to_string() - }; - let model = overlay - .and_then(|overlay| overlay.priority_model(row.turn_id.as_deref(), &model)) - .unwrap_or(model); + let cost_usd = priced_model(&row.pricing, row.turn_id.as_deref(), overlay).and_then(|model| { let date = timestamp.map(|value| value.with_timezone(&Local).date_naive())?; let cached = u64::try_from(row.cached.max(0)).unwrap_or(0).min(input); if model.ends_with("-priority") { diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 31ad34abfa..62142dee19 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -522,7 +522,8 @@ pub use codex::priority::{ CodexPriorityTurnMetadata, CodexPriorityTurnsCursor, }; pub(crate) use codex::source_rows::{ - read_source_rows, recover_rows, row_cache, row_cache_matches, row_cache_needs_recovery, + priced_model, read_source_rows, recover_rows, row_cache, row_cache_matches, + row_cache_needs_recovery, }; impl JsonlScanner { diff --git a/rust/src/core/jsonl_scanner/codex/priority.rs b/rust/src/core/jsonl_scanner/codex/priority.rs index 12fdcf0d87..2b865a15c8 100644 --- a/rust/src/core/jsonl_scanner/codex/priority.rs +++ b/rust/src/core/jsonl_scanner/codex/priority.rs @@ -11,7 +11,7 @@ //! baked into the persisted source-row pricing evidence, so cached-price //! recovery keeps comparing like with like. -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::path::Path; use serde::{Deserialize, Serialize}; @@ -54,9 +54,6 @@ pub struct CodexPriorityTurnsCursor { pub file_identity: Option, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub anchors: Vec, - /// Priority turns keyed by turn id. - #[serde(default)] - pub turns: HashMap, /// Source trace rows behind each Priority turn, keyed by rowid. #[serde(default)] pub request_sources: HashMap>, @@ -67,7 +64,7 @@ pub struct CodexPriorityTurnsCursor { #[serde(default)] pub completed_models: HashMap>, #[serde(default)] - pub completed_order: Vec, + pub completed_order: VecDeque, } /// A borrowed view that prices session rows of one file as Priority. @@ -80,13 +77,18 @@ impl CodexPriorityTurnsCursor { /// the Codex home this trace database belongs to. Turn sets are never /// shared across `CODEX_HOME` scopes. pub(crate) fn overlay_for_file(&self, file_path: &Path) -> Option> { - if self.turns.is_empty() { + if self.request_sources.is_empty() { return None; } let home = Path::new(&self.database_path).parent()?; path_is_under(file_path, home).then_some(CodexPriorityOverlay { cursor: self }) } + /// Latest retained Priority request for `turn_id`. + pub(crate) fn turn(&self, turn_id: &str) -> Option<&CodexPriorityTurnMetadata> { + self.request_sources.get(turn_id)?.values().next_back() + } + /// Model evidence for a Priority turn: the latest completion, else the /// request model. pub(crate) fn turn_model(&self, turn_id: &str) -> Option<&str> { @@ -94,7 +96,7 @@ impl CodexPriorityTurnsCursor { .get(turn_id) .and_then(|models| models.values().next_back()) .map(String::as_str) - .or_else(|| self.turns.get(turn_id)?.model.as_deref()) + .or_else(|| self.turn(turn_id)?.model.as_deref()) } } @@ -104,7 +106,8 @@ impl CodexPriorityOverlay<'_> { /// lane keeps Standard pricing, as upstream does. pub(crate) fn priority_model(&self, turn_id: Option<&str>, row_model: &str) -> Option { let turn_id = turn_id?; - if !self.cursor.turns.contains_key(turn_id) || row_model.ends_with("-priority") { + self.cursor.turn(turn_id)?; + if row_model.ends_with("-priority") { return None; } let priced = self diff --git a/rust/src/core/jsonl_scanner/codex/source_rows.rs b/rust/src/core/jsonl_scanner/codex/source_rows.rs index 8f888aa1c0..a09ab9f615 100644 --- a/rust/src/core/jsonl_scanner/codex/source_rows.rs +++ b/rust/src/core/jsonl_scanner/codex/source_rows.rs @@ -51,6 +51,29 @@ pub(crate) fn model_of_pricing_mode(model: &str) -> String { } } +/// Apply the row's stored pricing mode and optional trace evidence. +pub(crate) fn priced_model( + pricing: &CodexSourcePricingEvidence, + turn_id: Option<&str>, + overlay: Option<&CodexPriorityOverlay<'_>>, +) -> Option { + let model = pricing + .pricing_model + .as_deref() + .filter(|model| !model.is_empty())?; + let model = + if pricing.pricing_mode.as_deref() == Some("priority") && !model.ends_with("-priority") { + format!("{model}-priority") + } else { + model.to_string() + }; + Some( + overlay + .and_then(|overlay| overlay.priority_model(turn_id, &model)) + .unwrap_or(model), + ) +} + /// Re-read the bounded reporting partition to obtain request-row order. /// The normal scanner still owns aggregate parsing and its byte budget; /// this path is used only after a complete file pass has established that diff --git a/rust/src/cost_scanner/codex/cache_days.rs b/rust/src/cost_scanner/codex/cache_days.rs index d5f3be7ff8..db6f438ab2 100644 --- a/rust/src/cost_scanner/codex/cache_days.rs +++ b/rust/src/cost_scanner/codex/cache_days.rs @@ -1,5 +1,7 @@ use super::*; -use crate::core::{CodexPriorityOverlay, CodexSourceUsageRow, CodexUsageRecord, CostUsagePricing}; +use crate::core::{ + CodexPriorityOverlay, CodexSourceUsageRow, CodexUsageRecord, CostUsagePricing, priced_model, +}; use std::collections::BTreeMap; type DayModels = HashMap>>; @@ -76,20 +78,12 @@ fn priority_days( parsed_days: &DayModels, overlay: &CodexPriorityOverlay<'_>, ) -> Option { - if !rows - .iter() - .any(|row| row_pricing_model(row, Some(overlay)).1) - { - return None; - } let plain = days_from_codex_source_rows_with_priority(rows, None); if day_token_totals(&plain) != day_token_totals(parsed_days) { return None; } - Some(days_from_codex_source_rows_with_priority( - rows, - Some(overlay), - )) + let overlaid = days_from_codex_source_rows_with_priority(rows, Some(overlay)); + (overlaid != plain).then_some(overlaid) } fn day_token_totals(days: &DayModels) -> BTreeMap<&str, (i64, i64, i64)> { @@ -105,35 +99,6 @@ fn day_token_totals(days: &DayModels) -> BTreeMap<&str, (i64, i64, i64)> { totals } -/// The model a source row prices under, and whether Priority trace evidence -/// changed it. -fn row_pricing_model( - row: &CodexSourceUsageRow, - overlay: Option<&CodexPriorityOverlay<'_>>, -) -> (String, bool) { - let model = match row.pricing.pricing_model.as_deref() { - Some(model) if !model.is_empty() => { - if row.pricing.pricing_mode.as_deref() == Some("priority") - && !model.ends_with("-priority") - { - format!("{model}-priority") - } else { - model.to_string() - } - } - _ => { - return ( - CostUsagePricing::CODEX_UNATTRIBUTED_MODEL.to_string(), - false, - ); - } - }; - match overlay.and_then(|overlay| overlay.priority_model(row.turn_id.as_deref(), &model)) { - Some(priority) => (priority, true), - None => (model, false), - } -} - pub(super) fn days_from_codex_source_rows(rows: &[CodexSourceUsageRow]) -> DayModels { days_from_codex_source_rows_with_priority(rows, None) } @@ -144,7 +109,8 @@ fn days_from_codex_source_rows_with_priority( ) -> DayModels { let mut days: DayModels = HashMap::new(); for row in rows { - let (model, _) = row_pricing_model(row, overlay); + let model = priced_model(&row.pricing, row.turn_id.as_deref(), overlay) + .unwrap_or_else(|| CostUsagePricing::CODEX_UNATTRIBUTED_MODEL.to_string()); let record = CodexUsageRecord { day_key: row.day_key.clone(), timestamp: row.timestamp, diff --git a/rust/src/cost_scanner/codex/priority_trace.rs b/rust/src/cost_scanner/codex/priority_trace.rs index 8a665290be..69274cde93 100644 --- a/rust/src/cost_scanner/codex/priority_trace.rs +++ b/rust/src/cost_scanner/codex/priority_trace.rs @@ -79,20 +79,16 @@ pub(super) fn resolve_priority_turns( let previous = previous.filter(|cursor| cursor.database_path == path_key); let Ok(metadata) = std::fs::metadata(database_path) else { - // A missing optional source is normal until this path has supplied - // evidence; once it has, the evidence stays but validation is pending. - let pending = previous.is_some(); - return PriorityTraceResolution::keep(previous, pending); + // A missing optional source is normal, but should be retried on the + // next cost scan even when no trace evidence has been cached yet. + return PriorityTraceResolution::keep(previous, true); }; let Some(identity) = JsonlScanner::codex_file_identity(database_path, &metadata) else { return PriorityTraceResolution::keep(previous, true); }; let conn = match open_readonly_sqlite_connection(database_path, DEFAULT_SQLITE_BUSY_TIMEOUT) { Ok(conn) => conn, - Err(error) => { - tracing::debug!(%error, "Codex trace database could not be opened"); - return PriorityTraceResolution::keep(previous, true); - } + Err(_) => return PriorityTraceResolution::keep(previous, true), }; // The file may have been replaced between the identity read and the open. let identity_after = std::fs::metadata(database_path) @@ -102,7 +98,6 @@ pub(super) fn resolve_priority_turns( return PriorityTraceResolution::keep(previous, true); } let Some(max_row_id) = max_logs_row_id(&conn) else { - tracing::debug!("Codex trace database has no readable logs table"); return PriorityTraceResolution::keep(previous, true); }; @@ -125,6 +120,9 @@ pub(super) fn resolve_priority_turns( AnchorValidation::Invalid => state = None, } } + if let Some(cursor) = &mut state { + advance_coverage(cursor, coverage_since_epoch); + } let had_state = state.is_some(); let mut resolved = state.unwrap_or_else(fresh); @@ -182,6 +180,32 @@ pub(super) fn resolve_priority_turns( PriorityTraceResolution::keep(Some(resolved), false) } +/// Move a retained cursor with the active scan window and discard turn +/// evidence that can no longer match any scanned session rows. +fn advance_coverage(state: &mut CodexPriorityTurnsCursor, coverage_since_epoch: i64) { + if coverage_since_epoch <= state.coverage_since_epoch { + return; + } + state.coverage_since_epoch = coverage_since_epoch; + let mut expired_turns = Vec::new(); + state.request_sources.retain(|turn_id, sources| { + sources.retain(|_, metadata| { + metadata + .timestamp + .is_none_or(|timestamp| timestamp >= coverage_since_epoch) + }); + if sources.is_empty() { + expired_turns.push(turn_id.clone()); + false + } else { + true + } + }); + for turn_id in expired_turns { + state.priority_completed_models.remove(&turn_id); + } +} + /// Best-effort rollback; the read-only connection has nothing to lose. fn rollback(conn: &Connection) { if let Err(error) = conn.execute_batch("rollback") { @@ -306,7 +330,7 @@ fn absorb_row( body: &str, ) { if let Some(completed) = parse_completed_trace_row(body) { - if state.turns.contains_key(&completed.turn_id) { + if state.turn(&completed.turn_id).is_some() { state .priority_completed_models .entry(completed.turn_id) @@ -325,7 +349,6 @@ fn absorb_row( return; }; let turn_id = parsed.turn_id.clone(); - state.turns.insert(turn_id.clone(), parsed.clone()); state .request_sources .entry(turn_id.clone()) @@ -345,9 +368,12 @@ fn store_pending_completed_models( models: BTreeMap, ) { if !state.completed_models.contains_key(turn_id) { - state.completed_order.push(turn_id.to_string()); + state.completed_order.push_back(turn_id.to_string()); if state.completed_order.len() > CODEX_PRIORITY_COMPLETED_MODEL_RETENTION_LIMIT { - let evicted = state.completed_order.remove(0); + let evicted = state + .completed_order + .pop_front() + .expect("queue length checked"); state.completed_models.remove(&evicted); } } @@ -369,11 +395,37 @@ enum AnchorLookup { Failed, } -fn anchor_digest(timestamp: i64, body: Option<&[u8]>) -> String { +fn hash_anchor_value(hasher: &mut Sha256, value: ValueRef<'_>) { + match value { + ValueRef::Null => hasher.update([0]), + ValueRef::Integer(value) => { + hasher.update([1]); + hasher.update(value.to_le_bytes()); + } + ValueRef::Real(value) => { + hasher.update([2]); + hasher.update(value.to_bits().to_le_bytes()); + } + ValueRef::Text(value) => { + hasher.update([3]); + hasher.update(u64::try_from(value.len()).unwrap_or(u64::MAX).to_le_bytes()); + hasher.update(value); + } + ValueRef::Blob(value) => { + hasher.update([4]); + hasher.update(u64::try_from(value.len()).unwrap_or(u64::MAX).to_le_bytes()); + hasher.update(value); + } + } +} + +fn row_anchor_digest(row: &rusqlite::Row<'_>, timestamp_index: usize, body_index: usize) -> String { let mut hasher = Sha256::new(); - hasher.update(format!("{timestamp}\n").as_bytes()); - if let Some(body) = body { - hasher.update(body); + if let Ok(timestamp) = row.get_ref(timestamp_index) { + hash_anchor_value(&mut hasher, timestamp); + } + if let Ok(body) = row.get_ref(body_index) { + hash_anchor_value(&mut hasher, body); } hasher .finalize() @@ -382,18 +434,6 @@ fn anchor_digest(timestamp: i64, body: Option<&[u8]>) -> String { .collect() } -fn row_anchor_digest(row: &rusqlite::Row<'_>, timestamp_index: usize, body_index: usize) -> String { - let timestamp = match row.get_ref(timestamp_index) { - Ok(ValueRef::Integer(value)) => value, - _ => 0, - }; - let body = match row.get_ref(body_index) { - Ok(ValueRef::Text(bytes) | ValueRef::Blob(bytes)) => Some(bytes), - _ => None, - }; - anchor_digest(timestamp, body) -} - /// Sample four rows across the scanned range (a quarter, half, three /// quarters and the last) so a rewritten database cannot look unchanged. fn capture_anchors(conn: &Connection, state: &mut CodexPriorityTurnsCursor) -> bool { @@ -537,14 +577,10 @@ fn prune_deleted_sources(conn: &Connection, state: &mut CodexPriorityTurnsCursor pruned = true; if kept.is_empty() { state.request_sources.remove(&turn_id); - state.turns.remove(&turn_id); if let Some(models) = state.priority_completed_models.remove(&turn_id) { store_pending_completed_models(state, &turn_id, models); } } else { - if let Some((_, latest)) = kept.iter().next_back() { - state.turns.insert(turn_id.clone(), latest.clone()); - } state.request_sources.insert(turn_id, kept); } } diff --git a/rust/src/cost_scanner/codex/priority_trace/tests.rs b/rust/src/cost_scanner/codex/priority_trace/tests.rs index 7fc57df42f..8e6b2df1dd 100644 --- a/rust/src/cost_scanner/codex/priority_trace/tests.rs +++ b/rust/src/cost_scanner/codex/priority_trace/tests.rs @@ -143,8 +143,8 @@ fn cold_scan_collects_priority_turns_and_completed_model() { let resolution = resolve(&db, None); let cursor = resolution.cursor.unwrap(); assert!(!resolution.validation_pending); - assert_eq!(cursor.turns.len(), 1); - assert!(cursor.turns.contains_key("turn-1")); + assert_eq!(cursor.request_sources.len(), 1); + assert!(cursor.turn("turn-1").is_some()); assert_eq!(cursor.last_row_id, 6); assert_eq!(cursor.anchors.len(), 4); assert_eq!(cursor.turn_model("turn-1"), Some("gpt-5.4")); @@ -159,14 +159,14 @@ fn incremental_scan_appends_only_new_rows() { db.insert(&[(2_001, request_body("turn-2", "priority", "gpt-5.4"))]); let second = resolve(&db, Some(first)).cursor.unwrap(); assert_eq!(second.last_row_id, 2); - assert_eq!(second.turns.len(), 2); + assert_eq!(second.request_sources.len(), 2); } #[test] fn completion_before_request_is_matched_when_the_request_arrives() { let db = TraceDb::new(&[(2_000, completed_body("turn-1", "gpt-5.4"))]); let first = resolve(&db, None).cursor.unwrap(); - assert!(first.turns.is_empty()); + assert!(first.request_sources.is_empty()); assert!(first.completed_models.contains_key("turn-1")); db.insert(&[(2_001, request_body("turn-1", "priority", "gpt-5.5"))]); @@ -181,13 +181,13 @@ fn rewritten_database_rebuilds_instead_of_reusing_evidence() { rows.push((2_000, request_body("old-turn", "priority", "gpt-5.5"))); let db = TraceDb::new(&rows); let first = resolve(&db, None).cursor.unwrap(); - assert!(first.turns.contains_key("old-turn")); + assert!(first.turn("old-turn").is_some()); db.execute("update logs set feedback_log_body = 'rewritten ' || id"); db.insert(&[(3_000, request_body("new-turn", "priority", "gpt-5.5"))]); let second = resolve(&db, Some(first)).cursor.unwrap(); - assert!(!second.turns.contains_key("old-turn")); - assert!(second.turns.contains_key("new-turn")); + assert!(second.turn("old-turn").is_none()); + assert!(second.turn("new-turn").is_some()); } #[test] @@ -198,12 +198,12 @@ fn deleted_source_rows_drop_their_turns() { (2_002, "tail row".to_string()), ]); let first = resolve(&db, None).cursor.unwrap(); - assert_eq!(first.turns.len(), 2); + assert_eq!(first.request_sources.len(), 2); db.execute("delete from logs where id = 1"); let second = resolve(&db, Some(first)).cursor.unwrap(); - assert!(!second.turns.contains_key("turn-1")); - assert!(second.turns.contains_key("turn-2")); + assert!(second.turn("turn-1").is_none()); + assert!(second.turn("turn-2").is_some()); } #[test] @@ -215,7 +215,7 @@ fn cancelled_cold_scan_reports_pending_without_evidence() { assert!( resolution .cursor - .is_none_or(|cursor| cursor.turns.is_empty()) + .is_none_or(|cursor| cursor.request_sources.is_empty()) ); } @@ -227,13 +227,13 @@ fn missing_database_keeps_prior_evidence_without_error() { let missing = db.path.with_file_name("absent.sqlite"); let none = resolve_priority_turns(&missing, None, 0, None); assert!(none.cursor.is_none()); - assert!(!none.validation_pending); + assert!(none.validation_pending); let mut previous = cursor; previous.database_path = missing.to_string_lossy().to_string(); let kept = resolve_priority_turns(&missing, Some(previous), 0, None); assert!(kept.validation_pending); - assert_eq!(kept.cursor.unwrap().turns.len(), 1); + assert_eq!(kept.cursor.unwrap().request_sources.len(), 1); } #[test] @@ -251,16 +251,19 @@ fn cursor_for_another_database_is_discarded() { let db = TraceDb::new(&[(2_000, request_body("turn-1", "priority", "gpt-5.5"))]); let mut stale = resolve(&db, None).cursor.unwrap(); stale.database_path = "elsewhere".to_string(); - stale.turns.insert( + stale.request_sources.insert( "ghost".to_string(), - CodexPriorityTurnMetadata { - turn_id: "ghost".to_string(), - ..CodexPriorityTurnMetadata::default() - }, + BTreeMap::from([( + 2, + CodexPriorityTurnMetadata { + turn_id: "ghost".to_string(), + ..CodexPriorityTurnMetadata::default() + }, + )]), ); let cursor = resolve(&db, Some(stale)).cursor.unwrap(); - assert!(!cursor.turns.contains_key("ghost")); - assert!(cursor.turns.contains_key("turn-1")); + assert!(cursor.turn("ghost").is_none()); + assert!(cursor.turn("turn-1").is_some()); } #[test] @@ -272,8 +275,48 @@ fn coverage_window_skips_older_history() { let cursor = resolve_priority_turns(&db.path, None, 1_000, None) .cursor .unwrap(); - assert!(!cursor.turns.contains_key("ancient")); - assert!(cursor.turns.contains_key("recent")); + assert!(cursor.turn("ancient").is_none()); + assert!(cursor.turn("recent").is_some()); +} + +#[test] +fn advancing_coverage_prunes_expired_turns_without_restarting_the_cursor() { + let db = TraceDb::new(&[ + (1_000, request_body("expired", "priority", "gpt-5.5")), + (2_000, request_body("current", "priority", "gpt-5.5")), + ]); + let first = resolve(&db, None).cursor.unwrap(); + let last_row_id = first.last_row_id; + assert!(first.turn("expired").is_some()); + + let updated = resolve_priority_turns(&db.path, Some(first), 1_500, None) + .cursor + .unwrap(); + + assert_eq!(updated.last_row_id, last_row_id); + assert_eq!(updated.coverage_since_epoch, 1_500); + assert!(updated.turn("expired").is_none()); + assert!(updated.turn("current").is_some()); +} + +#[test] +fn anchor_digest_tracks_fractional_sqlite_timestamps() { + let db = TraceDb::new(&noise(6)); + let first = resolve(&db, None).cursor.unwrap(); + let anchor = first.anchors[0].clone(); + + db.execute(&format!( + "update logs set ts = ts + 0.25 where id = {}", + anchor.row_id + )); + + let updated = resolve(&db, Some(first)).cursor.unwrap(); + let updated_anchor = updated + .anchors + .iter() + .find(|candidate| candidate.row_id == anchor.row_id) + .unwrap(); + assert_ne!(updated_anchor.digest, anchor.digest); } #[test] From 7cfde609534dfc53f31f7d4a575250aff1d17b57 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:18:29 +0700 Subject: [PATCH 131/259] Address thermo review --- rust/src/browser/leveldb/log.rs | 12 +++++- rust/src/browser/leveldb/mod.rs | 20 +++++++--- rust/src/browser/leveldb/table.rs | 65 ++++++++++++++++++++++--------- 3 files changed, 71 insertions(+), 26 deletions(-) diff --git a/rust/src/browser/leveldb/log.rs b/rust/src/browser/leveldb/log.rs index 98cf322948..1256a818cb 100644 --- a/rust/src/browser/leveldb/log.rs +++ b/rust/src/browser/leveldb/log.rs @@ -84,6 +84,7 @@ fn read_batch(batch: &[u8], emit: &mut impl FnMut(Record)) { return; }; let mut rest = batch.get(BATCH_HEADER_SIZE..).unwrap_or_default(); + let mut records = Vec::new(); for index in 0..u64::from(count) { let Some((&op, after_op)) = rest.split_first() else { return; @@ -105,10 +106,17 @@ fn read_batch(batch: &[u8], emit: &mut impl FnMut(Record)) { } _ => return, }; - emit(Record { + let Some(sequence) = sequence.checked_add(index) else { + return; + }; + records.push(Record { key: key.to_vec(), - sequence: sequence.wrapping_add(index), + sequence, value, }); } + if !rest.is_empty() { + return; + } + records.into_iter().for_each(emit); } diff --git a/rust/src/browser/leveldb/mod.rs b/rust/src/browser/leveldb/mod.rs index 19af9081dd..aa95a9cb42 100644 --- a/rust/src/browser/leveldb/mod.rs +++ b/rust/src/browser/leveldb/mod.rs @@ -13,8 +13,9 @@ //! - Checksums are not verified and only the bytewise comparator is assumed (which is what //! `Local Storage` uses; a full scan does not depend on ordering anyway). //! - Blocks using a compression type other than none/Snappy are skipped. -//! - Files larger than [`MAX_FILE_BYTES`] and blocks that inflate beyond [`MAX_BLOCK_BYTES`] are -//! skipped so a corrupt profile cannot exhaust memory. +//! - Files larger than [`MAX_FILE_BYTES`] and decoded table blocks larger than +//! [`MAX_BLOCK_BYTES`] are skipped to bound each input allocation. The returned snapshot still +//! grows with the database's live data. pub mod local_storage; mod log; @@ -27,11 +28,12 @@ mod tests; use std::collections::BTreeMap; use std::ffi::OsStr; +use std::io::Read; use std::path::Path; /// Largest single log or table file that will be read into memory. pub const MAX_FILE_BYTES: u64 = 64 * 1024 * 1024; -/// Largest decompressed table block accepted. +/// Largest decoded table block accepted. pub const MAX_BLOCK_BYTES: usize = 16 * 1024 * 1024; /// One live key/value pair of the database. @@ -124,11 +126,19 @@ impl FileKind { } fn read_bounded_file(path: &Path) -> std::io::Result> { - let size = std::fs::metadata(path)?.len(); + let file = std::fs::File::open(path)?; + let size = file.metadata()?.len(); if size > MAX_FILE_BYTES { return Err(std::io::Error::other(format!( "file is {size} bytes, over the {MAX_FILE_BYTES} byte limit" ))); } - std::fs::read(path) + let mut data = Vec::new(); + file.take(MAX_FILE_BYTES + 1).read_to_end(&mut data)?; + if data.len() as u64 > MAX_FILE_BYTES { + return Err(std::io::Error::other(format!( + "file grew beyond the {MAX_FILE_BYTES} byte limit while being read" + ))); + } + Ok(data) } diff --git a/rust/src/browser/leveldb/table.rs b/rust/src/browser/leveldb/table.rs index 0bf805c6bf..30e8bacfe4 100644 --- a/rust/src/browser/leveldb/table.rs +++ b/rust/src/browser/leveldb/table.rs @@ -31,6 +31,8 @@ pub(super) enum TableError { BadFooter, #[error("block handle points outside the file")] BadHandle, + #[error("table block is {0} bytes, over the {1} byte limit")] + BlockTooLarge(usize, usize), #[error("unsupported block compression type {0}")] UnsupportedCompression(u8), #[error("block is malformed")] @@ -73,29 +75,34 @@ pub(super) fn read_table(data: &[u8], emit: &mut impl FnMut(Record)) -> Result = BlockEntries::new(&index_block)?.collect::>()?; let mut skipped_blocks = 0usize; - for (_separator, handle_bytes) in BlockEntries::new(&index_block)? { + for (_separator, handle_bytes) in index_entries { let entry = BlockHandle::parse(handle_bytes) - .map(|(handle, _)| handle) + .and_then(|(handle, used)| (used == handle_bytes.len()).then_some(handle)) .ok_or(TableError::BadBlock); - let block = entry.and_then(|handle| read_block(data, handle)); + let block = entry.and_then(|handle| read_block(data, handle, index.offset)); let Ok(block) = block else { skipped_blocks += 1; continue; }; - if emit_block_records(&block, emit).is_err() { - skipped_blocks += 1; + match decode_block_records(&block) { + Ok(records) => records.into_iter().for_each(&mut *emit), + Err(_) => skipped_blocks += 1, } } Ok(skipped_blocks) } -fn emit_block_records(block: &[u8], emit: &mut impl FnMut(Record)) -> Result<(), TableError> { - for (internal_key, value) in BlockEntries::new(block)? { +fn decode_block_records(block: &[u8]) -> Result, TableError> { + let mut records = Vec::new(); + for entry in BlockEntries::new(block)? { + let (internal_key, value) = entry?; let Some(split) = internal_key.len().checked_sub(8) else { return Err(TableError::BadBlock); }; @@ -104,25 +111,28 @@ fn emit_block_records(block: &[u8], emit: &mut impl FnMut(Record)) -> Result<(), let value = match packed & 0xff { KIND_VALUE => Some(value.to_vec()), KIND_DELETE => None, - _ => continue, + _ => return Err(TableError::BadBlock), }; - emit(Record { + records.push(Record { key: user_key.to_vec(), sequence: packed >> 8, value, }); } - Ok(()) + Ok(records) } /// Return the decompressed contents of the block at `handle` (without its trailer). -fn read_block(data: &[u8], handle: BlockHandle) -> Result, TableError> { +fn read_block(data: &[u8], handle: BlockHandle, upper_bound: usize) -> Result, TableError> { + if handle.offset >= upper_bound { + return Err(TableError::BadHandle); + } let end = handle .offset .checked_add(handle.size) .and_then(|end| end.checked_add(BLOCK_TRAILER_SIZE)) .ok_or(TableError::BadHandle)?; - if data.len() < end { + if end > upper_bound || data.len() < end { return Err(TableError::BadHandle); } let contents_end = end - BLOCK_TRAILER_SIZE; @@ -131,7 +141,8 @@ fn read_block(data: &[u8], handle: BlockHandle) -> Result, TableError> { .ok_or(TableError::BadHandle)?; let compression = data[contents_end]; match compression { - COMPRESSION_NONE => Ok(raw.to_vec()), + COMPRESSION_NONE if raw.len() <= MAX_BLOCK_BYTES => Ok(raw.to_vec()), + COMPRESSION_NONE => Err(TableError::BlockTooLarge(raw.len(), MAX_BLOCK_BYTES)), COMPRESSION_SNAPPY => Ok(snappy::decompress(raw, MAX_BLOCK_BYTES)?), other => Err(TableError::UnsupportedCompression(other)), } @@ -148,6 +159,9 @@ impl<'a> BlockEntries<'a> { fn new(block: &'a [u8]) -> Result { let restart_count = read_u32_le(block, block.len().saturating_sub(4)).ok_or(TableError::BadBlock)? as usize; + if restart_count == 0 { + return Err(TableError::BadBlock); + } let restarts_len = restart_count .checked_mul(4) .and_then(|len| len.checked_add(4)) @@ -156,6 +170,19 @@ impl<'a> BlockEntries<'a> { .len() .checked_sub(restarts_len) .ok_or(TableError::BadBlock)?; + let mut previous_restart = None; + for index in 0..restart_count { + let restart = + read_u32_le(block, entries_end + index * 4).ok_or(TableError::BadBlock)? as usize; + if (index == 0 && restart != 0) + || restart > entries_end + || (entries_end > 0 && restart == entries_end) + || previous_restart.is_some_and(|previous| restart <= previous) + { + return Err(TableError::BadBlock); + } + previous_restart = Some(restart); + } Ok(Self { entries: &block[..entries_end], key: Vec::new(), @@ -165,7 +192,7 @@ impl<'a> BlockEntries<'a> { } impl<'a> Iterator for BlockEntries<'a> { - type Item = (Vec, &'a [u8]); + type Item = Result<(Vec, &'a [u8]), TableError>; fn next(&mut self) -> Option { if self.failed || self.entries.is_empty() { @@ -176,11 +203,11 @@ impl<'a> Iterator for BlockEntries<'a> { Some((key, value, rest)) => { self.key.clone_from(&key); self.entries = rest; - Some((key, value)) + Some(Ok((key, value))) } None => { self.failed = true; - None + Some(Err(TableError::BadBlock)) } } } From 2542a90f3c8c2dc059ae4c3991dfbb92a946a70b Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:56:51 +0700 Subject: [PATCH 132/259] Address thermo review --- rust/src/browser/leveldb/local_storage.rs | 41 ++++- rust/src/browser/leveldb/log.rs | 108 ++++++++--- rust/src/browser/leveldb/mod.rs | 215 ++++++++++++++++++---- rust/src/browser/leveldb/snappy.rs | 16 +- rust/src/browser/leveldb/table.rs | 141 ++++++++++++-- rust/src/providers/kimi/local_storage.rs | 30 ++- rust/src/providers/kimi/web.rs | 12 +- 7 files changed, 467 insertions(+), 96 deletions(-) diff --git a/rust/src/browser/leveldb/local_storage.rs b/rust/src/browser/leveldb/local_storage.rs index a2e3d5236e..e9bc81b1ae 100644 --- a/rust/src/browser/leveldb/local_storage.rs +++ b/rust/src/browser/leveldb/local_storage.rs @@ -6,7 +6,7 @@ //! the same format byte. Other keys in the database (`VERSION`, `META:`, //! `METAACCESS:`) are bookkeeping and are ignored. -use super::{Entry, LevelDbError, read_entries}; +use super::{Entry, LevelDbError, MAX_TOTAL_SCAN_BYTES, read_entries_with_budget}; use std::path::{Path, PathBuf}; const KEY_PREFIX: u8 = b'_'; @@ -15,12 +15,22 @@ const FORMAT_UTF16LE: u8 = 0; const FORMAT_LATIN1: u8 = 1; /// One decoded `localStorage` item. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct LocalStorageEntry { pub key: String, pub value: String, } +impl std::fmt::Debug for LocalStorageEntry { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LocalStorageEntry") + .field("key", &self.key) + .field("value", &"[REDACTED]") + .finish() + } +} + /// Directory holding Local Storage for a Chromium profile directory (`Default`, `Profile 1`, ...). pub fn local_storage_dir(profile_dir: &Path) -> PathBuf { profile_dir.join("Local Storage").join("leveldb") @@ -32,14 +42,23 @@ pub fn read_local_storage_entries( dir: &Path, origin: &str, ) -> Result, LevelDbError> { - Ok(decode_origin_entries(&read_entries(dir)?, origin)) + read_local_storage_entries_with_budget(dir, origin, MAX_TOTAL_SCAN_BYTES) + .map(|(entries, _)| entries) +} + +pub(crate) fn read_local_storage_entries_with_budget( + dir: &Path, + origin: &str, + max_total_bytes: u64, +) -> Result<(Vec, u64), LevelDbError> { + let prefix = origin_key_prefix(origin); + let (entries, scanned_bytes) = + read_entries_with_budget(dir, max_total_bytes, Some(prefix.as_slice()))?; + Ok((decode_origin_entries(&entries, origin), scanned_bytes)) } pub(super) fn decode_origin_entries(entries: &[Entry], origin: &str) -> Vec { - let mut prefix = Vec::with_capacity(origin.len() + 2); - prefix.push(KEY_PREFIX); - prefix.extend_from_slice(origin.trim_end_matches('/').as_bytes()); - prefix.push(ORIGIN_TERMINATOR); + let prefix = origin_key_prefix(origin); entries .iter() @@ -51,6 +70,14 @@ pub(super) fn decode_origin_entries(entries: &[Entry], origin: &str) -> Vec Vec { + let mut prefix = Vec::with_capacity(origin.len() + 2); + prefix.push(KEY_PREFIX); + prefix.extend_from_slice(origin.trim_end_matches('/').as_bytes()); + prefix.push(ORIGIN_TERMINATOR); + prefix +} + /// Decode a format-byte-prefixed Chromium string; `None` for an unknown format or bad UTF-16. fn decode_text(bytes: &[u8]) -> Option { let (&format, data) = bytes.split_first()?; diff --git a/rust/src/browser/leveldb/log.rs b/rust/src/browser/leveldb/log.rs index 98cf322948..d8e01b4a14 100644 --- a/rust/src/browser/leveldb/log.rs +++ b/rust/src/browser/leveldb/log.rs @@ -9,12 +9,13 @@ //! malformed tail ends the scan quietly and everything before it is kept. Record checksums are //! not verified: this is a best-effort read of another program's cache, not a database recovery. -use super::Record; use super::varint::{read_length_prefixed, read_u32_le, read_u64_le}; +use super::{MAX_RECORDS_PER_DIRECTORY, Record}; const BLOCK_SIZE: usize = 32 * 1024; const HEADER_SIZE: usize = 7; const BATCH_HEADER_SIZE: usize = 12; +const MAX_LOG_RECORD_BYTES: usize = 16 * 1024 * 1024; const TYPE_ZERO: u8 = 0; const TYPE_FULL: u8 = 1; @@ -24,9 +25,19 @@ const TYPE_LAST: u8 = 4; const OP_DELETE: u8 = 0; const OP_PUT: u8 = 1; +const MAX_SEQUENCE: u64 = (1 << 56) - 1; /// Feed every put/delete found in `data` to `emit`. +#[cfg(test)] pub(super) fn read_log(data: &[u8], emit: &mut impl FnMut(Record)) { + let _ = read_log_until(data, &mut |record| { + emit(record); + true + }); +} + +/// Read a log until it is malformed or `emit` asks the scan to stop. +pub(super) fn read_log_until(data: &[u8], emit: &mut impl FnMut(Record) -> bool) -> bool { let mut assembled: Vec = Vec::new(); let mut in_fragmented = false; @@ -38,7 +49,7 @@ pub(super) fn read_log(data: &[u8], emit: &mut impl FnMut(Record)) { continue; } let Some(header) = data.get(offset..offset + HEADER_SIZE) else { - return; + return false; }; let length = usize::from(u16::from_le_bytes([header[4], header[5]])); let kind = header[6]; @@ -49,10 +60,10 @@ pub(super) fn read_log(data: &[u8], emit: &mut impl FnMut(Record)) { } let payload_start = offset + HEADER_SIZE; if HEADER_SIZE + length > block_remaining { - return; + return false; } let Some(payload) = data.get(payload_start..payload_start + length) else { - return; + return false; }; offset = payload_start + length; @@ -60,55 +71,106 @@ pub(super) fn read_log(data: &[u8], emit: &mut impl FnMut(Record)) { TYPE_FULL => { assembled.clear(); in_fragmented = false; - read_batch(payload, emit); + if !read_batch(payload, emit) { + return false; + } } TYPE_FIRST => { assembled.clear(); - assembled.extend_from_slice(payload); + if !append_fragment(&mut assembled, payload) { + return false; + } in_fragmented = true; } - TYPE_MIDDLE if in_fragmented => assembled.extend_from_slice(payload), + TYPE_MIDDLE if in_fragmented => { + if !append_fragment(&mut assembled, payload) { + return false; + } + } TYPE_LAST if in_fragmented => { - assembled.extend_from_slice(payload); + if !append_fragment(&mut assembled, payload) { + return false; + } in_fragmented = false; - read_batch(&assembled, emit); + if !read_batch(&assembled, emit) { + return false; + } assembled.clear(); } - _ => return, + _ => return false, } } + true +} + +fn append_fragment(assembled: &mut Vec, payload: &[u8]) -> bool { + let Some(new_len) = assembled.len().checked_add(payload.len()) else { + return false; + }; + if new_len > MAX_LOG_RECORD_BYTES || assembled.try_reserve(payload.len()).is_err() { + return false; + } + assembled.extend_from_slice(payload); + true +} + +fn read_batch(batch: &[u8], emit: &mut impl FnMut(Record) -> bool) -> bool { + if !visit_batch(batch, &mut |_, _, _| true) { + return true; + } + visit_batch(batch, &mut |sequence, key, value| { + emit(Record { + key: key.to_vec(), + sequence, + value: value.map(|value| value.to_vec()), + }) + }) } -fn read_batch(batch: &[u8], emit: &mut impl FnMut(Record)) { +/// Validate the whole batch before applying any of its operations. A torn or malformed write +/// batch must not leave a valid-looking prefix in the returned database state. +fn visit_batch<'a>( + batch: &'a [u8], + emit: &mut impl FnMut(u64, &'a [u8], Option<&'a [u8]>) -> bool, +) -> bool { let (Some(sequence), Some(count)) = (read_u64_le(batch, 0), read_u32_le(batch, 8)) else { - return; + return false; + }; + if sequence > MAX_SEQUENCE || u64::from(count) > MAX_RECORDS_PER_DIRECTORY as u64 { + return false; + } + let mut rest = match batch.get(BATCH_HEADER_SIZE..) { + Some(rest) => rest, + None => return false, }; - let mut rest = batch.get(BATCH_HEADER_SIZE..).unwrap_or_default(); for index in 0..u64::from(count) { let Some((&op, after_op)) = rest.split_first() else { - return; + return false; }; let Some((key, after_key)) = read_length_prefixed(after_op) else { - return; + return false; }; let value = match op { OP_PUT => { let Some((value, after_value)) = read_length_prefixed(after_key) else { - return; + return false; }; rest = after_value; - Some(value.to_vec()) + Some(value) } OP_DELETE => { rest = after_key; None } - _ => return, + _ => return false, }; - emit(Record { - key: key.to_vec(), - sequence: sequence.wrapping_add(index), - value, - }); + let Some(record_sequence) = sequence.checked_add(index).filter(|s| *s <= MAX_SEQUENCE) + else { + return false; + }; + if !emit(record_sequence, key, value) { + return false; + } } + rest.is_empty() } diff --git a/rust/src/browser/leveldb/mod.rs b/rust/src/browser/leveldb/mod.rs index 19af9081dd..ea5bbefd13 100644 --- a/rust/src/browser/leveldb/mod.rs +++ b/rust/src/browser/leveldb/mod.rs @@ -13,8 +13,9 @@ //! - Checksums are not verified and only the bytewise comparator is assumed (which is what //! `Local Storage` uses; a full scan does not depend on ordering anyway). //! - Blocks using a compression type other than none/Snappy are skipped. -//! - Files larger than [`MAX_FILE_BYTES`] and blocks that inflate beyond [`MAX_BLOCK_BYTES`] are -//! skipped so a corrupt profile cannot exhaust memory. +//! - Files larger than [`MAX_FILE_BYTES`], scans larger than [`MAX_TOTAL_SCAN_BYTES`], and blocks +//! that decode beyond [`MAX_BLOCK_BYTES`] are skipped so a corrupt profile cannot exhaust +//! memory or monopolize a refresh. pub mod local_storage; mod log; @@ -27,25 +28,44 @@ mod tests; use std::collections::BTreeMap; use std::ffi::OsStr; +use std::io::Read; use std::path::Path; /// Largest single log or table file that will be read into memory. pub const MAX_FILE_BYTES: u64 = 64 * 1024 * 1024; +/// Largest total number of bytes scanned across a LevelDB directory. +pub const MAX_TOTAL_SCAN_BYTES: u64 = 256 * 1024 * 1024; /// Largest decompressed table block accepted. pub const MAX_BLOCK_BYTES: usize = 16 * 1024 * 1024; +const MAX_DIRECTORY_ENTRIES: usize = 8192; +const MAX_RECORDS_PER_DIRECTORY: usize = 250_000; +const MAX_LIVE_ENTRIES: usize = 100_000; +const MAX_LIVE_ENTRY_BYTES: usize = 64 * 1024 * 1024; /// One live key/value pair of the database. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Clone, PartialEq, Eq)] pub struct Entry { pub key: Vec, pub value: Vec, } +impl std::fmt::Debug for Entry { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("Entry") + .field("key", &"[REDACTED]") + .field("value", &"[REDACTED]") + .finish() + } +} + /// Failure to read a LevelDB directory at all (per-file problems are skipped, not reported). #[derive(Debug, thiserror::Error)] pub enum LevelDbError { #[error("cannot read LevelDB directory: {0}")] Io(#[from] std::io::Error), + #[error("LevelDB directory exceeds safe scan limits")] + ResourceLimit, } /// A put (`value: Some`) or delete (`value: None`) with the sequence number it was written at. @@ -55,53 +75,154 @@ pub(crate) struct Record { pub(crate) value: Option>, } +#[derive(Default)] +struct EntryAccumulator<'a> { + key_prefix: Option<&'a [u8]>, + newest: BTreeMap, (u64, Option>)>, + record_count: usize, + retained_bytes: usize, + resource_limit_exceeded: bool, +} + +impl<'a> EntryAccumulator<'a> { + fn absorb(&mut self, record: Record) -> bool { + if self.record_count == MAX_RECORDS_PER_DIRECTORY { + self.resource_limit_exceeded = true; + return false; + } + self.record_count += 1; + if self + .key_prefix + .is_some_and(|prefix| !record.key.starts_with(prefix)) + { + return true; + } + + let current_entry_count = self.newest.len(); + match self.newest.entry(record.key) { + std::collections::btree_map::Entry::Occupied(mut entry) => { + let existing = entry.get_mut(); + if existing.0 >= record.sequence { + return true; + } + let old_value_bytes = existing.1.as_ref().map_or(0, Vec::len); + let new_value_bytes = record.value.as_ref().map_or(0, Vec::len); + let next_retained_bytes = self + .retained_bytes + .checked_sub(old_value_bytes) + .and_then(|bytes| bytes.checked_add(new_value_bytes)); + let Some(next_retained_bytes) = next_retained_bytes else { + self.resource_limit_exceeded = true; + return false; + }; + if next_retained_bytes > MAX_LIVE_ENTRY_BYTES { + self.resource_limit_exceeded = true; + return false; + } + self.retained_bytes = next_retained_bytes; + *existing = (record.sequence, record.value); + true + } + std::collections::btree_map::Entry::Vacant(entry) => { + let Some(entry_bytes) = entry + .key() + .len() + .checked_add(record.value.as_ref().map_or(0, Vec::len)) + else { + self.resource_limit_exceeded = true; + return false; + }; + let Some(next_retained_bytes) = self.retained_bytes.checked_add(entry_bytes) else { + self.resource_limit_exceeded = true; + return false; + }; + if current_entry_count == MAX_LIVE_ENTRIES + || next_retained_bytes > MAX_LIVE_ENTRY_BYTES + { + self.resource_limit_exceeded = true; + return false; + } + self.retained_bytes = next_retained_bytes; + entry.insert((record.sequence, record.value)); + true + } + } + } + + fn into_entries(self) -> Vec { + self.newest + .into_iter() + .filter_map(|(key, (_, value))| value.map(|value| Entry { key, value })) + .collect() + } +} + /// Read every live entry of the LevelDB directory `dir`, sorted by key. /// /// Keys that were deleted (or whose newest record is a delete) are omitted. Unreadable or corrupt -/// files are skipped with a debug log so the remaining data is still returned. +/// files are skipped with a debug log; scans over the configured resource limits fail as a whole. pub fn read_entries(dir: &Path) -> Result, LevelDbError> { - let mut newest: BTreeMap, (u64, Option>)> = BTreeMap::new(); - let mut absorb = |record: Record| match newest.get_mut(&record.key) { - Some(existing) if existing.0 >= record.sequence => {} - Some(existing) => *existing = (record.sequence, record.value), - None => { - newest.insert(record.key, (record.sequence, record.value)); - } + read_entries_with_budget(dir, MAX_TOTAL_SCAN_BYTES, None).map(|(entries, _)| entries) +} + +pub(crate) fn read_entries_with_budget( + dir: &Path, + max_total_bytes: u64, + key_prefix: Option<&[u8]>, +) -> Result<(Vec, u64), LevelDbError> { + let mut entries = EntryAccumulator { + key_prefix, + ..EntryAccumulator::default() }; + let mut scanned_bytes = 0u64; - for dir_entry in std::fs::read_dir(dir)? { + for (directory_entry_count, dir_entry) in std::fs::read_dir(dir)?.enumerate() { + if directory_entry_count == MAX_DIRECTORY_ENTRIES { + return Err(LevelDbError::ResourceLimit); + } let Ok(dir_entry) = dir_entry else { continue }; let path = dir_entry.path(); let Some(kind) = FileKind::of(&path) else { continue; }; - let data = match read_bounded_file(&path) { + let remaining_bytes = max_total_bytes.saturating_sub(scanned_bytes); + let data = match read_bounded_file(&path, remaining_bytes) { Ok(data) => data, - Err(error) => { + Err(BoundedReadError::FileTooLarge) => continue, + Err(BoundedReadError::BudgetExceeded) => return Err(LevelDbError::ResourceLimit), + Err(BoundedReadError::Io(error)) => { tracing::debug!(file = ?path.file_name(), %error, "skipping unreadable LevelDB file"); continue; } }; + let file_bytes = u64::try_from(data.len()).map_err(|_| LevelDbError::ResourceLimit)?; + scanned_bytes = scanned_bytes + .checked_add(file_bytes) + .ok_or(LevelDbError::ResourceLimit)?; match kind { - FileKind::Log => log::read_log(&data, &mut absorb), - FileKind::Table => match table::read_table(&data, &mut absorb) { - Ok(0) => {} - Ok(skipped) => tracing::debug!( - file = ?path.file_name(), - skipped, - "skipped undecodable LevelDB table blocks" - ), - Err(error) => { - tracing::debug!(file = ?path.file_name(), %error, "skipping malformed LevelDB table"); + FileKind::Log => { + let _ = log::read_log_until(&data, &mut |record| entries.absorb(record)); + } + FileKind::Table => { + match table::read_table_until(&data, &mut |record| entries.absorb(record)) { + Ok((0, _)) => {} + Ok((skipped, _)) => tracing::debug!( + file = ?path.file_name(), + skipped, + "skipped undecodable LevelDB table blocks" + ), + Err(error) => { + tracing::debug!(file = ?path.file_name(), %error, "skipping malformed LevelDB table"); + } } - }, + } + } + if entries.resource_limit_exceeded { + return Err(LevelDbError::ResourceLimit); } } - Ok(newest - .into_iter() - .filter_map(|(key, (_, value))| value.map(|value| Entry { key, value })) - .collect()) + Ok((entries.into_entries(), scanned_bytes)) } #[derive(Clone, Copy)] @@ -123,12 +244,36 @@ impl FileKind { } } -fn read_bounded_file(path: &Path) -> std::io::Result> { - let size = std::fs::metadata(path)?.len(); +enum BoundedReadError { + Io(std::io::Error), + FileTooLarge, + BudgetExceeded, +} + +fn read_bounded_file(path: &Path, remaining_bytes: u64) -> Result, BoundedReadError> { + let file = std::fs::File::open(path).map_err(BoundedReadError::Io)?; + let size = file.metadata().map_err(BoundedReadError::Io)?.len(); if size > MAX_FILE_BYTES { - return Err(std::io::Error::other(format!( - "file is {size} bytes, over the {MAX_FILE_BYTES} byte limit" - ))); + return Err(BoundedReadError::FileTooLarge); + } + if size > remaining_bytes { + return Err(BoundedReadError::BudgetExceeded); + } + + let read_limit = MAX_FILE_BYTES.min(remaining_bytes); + let reserve = usize::try_from(size).map_err(|_| BoundedReadError::FileTooLarge)?; + let mut data = Vec::new(); + data.try_reserve_exact(reserve) + .map_err(|error| BoundedReadError::Io(std::io::Error::other(error)))?; + file.take(read_limit + 1) + .read_to_end(&mut data) + .map_err(BoundedReadError::Io)?; + let data_len = u64::try_from(data.len()).map_err(|_| BoundedReadError::FileTooLarge)?; + if data_len > MAX_FILE_BYTES { + return Err(BoundedReadError::FileTooLarge); + } + if data_len > remaining_bytes { + return Err(BoundedReadError::BudgetExceeded); } - std::fs::read(path) + Ok(data) } diff --git a/rust/src/browser/leveldb/snappy.rs b/rust/src/browser/leveldb/snappy.rs index 3436a621df..d61542e015 100644 --- a/rust/src/browser/leveldb/snappy.rs +++ b/rust/src/browser/leveldb/snappy.rs @@ -17,6 +17,8 @@ pub enum SnappyError { Truncated, #[error("snappy copy references data before the start of the output")] BadOffset, + #[error("snappy output buffer could not be allocated")] + AllocationFailed, #[error("snappy stream length does not match its preamble")] LengthMismatch, } @@ -29,7 +31,9 @@ pub fn decompress(input: &[u8], max_len: usize) -> Result, SnappyError> return Err(SnappyError::TooLarge(expected, max_len)); } - let mut out = Vec::with_capacity(expected); + let mut out = Vec::new(); + out.try_reserve_exact(expected) + .map_err(|_| SnappyError::AllocationFailed)?; while pos < input.len() { let tag = input[pos]; pos += 1; @@ -50,7 +54,10 @@ pub fn decompress(input: &[u8], max_len: usize) -> Result, SnappyError> .checked_add(len) .and_then(|end| input.get(pos..end)) .ok_or(SnappyError::Truncated)?; - if out.len() + literal.len() > expected { + let Some(new_len) = out.len().checked_add(literal.len()) else { + return Err(SnappyError::LengthMismatch); + }; + if new_len > expected { return Err(SnappyError::LengthMismatch); } out.extend_from_slice(literal); @@ -105,7 +112,10 @@ fn copy_within_output( if offset == 0 || offset > out.len() { return Err(SnappyError::BadOffset); } - if out.len() + len > expected { + let Some(new_len) = out.len().checked_add(len) else { + return Err(SnappyError::LengthMismatch); + }; + if new_len > expected { return Err(SnappyError::LengthMismatch); } let start = out.len() - offset; diff --git a/rust/src/browser/leveldb/table.rs b/rust/src/browser/leveldb/table.rs index 0bf805c6bf..314cdee15c 100644 --- a/rust/src/browser/leveldb/table.rs +++ b/rust/src/browser/leveldb/table.rs @@ -13,7 +13,7 @@ use super::snappy; use super::varint::{read_u32_le, read_u64_le, read_varint32, read_varint64}; -use super::{MAX_BLOCK_BYTES, Record}; +use super::{MAX_BLOCK_BYTES, MAX_RECORDS_PER_DIRECTORY, Record}; const FOOTER_SIZE: usize = 48; const TABLE_MAGIC: u64 = 0xdb47_7524_8b80_fb57; @@ -35,6 +35,10 @@ pub(super) enum TableError { UnsupportedCompression(u8), #[error("block is malformed")] BadBlock, + #[error("block exceeds the configured size limit")] + BlockTooLarge, + #[error("block output buffer could not be allocated")] + AllocationFailed, #[error("snappy block failed to decode: {0}")] Snappy(#[from] snappy::SnappyError), } @@ -64,7 +68,19 @@ impl BlockHandle { /// /// A data block that cannot be decoded is skipped (and counted in the returned failure count) so /// one bad block does not hide the rest of the table; a bad footer or index block is an error. +#[cfg(test)] pub(super) fn read_table(data: &[u8], emit: &mut impl FnMut(Record)) -> Result { + read_table_until(data, &mut |record| { + emit(record); + true + }) + .map(|(skipped, _)| skipped) +} + +pub(super) fn read_table_until( + data: &[u8], + emit: &mut impl FnMut(Record) -> bool, +) -> Result<(usize, bool), TableError> { let footer_start = data .len() .checked_sub(FOOTER_SIZE) @@ -77,52 +93,103 @@ pub(super) fn read_table(data: &[u8], emit: &mut impl FnMut(Record)) -> Result {} + Ok(false) => return Ok((skipped_blocks, false)), + Err(_) => skipped_blocks += 1, } } - Ok(skipped_blocks) + Ok((skipped_blocks, true)) } -fn emit_block_records(block: &[u8], emit: &mut impl FnMut(Record)) -> Result<(), TableError> { - for (internal_key, value) in BlockEntries::new(block)? { +fn validate_index_block(block: &[u8]) -> Result<(), TableError> { + let mut entries = BlockEntries::new(block)?; + let mut entry_count = 0usize; + for (_separator, handle_bytes) in entries.by_ref() { + entry_count += 1; + if entry_count > MAX_RECORDS_PER_DIRECTORY { + return Err(TableError::BadBlock); + } + let (_, used) = BlockHandle::parse(handle_bytes).ok_or(TableError::BadBlock)?; + if used != handle_bytes.len() { + return Err(TableError::BadBlock); + } + } + if entries.failed { + return Err(TableError::BadBlock); + } + Ok(()) +} + +fn validate_data_block(block: &[u8]) -> Result<(), TableError> { + let mut entries = BlockEntries::new(block)?; + let mut entry_count = 0usize; + for (internal_key, _value) in entries.by_ref() { + entry_count += 1; + if entry_count > MAX_RECORDS_PER_DIRECTORY { + return Err(TableError::BadBlock); + } + if internal_key.len() < 8 { + return Err(TableError::BadBlock); + } + } + if entries.failed { + return Err(TableError::BadBlock); + } + Ok(()) +} + +fn emit_block_records( + block: &[u8], + emit: &mut impl FnMut(Record) -> bool, +) -> Result { + let mut entries = BlockEntries::new(block)?; + for (internal_key, value) in entries.by_ref() { let Some(split) = internal_key.len().checked_sub(8) else { return Err(TableError::BadBlock); }; let (user_key, trailer) = internal_key.split_at(split); let packed = read_u64_le(trailer, 0).ok_or(TableError::BadBlock)?; let value = match packed & 0xff { - KIND_VALUE => Some(value.to_vec()), + KIND_VALUE => Some(copy_bytes(value)?), KIND_DELETE => None, _ => continue, }; - emit(Record { - key: user_key.to_vec(), + if !emit(Record { + key: copy_bytes(user_key)?, sequence: packed >> 8, value, - }); + }) { + return Ok(false); + } } - Ok(()) + if entries.failed { + return Err(TableError::BadBlock); + } + Ok(true) } /// Return the decompressed contents of the block at `handle` (without its trailer). fn read_block(data: &[u8], handle: BlockHandle) -> Result, TableError> { + let table_data_end = data + .len() + .checked_sub(FOOTER_SIZE) + .ok_or(TableError::BadHandle)?; let end = handle .offset .checked_add(handle.size) .and_then(|end| end.checked_add(BLOCK_TRAILER_SIZE)) .ok_or(TableError::BadHandle)?; - if data.len() < end { + if end > table_data_end { return Err(TableError::BadHandle); } let contents_end = end - BLOCK_TRAILER_SIZE; @@ -131,12 +198,22 @@ fn read_block(data: &[u8], handle: BlockHandle) -> Result, TableError> { .ok_or(TableError::BadHandle)?; let compression = data[contents_end]; match compression { - COMPRESSION_NONE => Ok(raw.to_vec()), + COMPRESSION_NONE if raw.len() <= MAX_BLOCK_BYTES => copy_bytes(raw), + COMPRESSION_NONE => Err(TableError::BlockTooLarge), COMPRESSION_SNAPPY => Ok(snappy::decompress(raw, MAX_BLOCK_BYTES)?), other => Err(TableError::UnsupportedCompression(other)), } } +fn copy_bytes(bytes: &[u8]) -> Result, TableError> { + let mut output = Vec::new(); + output + .try_reserve_exact(bytes.len()) + .map_err(|_| TableError::AllocationFailed)?; + output.extend_from_slice(bytes); + Ok(output) +} + /// Iterator over the `(key, value)` entries of a decoded block, undoing prefix compression. struct BlockEntries<'a> { entries: &'a [u8], @@ -148,6 +225,9 @@ impl<'a> BlockEntries<'a> { fn new(block: &'a [u8]) -> Result { let restart_count = read_u32_le(block, block.len().saturating_sub(4)).ok_or(TableError::BadBlock)? as usize; + if restart_count == 0 || restart_count > MAX_RECORDS_PER_DIRECTORY { + return Err(TableError::BadBlock); + } let restarts_len = restart_count .checked_mul(4) .and_then(|len| len.checked_add(4)) @@ -156,6 +236,22 @@ impl<'a> BlockEntries<'a> { .len() .checked_sub(restarts_len) .ok_or(TableError::BadBlock)?; + let mut restart_position = entries_end; + let mut previous_restart = 0usize; + for index in 0..restart_count { + let restart = + usize::try_from(read_u32_le(block, restart_position).ok_or(TableError::BadBlock)?) + .map_err(|_| TableError::BadBlock)?; + if (index == 0 && restart != 0) + || restart < previous_restart + || restart > entries_end + || (entries_end > 0 && restart == entries_end) + { + return Err(TableError::BadBlock); + } + previous_restart = restart; + restart_position += 4; + } Ok(Self { entries: &block[..entries_end], key: Vec::new(), @@ -174,7 +270,12 @@ impl<'a> Iterator for BlockEntries<'a> { let parsed = parse_entry(self.entries, &self.key); match parsed { Some((key, value, rest)) => { - self.key.clone_from(&key); + self.key.clear(); + if self.key.try_reserve(key.len()).is_err() { + self.failed = true; + return None; + } + self.key.extend_from_slice(&key); self.entries = rest; Some((key, value)) } @@ -201,7 +302,9 @@ fn parse_entry<'a>( if shared > previous_key.len() || body.len() < payload_len { return None; } - let mut key = Vec::with_capacity(shared + non_shared); + let key_len = shared.checked_add(non_shared)?; + let mut key = Vec::new(); + key.try_reserve_exact(key_len).ok()?; key.extend_from_slice(&previous_key[..shared]); key.extend_from_slice(&body[..non_shared]); Some((key, &body[non_shared..payload_len], &body[payload_len..])) diff --git a/rust/src/providers/kimi/local_storage.rs b/rust/src/providers/kimi/local_storage.rs index 9314082b13..802f63aff7 100644 --- a/rust/src/providers/kimi/local_storage.rs +++ b/rust/src/providers/kimi/local_storage.rs @@ -12,18 +12,23 @@ use std::time::{SystemTime, UNIX_EPOCH}; use super::KimiRegion; use crate::browser::detection::BrowserDetector; use crate::browser::leveldb::local_storage::{ - LocalStorageEntry, local_storage_dir, read_local_storage_entries, + LocalStorageEntry, local_storage_dir, read_local_storage_entries_with_budget, }; +use crate::browser::leveldb::{LevelDbError, MAX_TOTAL_SCAN_BYTES}; use crate::codex_accounts::api::jwt_payload; const ACCESS_TOKEN_KEY: &str = "access_token"; +const MAX_PROFILE_DIRECTORIES: usize = 128; +const MAX_STORED_TOKEN_BYTES: usize = 32 * 1024; +const MAX_JWT_BYTES: usize = 16 * 1024; /// Current Kimi web access tokens stored by Chromium browsers for `region`, in browser and /// profile detection order, without duplicates. pub(super) fn local_storage_tokens(region: KimiRegion) -> Vec { - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0.0, |elapsed| elapsed.as_secs_f64()); + let Ok(elapsed) = SystemTime::now().duration_since(UNIX_EPOCH) else { + return Vec::new(); + }; + let now = elapsed.as_secs_f64(); tokens_from_profiles(&chromium_profile_dirs(), region.web_base_url(), now) } @@ -32,6 +37,7 @@ fn chromium_profile_dirs() -> Vec { .into_iter() .filter(|browser| browser.browser_type.is_chromium_based()) .flat_map(|browser| browser.profiles) + .take(MAX_PROFILE_DIRECTORIES) .map(|profile| profile.path) .collect() } @@ -39,19 +45,25 @@ fn chromium_profile_dirs() -> Vec { fn tokens_from_profiles(profiles: &[PathBuf], origin: &str, now_unix: f64) -> Vec { let mut seen = HashSet::new(); let mut tokens = Vec::new(); - for profile in profiles { + let mut remaining_bytes = MAX_TOTAL_SCAN_BYTES; + for profile in profiles.iter().take(MAX_PROFILE_DIRECTORIES) { let dir = local_storage_dir(profile); if !dir.is_dir() { continue; } - match read_local_storage_entries(&dir, origin) { - Ok(entries) => { + match read_local_storage_entries_with_budget(&dir, origin, remaining_bytes) { + Ok((entries, scanned_bytes)) => { + remaining_bytes = remaining_bytes.saturating_sub(scanned_bytes); for token in access_tokens(&entries, now_unix) { if seen.insert(token.clone()) { tokens.push(token); } } } + Err(LevelDbError::ResourceLimit) => { + tracing::debug!("Kimi local storage scan limit reached"); + break; + } Err(error) => tracing::debug!(%error, "Kimi local storage is not readable"), } } @@ -62,6 +74,7 @@ fn access_tokens(entries: &[LocalStorageEntry], now_unix: f64) -> impl Iterator< entries .iter() .filter(|entry| entry.key == ACCESS_TOKEN_KEY) + .filter(|entry| entry.value.len() <= MAX_STORED_TOKEN_BYTES) .map(|entry| normalized_value(&entry.value)) .filter(move |token| is_current_jwt(token, now_unix)) } @@ -72,7 +85,8 @@ fn normalized_value(value: &str) -> String { } fn is_current_jwt(token: &str, now_unix: f64) -> bool { - token.split('.').count() == 3 + token.len() <= MAX_JWT_BYTES + && token.split('.').count() == 3 && token .bytes() .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')) diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 3e9cd3e4ca..b3bcb15c5c 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -80,12 +80,22 @@ enum WebTokenSource { LocalStorage, } -#[derive(Debug, PartialEq, Eq)] +#[derive(PartialEq, Eq)] struct WebTokenCandidate { token: String, source: WebTokenSource, } +impl std::fmt::Debug for WebTokenCandidate { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WebTokenCandidate") + .field("token", &"[REDACTED]") + .field("source", &self.source) + .finish() + } +} + fn resolve_web_tokens(input: WebTokenInput) -> Vec { if let Some(header) = input.manual_header && let Ok(token) = KimiProvider::auth_token_from_cookie_header(header) From f37043a354678ca9726e383b29407d16c1085e28 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:01:56 +0700 Subject: [PATCH 133/259] Address thermo review --- rust/src/browser/detection.rs | 48 +++--- rust/src/browser/leveldb/local_storage.rs | 15 +- rust/src/browser/storage_discovery.rs | 34 ++-- rust/src/providers/minimax/local_storage.rs | 167 +++++++++++++++----- 4 files changed, 176 insertions(+), 88 deletions(-) diff --git a/rust/src/browser/detection.rs b/rust/src/browser/detection.rs index 0dc39f27d7..cac201ab5e 100755 --- a/rust/src/browser/detection.rs +++ b/rust/src/browser/detection.rs @@ -206,36 +206,26 @@ impl BrowserDetector { } /// Detect Chromium-based browser profiles - fn detect_chromium_profiles(user_data_dir: &PathBuf) -> Vec { - let mut profiles = Vec::new(); - - // Default profile - let default_path = user_data_dir.join("Default"); - if default_path.exists() { - profiles.push(BrowserProfile { - name: "Default".to_string(), - path: default_path, - is_default: true, - }); - } - - // Additional profiles (Profile 1, Profile 2, etc.) - if let Ok(entries) = std::fs::read_dir(user_data_dir) { - for entry in entries.flatten() { - let name = entry.file_name().to_string_lossy().to_string(); - if name.starts_with("Profile ") { - let path = entry.path(); - if path.is_dir() { - profiles.push(BrowserProfile { - name, - path, - is_default: false, - }); - } - } - } - } + pub(super) fn detect_chromium_profiles(user_data_dir: &Path) -> Vec { + let Ok(entries) = std::fs::read_dir(user_data_dir) else { + return Vec::new(); + }; + let mut profiles: Vec<_> = entries + .flatten() + .filter_map(|entry| { + let name = entry.file_name().into_string().ok()?; + let path = entry.path(); + let is_profile = + name == "Default" || name.starts_with("Profile ") || name.starts_with("user-"); + (is_profile && path.is_dir()).then(|| BrowserProfile { + is_default: name == "Default", + name, + path, + }) + }) + .collect(); + profiles.sort_by(|left, right| left.name.cmp(&right.name)); profiles } diff --git a/rust/src/browser/leveldb/local_storage.rs b/rust/src/browser/leveldb/local_storage.rs index a2e3d5236e..04ede1d75e 100644 --- a/rust/src/browser/leveldb/local_storage.rs +++ b/rust/src/browser/leveldb/local_storage.rs @@ -32,7 +32,20 @@ pub fn read_local_storage_entries( dir: &Path, origin: &str, ) -> Result, LevelDbError> { - Ok(decode_origin_entries(&read_entries(dir)?, origin)) + read_local_storage_entries_for_origins(dir, &[origin]) +} + +/// Read `localStorage` items for any of `origins` with one LevelDB scan, sorted by key within +/// each origin. A trailing slash on an origin is ignored. +pub fn read_local_storage_entries_for_origins( + dir: &Path, + origins: &[&str], +) -> Result, LevelDbError> { + let entries = read_entries(dir)?; + Ok(origins + .iter() + .flat_map(|origin| decode_origin_entries(&entries, origin)) + .collect()) } pub(super) fn decode_origin_entries(entries: &[Entry], origin: &str) -> Vec { diff --git a/rust/src/browser/storage_discovery.rs b/rust/src/browser/storage_discovery.rs index 49d3a63d53..ab6214e43b 100644 --- a/rust/src/browser/storage_discovery.rs +++ b/rust/src/browser/storage_discovery.rs @@ -55,11 +55,7 @@ pub fn discover(kind: StorageKind) -> Vec { .iter() .filter(|browser| browser.browser_type.is_chromium_based()) .flat_map(|browser| { - candidates_in_user_data_dir( - &browser.user_data_dir, - browser.browser_type.display_name(), - kind, - ) + candidates_in_profiles(&browser.profiles, browser.browser_type.display_name(), kind) }) .collect() } @@ -68,16 +64,26 @@ pub fn discover(kind: StorageKind) -> Vec { /// /// Only `Default`, `Profile *`, and `user-*` directories count as profiles; guest and system /// profiles are ignored. -pub fn candidates_in_user_data_dir( +#[cfg(test)] +fn candidates_in_user_data_dir( user_data_dir: &Path, label_prefix: &str, kind: StorageKind, ) -> Vec { - sorted_child_dirs(user_data_dir, is_profile_dir_name) - .into_iter() - .flat_map(|(name, profile_dir)| { - let label = format!("{label_prefix} {name}{}", kind.label_suffix()); - profile_store_paths(&profile_dir, kind) + let profiles = BrowserDetector::detect_chromium_profiles(user_data_dir); + candidates_in_profiles(&profiles, label_prefix, kind) +} + +fn candidates_in_profiles( + profiles: &[super::detection::BrowserProfile], + label_prefix: &str, + kind: StorageKind, +) -> Vec { + profiles + .iter() + .flat_map(|profile| { + let label = format!("{label_prefix} {}{}", profile.name, kind.label_suffix()); + profile_store_paths(&profile.path, kind) .into_iter() .map(move |path| StorageCandidate { label: label.clone(), @@ -106,17 +112,13 @@ fn profile_store_paths(profile_dir: &Path, kind: StorageKind) -> Vec { } fn existing(path: PathBuf) -> Vec { - if path.exists() { + if path.is_dir() { vec![path] } else { Vec::new() } } -fn is_profile_dir_name(name: &str) -> bool { - name == "Default" || name.starts_with("Profile ") || name.starts_with("user-") -} - /// Child directories of `dir` whose (non-hidden, valid UTF-8) name passes `keep`, sorted by name. /// A missing or unreadable `dir` yields nothing. fn sorted_child_dirs(dir: &Path, keep: impl Fn(&str) -> bool) -> Vec<(String, PathBuf)> { diff --git a/rust/src/providers/minimax/local_storage.rs b/rust/src/providers/minimax/local_storage.rs index 41ef8c8d29..3225be47e1 100755 --- a/rust/src/providers/minimax/local_storage.rs +++ b/rust/src/providers/minimax/local_storage.rs @@ -4,6 +4,10 @@ //! Storage directories come from `browser::storage_discovery` (every installed Chromium-family //! browser and profile: Local Storage, then Session Storage, then MiniMax IndexedDB). +use crate::browser::leveldb::local_storage::{ + LocalStorageEntry, read_local_storage_entries_for_origins, +}; +use crate::browser::leveldb::{self, Entry}; use crate::browser::storage_discovery::{self, StorageCandidate, StorageKind}; use serde::{Deserialize, Serialize}; use std::path::Path; @@ -52,6 +56,23 @@ const INDEXED_DB_ORIGIN_PREFIXES: &[&str] = &[ "https_www.minimaxi.com_", ]; +const MINIMAX_LOCAL_STORAGE_ORIGINS: &[&str] = &[ + "https://platform.minimax.io", + "https://www.minimax.io", + "https://minimax.io", + "https://platform.minimaxi.com", + "https://www.minimaxi.com", + "https://minimaxi.com", +]; + +const MINIMAX_SESSION_PATTERNS: &[&str] = &[ + "minimax_user", + "minimax_session", + "platform.minimaxi.com", + "mm_token", + "mm_user_info", +]; + /// Stores tried in order; a later store is read only when earlier ones yield no session. const STORAGE_ORDER: [StorageKind; 3] = [ StorageKind::LocalStorage, @@ -77,45 +98,44 @@ impl MiniMaxLocalStorageImporter { discover: impl Fn(StorageKind) -> Vec, ) -> Result { let mut found_any_store = false; + let mut last_error = None; for kind in STORAGE_ORDER { for candidate in discover(kind) { found_any_store = true; - if let Ok(session) = Self::extract_from_path(&candidate.path, &candidate.label) { - return Ok(session); + match Self::extract_from_path(&candidate.path, &candidate.label, kind) { + Ok(session) => return Ok(session), + Err(ImportError::StorageNotFound) => {} + Err(error) => last_error = Some(error), } } } Err(if found_any_store { - ImportError::StorageNotFound + last_error.unwrap_or(ImportError::StorageNotFound) } else { ImportError::BrowserNotFound }) } - /// Extract session from a storage directory - fn extract_from_path(path: &Path, source_label: &str) -> Result { - // Look for .ldb or .log files - let entries = - std::fs::read_dir(path).map_err(|e| ImportError::AccessDenied(e.to_string()))?; - - let mut minimax_data: Option = None; - - for entry in entries.flatten() { - let entry_path = entry.path(); - if let Some(ext) = entry_path.extension() - && (ext == "ldb" || ext == "log") - { - // Read file and search for MiniMax data - if let Ok(contents) = std::fs::read(&entry_path) { - // Search for MiniMax-related JSON in the binary content - if let Some(data) = Self::extract_minimax_json(&contents) { - minimax_data = Some(data); - break; - } - } + /// Extract a MiniMax session from one decoded Chromium storage directory. + fn extract_from_path( + path: &Path, + source_label: &str, + kind: StorageKind, + ) -> Result { + let minimax_data = match kind { + StorageKind::LocalStorage => { + let entries = + read_local_storage_entries_for_origins(path, MINIMAX_LOCAL_STORAGE_ORIGINS) + .map_err(|error| ImportError::AccessDenied(error.to_string()))?; + entries.iter().find_map(Self::extract_local_storage_json) } - } + StorageKind::SessionStorage | StorageKind::IndexedDb { .. } => { + let entries = leveldb::read_entries(path) + .map_err(|error| ImportError::AccessDenied(error.to_string()))?; + entries.iter().find_map(Self::extract_minimax_entry) + } + }; match minimax_data { Some(json) => Self::parse_session_from_json(&json, source_label), @@ -123,21 +143,38 @@ impl MiniMaxLocalStorageImporter { } } - /// Extract MiniMax JSON from binary localStorage data - fn extract_minimax_json(data: &[u8]) -> Option { - // Convert to string, handling binary data - let content = String::from_utf8_lossy(data); + fn extract_local_storage_json(entry: &LocalStorageEntry) -> Option { + if MINIMAX_SESSION_PATTERNS + .iter() + .any(|pattern| entry.key.contains(pattern)) + { + return serde_json::from_str(&entry.value) + .ok() + .or_else(|| Self::extract_minimax_json(entry.value.as_bytes())); + } - // Look for patterns that indicate MiniMax session data - let patterns = [ - "minimax_user", - "minimax_session", - "platform.minimaxi.com", - "mm_token", - "mm_user_info", - ]; + Self::extract_minimax_json(entry.value.as_bytes()) + } + + fn extract_minimax_entry(entry: &Entry) -> Option { + let key = String::from_utf8_lossy(&entry.key); + let value = String::from_utf8_lossy(&entry.value); + if MINIMAX_SESSION_PATTERNS + .iter() + .any(|pattern| key.contains(pattern)) + { + return serde_json::from_str(&value) + .ok() + .or_else(|| Self::extract_minimax_json(&entry.value)); + } - for pattern in patterns { + Self::extract_minimax_json(&entry.value).or_else(|| Self::extract_minimax_json(&entry.key)) + } + + /// Find MiniMax JSON embedded in a decoded LevelDB key or value. + fn extract_minimax_json(data: &[u8]) -> Option { + let content = String::from_utf8_lossy(data); + for pattern in MINIMAX_SESSION_PATTERNS { if let Some(parsed) = Self::extract_json_after_pattern(&content, pattern) { return Some(parsed); } @@ -162,8 +199,19 @@ impl MiniMaxLocalStorageImporter { fn matching_json_object_end(content: &str) -> Option { let mut depth = 0; + let mut in_string = false; + let mut escaped = false; for (i, c) in content.char_indices() { + if in_string { + match c { + '\\' if !escaped => escaped = true, + '"' if !escaped => in_string = false, + _ => escaped = false, + } + continue; + } match c { + '"' => in_string = true, '{' => depth += 1, '}' => { depth -= 1; @@ -285,16 +333,51 @@ mod tests { fn write_store(dir: &Path, token: Option<&str>) -> std::path::PathBuf { std::fs::create_dir_all(dir).unwrap(); - let body = match token { + let (key, value, local_key, local_value) = match token { Some(token) => { - format!("minimax_user{{\"access_token\":\"{token}\",\"user_id\":\"1\"}}") + let json = format!(r#"{{"access_token":"{token}","user_id":"1"}}"#); + let mut local_value = vec![1]; + local_value.extend_from_slice(json.as_bytes()); + ( + b"minimax_user".to_vec(), + json.as_bytes().to_vec(), + b"_https://platform.minimax.io\0\x01minimax_user".to_vec(), + local_value, + ) } - None => "unrelated".to_string(), + None => ( + b"unrelated".to_vec(), + b"unrelated".to_vec(), + b"_https://platform.minimax.io\0\x01unrelated".to_vec(), + b"\x01unrelated".to_vec(), + ), }; - std::fs::write(dir.join("000003.log"), body).unwrap(); + + let mut batch = Vec::new(); + batch.extend_from_slice(&1u64.to_le_bytes()); + batch.extend_from_slice(&2u32.to_le_bytes()); + for (key, value) in [(key, value), (local_key, local_value)] { + batch.push(1); // put + append_varint(&mut batch, key.len()); + batch.extend_from_slice(&key); + append_varint(&mut batch, value.len()); + batch.extend_from_slice(&value); + } + let mut log = vec![0; 4]; // the best-effort reader does not verify checksums + log.extend_from_slice(&u16::try_from(batch.len()).unwrap().to_le_bytes()); + log.push(1); // full physical record + log.extend_from_slice(&batch); + std::fs::write(dir.join("000003.log"), log).unwrap(); dir.to_path_buf() } + fn append_varint(output: &mut Vec, mut value: usize) { + while value >= 0x80 { + output.push(u8::try_from(value & 0x7f).unwrap() | 0x80); + value >>= 7; + } + output.push(u8::try_from(value).unwrap()); + } fn discover_from( local: Vec, session: Vec, From 195870a6f3e79003bc2c9dd948b6fb94bfdc4f0d Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:07:40 +0700 Subject: [PATCH 134/259] Address thermo review --- rust/src/cost_reporting_period.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/rust/src/cost_reporting_period.rs b/rust/src/cost_reporting_period.rs index c1721afaa3..8f0e41e3a4 100644 --- a/rust/src/cost_reporting_period.rs +++ b/rust/src/cost_reporting_period.rs @@ -241,7 +241,17 @@ impl Serialize for CostReportingPeriod { /// Settings files must keep loading, so an unreadable value reads as the default. impl<'de> Deserialize<'de> for CostReportingPeriod { fn deserialize>(deserializer: D) -> Result { - let raw = Option::::deserialize(deserializer)?; + #[derive(Deserialize)] + #[serde(untagged)] + enum SavedPeriod { + Text(String), + Unreadable(serde::de::IgnoredAny), + } + + let raw = match SavedPeriod::deserialize(deserializer)? { + SavedPeriod::Text(raw) => Some(raw), + SavedPeriod::Unreadable(_) => None, + }; Ok(Self::migrated(raw.as_deref(), None)) } } @@ -452,5 +462,7 @@ mod tests { assert_eq!(read("\"all\""), CostReportingPeriod::AllAvailable); assert_eq!(read("\"nonsense\""), CostReportingPeriod::Rolling(30)); assert_eq!(read("null"), CostReportingPeriod::Rolling(30)); + assert_eq!(read("42"), CostReportingPeriod::Rolling(30)); + assert_eq!(read(r#"{"a":1}"#), CostReportingPeriod::Rolling(30)); } } From 8f096d88bad5bd3fb11fe45dc20823cd3b2bcd61 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:11:22 +0700 Subject: [PATCH 135/259] Address thermo review --- rust/src/cli/cost.rs | 25 ++++++++++++++++++------- rust/src/cli/cost_period.rs | 7 +++++-- rust/src/cli/serve/data.rs | 31 ++++++++++++++++++------------- 3 files changed, 41 insertions(+), 22 deletions(-) diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 96d85a8f0a..97f328f1c6 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -100,11 +100,12 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { let group_by = CostGroupBy::from_arg(args.group_by.as_deref()); let use_color = !args.no_color && is_terminal(); + let settings = Settings::load(); let host_summary = args.remote.is_some() || args.summary_only; let period = resolve_period( args.days, args.period.as_deref(), - Settings::load().cost_reporting_period, + settings.cost_reporting_period, host_summary, )?; let days = window_days(period); @@ -221,7 +222,7 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { print_text_output(&results, use_color, period, group_by); } OutputFormat::Json => { - print_json_output(&results, args.pretty, period, days)?; + print_json_output(&results, args.pretty, period, days, &settings)?; } } @@ -485,8 +486,8 @@ fn build_json_payloads( results: &[CostResult], period: CostReportingPeriod, days: u32, + settings: &Settings, ) -> Vec { - let settings = Settings::load(); results .iter() .map(|r| { @@ -546,8 +547,9 @@ fn print_json_output( pretty: bool, period: CostReportingPeriod, days: u32, + settings: &Settings, ) -> anyhow::Result<()> { - let payloads = build_json_payloads(results, period, days); + let payloads = build_json_payloads(results, period, days, settings); let output = if pretty { serde_json::to_string_pretty(&payloads)? @@ -727,7 +729,9 @@ mod tests { supported: true, token_history: None, }]; - let payloads = build_json_payloads(&results, CostReportingPeriod::MonthToDate, 12); + let settings = Settings::default(); + let payloads = + build_json_payloads(&results, CostReportingPeriod::MonthToDate, 12, &settings); let payload = &payloads[0]; assert_eq!(payload["reportingPeriod"], "month-to-date"); assert_eq!(payload["historyLabel"], "Month to date"); @@ -738,7 +742,8 @@ mod tests { assert_eq!(payload["cost"]["total_usd"], 2.5); assert_eq!(payload["tokens"]["input"], 10); - let rolling = build_json_payloads(&results, CostReportingPeriod::Rolling(30), 30); + let rolling = + build_json_payloads(&results, CostReportingPeriod::Rolling(30), 30, &settings); assert_eq!(rolling[0]["reportingPeriod"], "rolling:30"); assert_eq!(rolling[0]["historyLabel"], "Last 30 days"); } @@ -757,7 +762,13 @@ mod tests { coverage: LocalHistoryCoverage::Complete, }), }]; - let payloads = build_json_payloads(&results, CostReportingPeriod::AllAvailable, 20_000); + let settings = Settings::default(); + let payloads = build_json_payloads( + &results, + CostReportingPeriod::AllAvailable, + 20_000, + &settings, + ); assert_eq!(payloads[0]["reportingPeriod"], "all"); assert_eq!(payloads[0]["historyLabel"], "All"); assert_eq!(payloads[0]["days_scanned"], 20_000); diff --git a/rust/src/cli/cost_period.rs b/rust/src/cli/cost_period.rs index 6335b5bc37..dee4c1274a 100644 --- a/rust/src/cli/cost_period.rs +++ b/rust/src/cli/cost_period.rs @@ -23,8 +23,11 @@ pub(super) fn resolve_period( return Ok(CostReportingPeriod::rolling(days)); } let period = match period { - Some(raw) => CostReportingPeriod::parse(raw) - .ok_or_else(|| anyhow::anyhow!("--period must be month-to-date or all"))?, + Some(raw) => match CostReportingPeriod::parse(raw) { + Some(CostReportingPeriod::MonthToDate) => CostReportingPeriod::MonthToDate, + Some(CostReportingPeriod::AllAvailable) => CostReportingPeriod::AllAvailable, + _ => anyhow::bail!("--period must be month-to-date or all"), + }, None => saved, }; if host_summary && period == CostReportingPeriod::AllAvailable { diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index b086745d90..5041815791 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -91,21 +91,26 @@ pub async fn cost_response(provider: Option<&str>) -> String { results.push(payload); continue; } - 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()), + let (supported, summary, daily) = match provider_id { + ProviderId::Codex => (true, scanner.scan_codex(), None), + ProviderId::Claude => { + let snapshot = scanner.scan_claude_chart_snapshot_with_cancel(None); + (true, snapshot.summary, Some(snapshot.daily_cost)) + } + ProviderId::Pi => (true, scanner.scan_pi(), None), + _ => (false, Default::default(), None), }; if supported { - // Daily spend history for the dashboard bar charts. The debounced - // helper reuses the cache the summary scan just warmed, so no - // second disk walk happens per request. All charts the most recent - // year; month to date charts exactly its own days. - let daily = daily_json(cost_scanner::get_daily_cost_history( - provider_id.cli_name(), - rolling_window_days(period), - )); + // Claude's snapshot derives the summary and chart rows in one + // transcript walk. Other providers use the shared daily-history + // path, capped at one year for All. + let daily = daily.unwrap_or_else(|| { + cost_scanner::get_daily_cost_history( + provider_id.cli_name(), + rolling_window_days(period), + ) + }); + let daily = daily_json(daily); let mut payload = json!({ "provider": provider_id.cli_name(), "supported": true, From bb22a6d066a99c5fb56cccb5c700515f5a24cd1a Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:24:18 +0700 Subject: [PATCH 136/259] Address thermo review --- rust/src/core/jsonl_scanner.rs | 40 +++++++++++++++++++++--- rust/src/core/jsonl_scanner/save_skip.rs | 27 +++++++--------- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index f5d2888a3e..4dbeba22df 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -169,6 +169,36 @@ impl CacheStamp { content_hash: hasher.finish(), } } + + /// Stamp the serialized cache while ignoring its scan timestamp. This + /// avoids allocating a second full-size JSON buffer just to normalize one + /// scalar before comparing cache payloads. + fn from_cache_payload(bytes: &[u8]) -> Option { + const FIELD: &[u8] = b"\"last_scan_unix_ms\":"; + let value_start = bytes + .windows(FIELD.len()) + .position(|window| window == FIELD)? + + FIELD.len(); + let mut value_end = value_start; + if bytes.get(value_end) == Some(&b'-') { + value_end += 1; + } + let digits_start = value_end; + while bytes.get(value_end).is_some_and(u8::is_ascii_digit) { + value_end += 1; + } + if value_end == digits_start { + return None; + } + + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + hasher.write(&bytes[..value_start]); + hasher.write(&bytes[value_end..]); + Some(Self { + byte_len: bytes.len() - (value_end - value_start), + content_hash: hasher.finish(), + }) + } } /// Terminal reason for a bounded Codex catch-up pause. @@ -249,10 +279,10 @@ pub struct CostUsageCache { /// and omitted from JSON so a stale reader cannot replace a newer cache. #[serde(skip)] pub(crate) loaded_stamp: Option>, - /// `last_scan_unix_ms` as decoded from disk, before any in-memory debounce - /// time from a skipped save. Baseline for the unchanged-payload check. + /// Stamp of the loaded cache payload with `last_scan_unix_ms` omitted. + /// This is separate from `loaded_stamp`, which still protects stale writes. #[serde(skip)] - pub(crate) loaded_last_scan_unix_ms: i64, + pub(crate) loaded_payload_stamp: Option, } /// Pricing evidence attached to one cached Codex request row. @@ -560,7 +590,7 @@ impl JsonlScanner { && let Ok(mut cache) = serde_json::from_str::(&contents) { let stamp = CacheStamp::from_bytes(contents.as_bytes()); - cache.loaded_last_scan_unix_ms = cache.last_scan_unix_ms; + cache.loaded_payload_stamp = CacheStamp::from_cache_payload(contents.as_bytes()); if let Some(scan_unix_ms) = save_skip::recorded_scan_time(&cache_path, &stamp) { cache.last_scan_unix_ms = scan_unix_ms; } @@ -901,7 +931,7 @@ impl JsonlScanner { }; if wrote { cache.loaded_stamp = Some(Some(CacheStamp::from_bytes(json.as_bytes()))); - cache.loaded_last_scan_unix_ms = cache.last_scan_unix_ms; + cache.loaded_payload_stamp = CacheStamp::from_cache_payload(json.as_bytes()); } // Best-effort temp cleanup (ignore errors — unique name avoids clashes). let _truncated_tmp = fs::File::create(&tmp_path).and_then(|f| f.set_len(0)); diff --git a/rust/src/core/jsonl_scanner/save_skip.rs b/rust/src/core/jsonl_scanner/save_skip.rs index 9c45b05b58..2c6231d87a 100644 --- a/rust/src/core/jsonl_scanner/save_skip.rs +++ b/rust/src/core/jsonl_scanner/save_skip.rs @@ -7,10 +7,13 @@ //! store is one JSON artifact, so the equivalent is to skip the write when the //! payload, excluding the scan timestamp, matches the decoded baseline. //! -//! Two pieces make that safe: +//! Three pieces make that safe: //! - Deterministic key order. `HashMap` iteration order differs between //! instances, so a reloaded cache would never re-encode to the same bytes. //! The `sorted_*` serializers give equal content equal bytes. +//! - A timestamp-free content stamp. The save path already has the encoded +//! JSON in memory, so the comparison hashes around the timestamp instead of +//! serializing the full cache into a second buffer. //! - An in-memory debounce mark. A skipped save records the new scan time for //! this process only, keyed by the on-disk stamp it applies to, so the //! scanner debounce still works. After a restart the older on-disk time @@ -78,25 +81,17 @@ pub(super) fn recorded_scan_time(cache_path: &Path, stamp: &CacheStamp) -> Optio /// the decoded on-disk baseline. `json` is the encoding of `cache` as-is. /// When it is unchanged the scan time is kept in memory and the caller skips /// the write. -pub(super) fn skip_unchanged_save( - cache_path: &Path, - cache: &mut CostUsageCache, - json: &str, -) -> bool { - let Some(Some(baseline)) = cache.loaded_stamp.clone() else { +pub(super) fn skip_unchanged_save(cache_path: &Path, cache: &CostUsageCache, json: &str) -> bool { + let Some(baseline) = cache.loaded_payload_stamp.as_ref() else { return false; }; - let unchanged = if cache.last_scan_unix_ms == cache.loaded_last_scan_unix_ms { - CacheStamp::from_bytes(json.as_bytes()) == baseline - } else { - let scan_unix_ms = cache.last_scan_unix_ms; - cache.last_scan_unix_ms = cache.loaded_last_scan_unix_ms; - let baseline_json = serde_json::to_vec(&*cache); - cache.last_scan_unix_ms = scan_unix_ms; - baseline_json.is_ok_and(|bytes| CacheStamp::from_bytes(&bytes) == baseline) + let Some(Some(stamp)) = cache.loaded_stamp.as_ref() else { + return false; }; + let unchanged = + CacheStamp::from_cache_payload(json.as_bytes()).is_some_and(|current| current == *baseline); if unchanged { - record_scan_time(cache_path, &baseline, cache.last_scan_unix_ms); + record_scan_time(cache_path, stamp, cache.last_scan_unix_ms); } unchanged } From 9170b79e964d888146cc1dade64e4edcb144c00a Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 19:25:58 +0700 Subject: [PATCH 137/259] Address thermo review --- .../src-tauri/src/usage_metric.rs | 1 + rust/src/cli/usage/render.rs | 30 ++++++++----------- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 095b181006..65a2408a15 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -63,6 +63,7 @@ fn same_window(left: &RateWindowSnapshot, right: &RateWindowSnapshot) -> bool { && left.window_minutes == right.window_minutes && left.resets_at == right.resets_at && left.reset_description == right.reset_description + && left.description_is_detail == right.description_is_detail && left.is_informational == right.is_informational } diff --git a/rust/src/cli/usage/render.rs b/rust/src/cli/usage/render.rs index 95f20b9d77..722a0c663a 100644 --- a/rust/src/cli/usage/render.rs +++ b/rust/src/cli/usage/render.rs @@ -326,15 +326,21 @@ fn append_window_line(lines: &mut Vec, label: &str, window: &RateWindow, let bar = render_progress_bar(window.used_percent, 20, use_color); let countdown = window.format_countdown(); + let reset_suffix = match (window.description_is_detail, countdown.as_deref()) { + (false, Some(countdown)) => format!(" (resets in {countdown})"), + _ => String::new(), + }; + lines.push(format!( + " {:<8} {} {} used{}", + format!("{}:", label), + bar, + format_percent(window.used_percent), + reset_suffix, + )); + if window.description_is_detail { // The description is a detail line (for example spend amounts), so // the reset stays on its own line and only exists with a real date. - lines.push(format!( - " {:<8} {} {} used", - format!("{}:", label), - bar, - format_percent(window.used_percent), - )); if let Some(countdown) = countdown { lines.push(format!(" resets in {countdown}")); } @@ -346,19 +352,7 @@ fn append_window_line(lines: &mut Vec, label: &str, window: &RateWindow, { lines.push(format!(" {detail}")); } - return; } - - let reset = countdown - .map(|c| format!(" (resets in {})", c)) - .unwrap_or_default(); - lines.push(format!( - " {:<8} {} {} used{}", - format!("{}:", label), - bar, - format_percent(window.used_percent), - reset - )); } fn append_secondary_window_line( From fce4584a6d0968c77ffabe15ddf4a68c59121094 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:13:06 +0700 Subject: [PATCH 138/259] Address thermo review --- rust/src/providers/bifrost/model_labels.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/src/providers/bifrost/model_labels.rs b/rust/src/providers/bifrost/model_labels.rs index 432361d0a4..8e8859a87a 100644 --- a/rust/src/providers/bifrost/model_labels.rs +++ b/rust/src/providers/bifrost/model_labels.rs @@ -25,7 +25,6 @@ pub(super) fn model_name(raw: &str) -> String { let name = strip_dotted_prefix(name, &VENDORS); let name = strip_revision(name); let name = strip_date(name); - let name = name.trim_end_matches([' ', '\t', '-']); if name.is_empty() { raw } else { name }.to_owned() } @@ -129,6 +128,7 @@ mod tests { ("gpt-4o", "gpt-4o"), ("global.qwen.qwen3-32b", "qwen3-32b"), ("usa.model", "usa.model"), + ("model-", "model-"), ("anthropic.", "anthropic."), ("-20250514", "-20250514"), ("model-v1:", "model-v1:"), From d472fc3fa1c1dd3a013fdc2d5c9440e332a8a4e6 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:46:42 +0700 Subject: [PATCH 139/259] Address thermo review --- .../src-tauri/src/commands/bridge.rs | 3 + .../src-tauri/src/commands/tests.rs | 10 +++ .../src/components/MenuCard.test.tsx | 74 +++++++++++++++++++ .../src/components/MenuCardDetails.tsx | 66 ++++++++--------- .../src/components/ProviderDisplayRow.tsx | 43 ++++++++++- .../sections/MenuBarMetricSection.test.tsx | 3 +- .../sections/MenuBarMetricSection.tsx | 6 +- .../providers/sections/UsageSection.tsx | 38 +++++++--- apps/desktop-tauri/src/types/bridge.ts | 1 + rust/src/core/display_detail.rs | 16 ++++ rust/src/core/usage_snapshot.rs | 18 +++++ 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/tests.rs | 7 ++ rust/src/locale/tr-TR.ftl | 3 + rust/src/locale/zh-CN.ftl | 3 + rust/src/locale/zh-TW.ftl | 3 + rust/src/providers/aixy/present.rs | 34 +++++---- rust/src/providers/aixy/tests.rs | 5 +- 21 files changed, 279 insertions(+), 66 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index a69be60f1c..4f3af6a351 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -206,6 +206,8 @@ pub struct ProviderDisplayProgressSnapshot { #[serde(rename_all = "camelCase")] pub struct ProviderDisplayDetailSnapshot { pub id: String, + #[serde(default)] + pub section_title: Option, pub title: String, pub value: String, pub secondary_value: Option, @@ -438,6 +440,7 @@ impl ProviderUsageSnapshot { .iter() .map(|detail| ProviderDisplayDetailSnapshot { id: detail.id().to_string(), + section_title: detail.section_title().map(ToOwned::to_owned), title: detail.title().to_string(), value: detail.value().to_string(), secondary_value: detail.secondary_value().map(ToOwned::to_owned), diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 94ac348f11..7383eb2503 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -1115,6 +1115,7 @@ fn provider_inventory_maps_to_the_bridge_without_token_ids() { .with_display_detail( ProviderDisplayDetail::new("credits", "Used this cycle", "12") .and_then(|row| row.with_secondary_value("Monthly refill: 100")) + .and_then(|row| row.with_section_title("Credit usage")) .and_then(|row| row.with_progress(12.0, 100.0)), ); let metadata = instantiate_provider(ProviderId::Grok).metadata().clone(); @@ -1129,6 +1130,15 @@ fn provider_inventory_maps_to_the_bridge_without_token_ids() { ); assert_eq!(snapshot.display_details.len(), 1); assert_eq!(snapshot.display_details[0].value, "12"); + assert_eq!( + snapshot.display_details[0].section_title.as_deref(), + Some("Credit usage") + ); + let snapshot_json = serde_json::to_value(&snapshot).unwrap(); + assert_eq!( + snapshot_json["displayDetails"][0]["sectionTitle"], + "Credit usage" + ); assert_eq!( snapshot.display_details[0].secondary_value.as_deref(), Some("Monthly refill: 100") diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index c8ec795546..02e9b99cad 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -601,6 +601,80 @@ describe("MenuCard", () => { expect(details.open).toBe(true); }); + it("renders provider detail rows once with separated secondary values and accessible progress", async () => { + const snapshot = provider(null); + snapshot.providerId = "aixy"; + snapshot.displayName = "Aixy"; + snapshot.displayDetails = [ + { + id: "budget-0", + sectionTitle: "Applicable budgets", + title: "Project · Monthly · Shared · Hard", + value: "$70.00 / $100.00 remaining", + secondaryValue: "$20.00 spent · $10.00 reserved", + progress: { used: 30, total: 100 }, + }, + ]; + + const { container } = renderCard(snapshot); + + await screen.findByRole("progressbar", { + name: "Project · Monthly · Shared · Hard progress", + }); + expect(container.querySelectorAll(".menu-card__provider-details")).toHaveLength(1); + expect(container.textContent).toContain( + "$70.00 / $100.00 remaining $20.00 spent · $10.00 reserved", + ); + expect( + screen.getByRole("progressbar", { + name: "Project · Monthly · Shared · Hard progress", + }), + ).toHaveAttribute("aria-valuenow", "30"); + }); + + it("hides provider display details in compact overview cards", () => { + const snapshot = provider(null); + snapshot.providerId = "aixy"; + snapshot.displayDetails = [ + { + id: "budget-0", + sectionTitle: "Applicable budgets", + title: "Project · Monthly · Shared · Hard", + value: "$70.00 / $100.00 remaining", + secondaryValue: "$20.00 spent · $10.00 reserved", + progress: { used: 30, total: 100 }, + }, + ]; + + const { container } = renderCard(snapshot, { compactOverview: true }); + + expect(container.querySelector(".menu-card__provider-details")).not.toBeInTheDocument(); + }); + + it("splits provider detail sections after 24 rows", async () => { + const snapshot = provider(null); + snapshot.providerId = "aixy"; + snapshot.displayDetails = Array.from({ length: 25 }, (_, index) => ({ + id: `budget-${index}`, + sectionTitle: "Applicable budgets", + title: `Budget ${index}`, + value: "$0.00 remaining", + secondaryValue: null, + progress: null, + })); + + const { container } = renderCard(snapshot); + + await screen.findAllByRole("heading", { name: "Applicable budgets" }); + const sections = container.querySelectorAll(".menu-card__provider-details"); + expect(sections).toHaveLength(2); + expect( + screen.getAllByRole("heading", { name: "Applicable budgets" }), + ).toHaveLength(2); + expect(sections[0].querySelectorAll(".menu-card__cost-line")).toHaveLength(24); + expect(sections[1].querySelectorAll(".menu-card__cost-line")).toHaveLength(1); + }); + it("places Claude accounts above metrics and the collapsed usage details", async () => { tauriMocks.claudeAccountsList.mockResolvedValue([ { id: "a", email: "a@example.com", organization: "Personal", isActive: true, isSaved: true }, diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 31634bd910..d446fa74c5 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -2,7 +2,6 @@ import { useState } from "react"; import type { CostSummaryDisplayStyle, DailyCostPoint, - ProviderDisplayDetail, PaceSnapshot, ProviderInventoryItem, ProviderChartData, @@ -22,6 +21,10 @@ import type { LocaleKey } from "../i18n/keys"; import { paceCategory } from "../surfaces/tray/paceCategory"; import { SimpleBarChart, StackedBarChart } from "./MiniBarChart"; import { InventoryItemRow } from "./InventoryRows"; +import { + groupProviderDisplayDetails, + ProviderDisplayRow, +} from "./ProviderDisplayRow"; import { QuotaWindowHistory } from "./QuotaWindowHistory"; import { getPaceBudget, type PaceBudget } from "../lib/paceBudget"; import PaceDetailsChart from "./PaceDetailsChart"; @@ -541,6 +544,9 @@ export default function MenuCardDetails({ ); const localCostHistory = chartData?.costHistory ?? []; const costStyle = display.costSummaryDisplayStyle ?? "detailed"; + const displayDetailGroups = groupProviderDisplayDetails( + provider.displayDetails ?? [], + ); const { hasMetrics, @@ -596,19 +602,28 @@ export default function MenuCardDetails({ )} {!provider.error && hasDisplayDetails && !compactOverview && ( -
- {provider.displayDetails?.map((detail, index) => ( - - ))} -
- )} - - {!provider.error && hasDisplayDetails && ( -
- {provider.displayDetails?.map((detail, index) => ( - - ))} -
+ displayDetailGroups.map((group) => ( +
+ {group.title && ( +
+ {group.title} +
+ )} + {group.rows.map((detail, index) => ( + + ))} +
+ )) )} {wayfinderUsage && !compactOverview && } @@ -786,26 +801,3 @@ export default function MenuCardDetails({
); } - -function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { - const progress = detail.progress; - const progressPercent = progress && Number.isFinite(progress.used) && Number.isFinite(progress.total) && progress.total > 0 - ? Math.max(0, Math.min(100, (progress.used / progress.total) * 100)) - : null; - - return ( -
-
- {detail.title}: {detail.value} - {detail.secondaryValue && ( - {detail.secondaryValue} - )} -
- {progressPercent != null && ( -
-
-
- )} -
- ); -} diff --git a/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx b/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx index 86982b66e4..b39fbe3ef3 100644 --- a/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx +++ b/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx @@ -1,5 +1,34 @@ import type { ProviderDisplayDetail } from "../types/bridge"; +const MAX_ROWS_PER_SECTION = 24; + +export interface ProviderDisplayDetailGroup { + id: number; + title: string | null; + rows: ProviderDisplayDetail[]; +} + +/** Group consecutive provider details and cap each rendered section. */ +export function groupProviderDisplayDetails( + details: ProviderDisplayDetail[], +): ProviderDisplayDetailGroup[] { + const groups: ProviderDisplayDetailGroup[] = []; + for (const detail of details) { + const title = detail.sectionTitle ?? null; + const current = groups[groups.length - 1]; + if ( + current && + current.title === title && + current.rows.length < MAX_ROWS_PER_SECTION + ) { + current.rows.push(detail); + } else { + groups.push({ id: groups.length, title, rows: [detail] }); + } + } + return groups; +} + /** * One transient provider detail line: "{title}: {value} [secondary]" * plus an optional clamped progress bar. @@ -34,11 +63,21 @@ export function ProviderDisplayRow({
{detail.title}: {detail.value} {detail.secondaryValue && secondaryClassName && ( - {detail.secondaryValue} + <> + {" "} + {detail.secondaryValue} + )}
{progressPercent != null && ( -
+
)} diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.test.tsx index 1c1315b9da..78494f71ec 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.test.tsx @@ -132,7 +132,7 @@ describe("MenuBarMetricSection", () => { render( key} onChange={vi.fn()} @@ -141,5 +141,6 @@ describe("MenuBarMetricSection", () => { const options = screen.getAllByRole("option").map((option) => option.textContent); expect(options).toEqual(["Automatic"]); + expect(screen.getByRole("combobox")).toHaveValue("automatic"); }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx index f8ff1eb89f..f0716229ce 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MenuBarMetricSection.tsx @@ -30,7 +30,9 @@ export function MenuBarMetricSection({ onChange, }: Props) { const [error, setError] = useState(null); - const selected = providerMetrics[provider.id] ?? "automatic"; + const selected = AUTOMATIC_ONLY_PROVIDERS.has(provider.id) + ? "automatic" + : providerMetrics[provider.id] ?? "automatic"; const options = metricOptions(provider, selected, t); const handleChange = (value: MetricPreference) => { @@ -83,7 +85,7 @@ function metricOptions( // Aixy's primary budget depends on which limits currently apply to the key, // so a fixed session/weekly lane would be misleading. Offer Automatic only. if (AUTOMATIC_ONLY_PROVIDERS.has(provider.id)) { - return withSelected(options, selected, t); + return options; } options.push({ value: "session", label: t("ProviderSessionLabel") }); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index b2d24dbdfc..e46e867a87 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -5,7 +5,10 @@ import type { RateWindowSnapshot, } from "../../../../types/bridge"; import { InventoryItemRow } from "../../../../components/InventoryRows"; -import { ProviderDisplayRow } from "../../../../components/ProviderDisplayRow"; +import { + groupProviderDisplayDetails, + ProviderDisplayRow, +} from "../../../../components/ProviderDisplayRow"; import type { LocaleKey } from "../../../../i18n/keys"; import { useFormattedResetTime } from "../../../../hooks/useFormattedResetTime"; import { isUsageItemVisible } from "../../../../lib/usageItemVisibility"; @@ -70,6 +73,7 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { const inventory = provider.inventory ?? []; const displayDetails = provider.displayDetails ?? []; + const displayDetailGroups = groupProviderDisplayDetails(displayDetails); if (bars.length === 0 && inventory.length === 0 && displayDetails.length === 0) { return null; } @@ -94,14 +98,28 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { lineClassName="provider-usage-inventory" /> ))} - {displayDetails.map((detail) => ( - + {displayDetailGroups.map((group) => ( +
+ {group.title && ( +
+ {group.title} +
+ )} + {group.rows.map((detail) => ( + + ))} +
))} ); @@ -163,4 +181,4 @@ function UsageBar({ )}
); -} \ No newline at end of file +} diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index fd8dd55109..233f7c0688 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -635,6 +635,7 @@ export interface ProviderDisplayProgress { /** Transient provider detail row; it is display-only and never quota math. */ export interface ProviderDisplayDetail { id: string; + sectionTitle: string | null; title: string; value: string; secondaryValue: string | null; diff --git a/rust/src/core/display_detail.rs b/rust/src/core/display_detail.rs index 1c80831aa3..cd5720d860 100644 --- a/rust/src/core/display_detail.rs +++ b/rust/src/core/display_detail.rs @@ -18,6 +18,7 @@ use crate::core::ProviderFetchResult; #[derive(Debug, Clone, PartialEq)] pub struct ProviderDisplayDetail { id: String, + section_title: Option, title: String, value: String, secondary_value: Option, @@ -40,6 +41,7 @@ impl ProviderDisplayDetail { ) -> Option { let row = Self { id: id.into(), + section_title: None, title: title.into(), value: value.into(), secondary_value: None, @@ -51,6 +53,16 @@ impl ProviderDisplayDetail { valid.then_some(row) } + /// Attach an optional display section title; rejects invalid text. + pub fn with_section_title(mut self, title: impl Into) -> Option { + let title = title.into(); + if !is_display_shape(&title, 128) { + return None; + } + self.section_title = Some(title); + Some(self) + } + /// Attach a validated secondary value; rejects invalid text. pub fn with_secondary_value(mut self, value: impl Into) -> Option { let value = value.into(); @@ -74,6 +86,10 @@ impl ProviderDisplayDetail { &self.id } + pub fn section_title(&self) -> Option<&str> { + self.section_title.as_deref() + } + pub fn title(&self) -> &str { &self.title } diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index ebbdcb8c20..3014476745 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -758,6 +758,24 @@ mod tests { assert_eq!(details[0].value(), "12"); } + #[test] + fn display_details_validate_optional_section_titles() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let row = ProviderDisplayDetail::new("budget", "Project", "$10 remaining") + .and_then(|row| row.with_section_title("Applicable budgets")); + let result = ProviderFetchResult::new(usage, "api").with_display_detail(row); + + assert_eq!( + result.display_details()[0].section_title(), + Some("Applicable budgets") + ); + assert!( + ProviderDisplayDetail::new("budget", "Project", "$10") + .and_then(|row| row.with_section_title("\n")) + .is_none() + ); + } + #[test] fn cost_snapshot_ignores_non_finite_values() { let cost = CostSnapshot::new(f64::NAN, "USD", "Monthly").with_limit(f64::INFINITY); diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index bae2c814c3..9d51ff351a 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -80,6 +80,9 @@ WayfinderGatewayTitle = Gateway de Wayfinder WayfinderGatewayLabel = URL del gateway WayfinderGatewayHelp = Usa HTTP solo para localhost o direcciones de loopback. HTTPS permite hosts remotos. WayfinderGatewayStatus = Gateway +AixyGatewayTitle = Gateway de Aixy +AixyGatewayLabel = URL base +AixyGatewayHelp = Déjalo vacío para usar el gateway alojado de Aixy o introduce la URL base de uno propio. HTTP solo se permite para localhost, redes privadas y hosts .local. WayfinderModels = Modelos WayfinderRequests = Solicitudes WayfinderTokens = Tokens diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 92618f72f7..2b1651507d 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -80,6 +80,9 @@ WayfinderGatewayTitle = Wayfinder ゲートウェイ WayfinderGatewayLabel = ゲートウェイ URL WayfinderGatewayHelp = HTTP は localhost またはループバックアドレスでのみ使用できます。リモートホストには HTTPS を使用します。 WayfinderGatewayStatus = ゲートウェイ +AixyGatewayTitle = Aixy ゲートウェイ +AixyGatewayLabel = ベース URL +AixyGatewayHelp = ホスト型 Aixy ゲートウェイを使う場合は空欄にするか、自分でホストするゲートウェイのベース URL を入力してください。HTTP は localhost、プライベートネットワーク、.local ホストでのみ使用できます。 WayfinderModels = モデル WayfinderRequests = リクエスト WayfinderTokens = トークン diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 7f59c6fc29..6988fded06 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -80,6 +80,9 @@ WayfinderGatewayTitle = Wayfinder 게이트웨이 WayfinderGatewayLabel = 게이트웨이 URL WayfinderGatewayHelp = HTTP는 localhost 또는 루프백 주소에서만 사용하세요. 원격 호스트에는 HTTPS를 사용합니다. WayfinderGatewayStatus = 게이트웨이 +AixyGatewayTitle = Aixy 게이트웨이 +AixyGatewayLabel = 기본 URL +AixyGatewayHelp = 호스팅된 Aixy 게이트웨이를 사용하려면 비워 두거나 자체 호스팅 게이트웨이의 기본 URL을 입력하세요. HTTP는 localhost, 사설 네트워크 및 .local 호스트에서만 허용됩니다. WayfinderModels = 모델 WayfinderRequests = 요청 WayfinderTokens = 토큰 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index 855f7a9eb6..afc216937e 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -64,6 +64,9 @@ WayfinderGatewayTitle = Шлюз Wayfinder WayfinderGatewayLabel = URL-адрес шлюза WayfinderGatewayHelp = Используйте HTTP только для локальных адресов или адресов обратной связи. HTTPS разрешен для удаленных хостов. WayfinderGatewayStatus = шлюз +AixyGatewayTitle = Шлюз Aixy +AixyGatewayLabel = Базовый URL +AixyGatewayHelp = Оставьте пустым для размещенного шлюза Aixy или укажите базовый URL собственного шлюза. HTTP разрешен только для localhost, частных сетей и хостов .local. WayfinderModels = Модели WayfinderRequests = Запросы WayfinderTokens = Токены diff --git a/rust/src/locale/tests.rs b/rust/src/locale/tests.rs index 9daf93d809..d11c4241ae 100644 --- a/rust/src/locale/tests.rs +++ b/rust/src/locale/tests.rs @@ -397,8 +397,15 @@ fn test_english_is_complete_and_other_languages_can_fallback() { .map(|(locale, resource)| (locale, resource_key_names(resource))) .collect(); let locale_key_names: HashSet<&str> = LocaleKey::ALL.iter().map(|(_, name)| *name).collect(); + let aixy_gateway_keys = ["AixyGatewayTitle", "AixyGatewayLabel", "AixyGatewayHelp"]; for (locale, keys) in &resource_keys { + for name in aixy_gateway_keys { + assert!( + keys.contains(name), + "missing Aixy gateway Fluent key {name} in {locale}" + ); + } for name in keys { assert!( locale_key_names.contains(name), diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index ca4fabe791..fa2f864995 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -86,6 +86,9 @@ WayfinderGatewayTitle = Wayfinder Ağ Geçidi WayfinderGatewayLabel = Ağ Geçidi URL'si WayfinderGatewayHelp = HTTP'yi yalnızca localhost veya geri döngü adresleri için kullanın. Uzak sunucular için HTTPS kullanılabilir. WayfinderGatewayStatus = Ağ Geçidi +AixyGatewayTitle = Aixy Ağ Geçidi +AixyGatewayLabel = Temel URL +AixyGatewayHelp = Barındırılan Aixy ağ geçidini kullanmak için boş bırakın veya kendi ağ geçidinizin temel URL'sini girin. HTTP yalnızca localhost, özel ağlar ve .local ana bilgisayarları için kullanılabilir. WayfinderModels = Modeller WayfinderRequests = İstekler WayfinderTokens = Tokenlar diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 98dd870db8..f8441f203d 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -80,6 +80,9 @@ WayfinderGatewayTitle = Wayfinder 网关 WayfinderGatewayLabel = 网关 URL WayfinderGatewayHelp = HTTP 仅可用于 localhost 或回环地址。远程主机请使用 HTTPS。 WayfinderGatewayStatus = 网关 +AixyGatewayTitle = Aixy 网关 +AixyGatewayLabel = 基础 URL +AixyGatewayHelp = 留空以使用托管的 Aixy 网关，或输入自托管网关的基础 URL。HTTP 仅允许用于 localhost、专用网络和 .local 主机。 WayfinderModels = 模型 WayfinderRequests = 请求 WayfinderTokens = 令牌 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index f893d5b423..e7af97178b 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -80,6 +80,9 @@ WayfinderGatewayTitle = Wayfinder 閘道 WayfinderGatewayLabel = 閘道 URL WayfinderGatewayHelp = HTTP 僅可用於 localhost 或迴路位址。遠端主機請使用 HTTPS。 WayfinderGatewayStatus = 閘道 +AixyGatewayTitle = Aixy 閘道 +AixyGatewayLabel = 基礎 URL +AixyGatewayHelp = 留空以使用託管的 Aixy 閘道，或輸入自架閘道的基礎 URL。HTTP 僅允許用於 localhost、私人網路和 .local 主機。 WayfinderModels = 模型 WayfinderRequests = 請求 WayfinderTokens = 權杖 diff --git a/rust/src/providers/aixy/present.rs b/rust/src/providers/aixy/present.rs index be49923ef9..87f5669478 100644 --- a/rust/src/providers/aixy/present.rs +++ b/rust/src/providers/aixy/present.rs @@ -81,22 +81,24 @@ fn rate_window(budget: &Budget) -> RateWindow { fn details(usage: &KeyUsage) -> Vec> { let mut rows = vec![ - ProviderDisplayDetail::new("key", "Key", &usage.key_label), - ProviderDisplayDetail::new("project", "Project", &usage.project_label), + ProviderDisplayDetail::new("key", "Key", &usage.key_label) + .and_then(|row| row.with_section_title("Aixy key")), + ProviderDisplayDetail::new("project", "Project", &usage.project_label) + .and_then(|row| row.with_section_title("Aixy key")), ProviderDisplayDetail::new( "observed", "Observed", usage.as_of.to_rfc3339_opts(SecondsFormat::Millis, true), - ), + ) + .and_then(|row| row.with_section_title("Aixy key")), ]; rows.extend(usage.budgets.iter().enumerate().map(budget_row)); match &usage.totals { Some(totals) => rows.extend(totals_rows(totals)), - None => rows.push(ProviderDisplayDetail::new( - "usage-7d", - "Usage (last 7 days)", - "Unavailable", - )), + None => rows.push( + ProviderDisplayDetail::new("usage-7d", "Usage (last 7 days)", "Unavailable") + .and_then(|row| row.with_section_title("Last 7 days · this key")), + ), } rows } @@ -105,7 +107,8 @@ fn budget_row((index, budget): (usize, &Budget)) -> Option Option Vec> { @@ -127,19 +131,22 @@ fn totals_rows(totals: &Totals) -> Vec> { "requests-7d", "Requests (last 7 days)", count(totals.requests), - ), + ) + .and_then(|row| row.with_section_title("Last 7 days · this key")), ProviderDisplayDetail::new( "tokens-7d", "Tokens (last 7 days)", count(totals.total_tokens), - ), + ) + .and_then(|row| row.with_section_title("Last 7 days · this key")), ProviderDisplayDetail::new( "spend-7d", "Attributed spend (last 7 days)", totals .spend_usd .map_or_else(|| "Unavailable".to_owned(), usd), - ), + ) + .and_then(|row| row.with_section_title("Last 7 days · this key")), ProviderDisplayDetail::new( "coverage-7d", "Cost coverage (last 7 days)", @@ -151,7 +158,8 @@ fn totals_rows(totals: &Totals) -> Vec> { ) .and_then(|row| { row.with_secondary_value(format!("{} partial", count(totals.partial_requests))) - }), + }) + .and_then(|row| row.with_section_title("Last 7 days · this key")), ] } diff --git a/rust/src/providers/aixy/tests.rs b/rust/src/providers/aixy/tests.rs index 72175225b8..f0b6f6ecb9 100644 --- a/rust/src/providers/aixy/tests.rs +++ b/rust/src/providers/aixy/tests.rs @@ -85,12 +85,14 @@ fn key_usage_keeps_overlapping_budgets_and_reservations_separate() { assert_eq!(result.source_label, "api"); assert_eq!(detail(&result, "key").value(), "Developer CLI"); + assert_eq!(detail(&result, "key").section_title(), Some("Aixy key")); assert_eq!(detail(&result, "project").value(), "Engineering"); assert_eq!( detail(&result, "observed").value(), "2026-09-24T12:00:00.000Z" ); let hard = detail(&result, "budget-0"); + assert_eq!(hard.section_title(), Some("Applicable budgets")); assert_eq!(hard.title(), "Project · Monthly · Shared · Hard"); assert_eq!(hard.value(), "$70.00 / $100.00 remaining"); assert_eq!( @@ -106,6 +108,7 @@ fn key_usage_keeps_overlapping_budgets_and_reservations_separate() { assert_eq!(detail(&result, "tokens-7d").value(), "1,200"); assert_eq!(detail(&result, "spend-7d").value(), "$1.25"); let coverage = detail(&result, "coverage-7d"); + assert_eq!(coverage.section_title(), Some("Last 7 days · this key")); assert_eq!(coverage.value(), "10 / 12 requests"); assert_eq!(coverage.secondary_value(), Some("2 partial")); } @@ -337,11 +340,11 @@ fn usage_url_supports_every_documented_base_form() { ), ("aixy.example.com", "https://aixy.example.com/v1/usage"), ("http://localhost:8080", "http://localhost:8080/v1/usage"), - ("http://10.1.2.3:8080/v1", "http://10.1.2.3:8080/v1/usage"), ( "http://gateway.local/team", "http://gateway.local/team/v1/usage", ), + ("http://10.1.2.3:8080/v1", "http://10.1.2.3:8080/v1/usage"), ] { assert_eq!(usage_url(base).unwrap().as_str(), expected, "{base}"); } From 4001969fe9e3439eacc8a4a96ba3a250397175a2 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:51:59 +0700 Subject: [PATCH 140/259] Address thermo review --- .../src/components/MenuCard.test.tsx | 23 +++++++++++ .../src/components/MenuCardDetails.tsx | 36 +++++++++++------- .../src/components/ProviderDisplayRow.tsx | 9 ++++- apps/desktop-tauri/src/i18n/keys.ts | 2 + .../src/lib/providerLabels.test.ts | 38 +++++++++++++++++++ apps/desktop-tauri/src/lib/providerLabels.ts | 29 ++++++++++++++ .../providers/sections/CostSection.test.tsx | 30 +++++++++++++++ .../providers/sections/CostSection.tsx | 7 ++++ .../providers/sections/UsageSection.tsx | 3 +- rust/src/locale.rs | 2 + rust/src/locale/en-US.ftl | 2 + rust/src/locale/es-MX.ftl | 2 + rust/src/locale/ja-JP.ftl | 2 + rust/src/locale/ko-KR.ftl | 2 + rust/src/locale/ru-RU.ftl | 2 + rust/src/locale/tests.rs | 2 + rust/src/locale/tr-TR.ftl | 2 + rust/src/locale/zh-CN.ftl | 2 + rust/src/locale/zh-TW.ftl | 2 + 19 files changed, 180 insertions(+), 17 deletions(-) create mode 100644 apps/desktop-tauri/src/lib/providerLabels.test.ts create mode 100644 apps/desktop-tauri/src/lib/providerLabels.ts create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.test.tsx diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index c8ec795546..2f85bd2176 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -119,6 +119,8 @@ describe("MenuCard", () => { buildBundle({ ActionCopyError: "Copy error", ApiSpendTitle: "API spend", + AtlasCloudAvailableBalance: "Available balance", + AtlasCloudBalance: "Atlas Cloud balance", DetailPaceRunsOutIn: "Runs out in", PanelEstimatedFromLocalLogs: "Estimated from local logs", PanelLeftSuffix: "left", @@ -601,6 +603,27 @@ describe("MenuCard", () => { expect(details.open).toBe(true); }); + it("renders provider display details once and hides them in compact overview", async () => { + const snapshot = provider(null); + snapshot.displayDetails = [ + { + id: "atlascloud-available", + title: "Available balance", + value: "$95.50", + secondaryValue: null, + progress: null, + }, + ]; + + const detailed = renderCard(snapshot); + expect(await screen.findByText("Available balance: $95.50")).toBeInTheDocument(); + expect(screen.getAllByText("Available balance: $95.50")).toHaveLength(1); + detailed.unmount(); + + renderCard(snapshot, { compactOverview: true }); + expect(screen.queryByText("Available balance: $95.50")).not.toBeInTheDocument(); + }); + it("places Claude accounts above metrics and the collapsed usage details", async () => { tauriMocks.claudeAccountsList.mockResolvedValue([ { id: "a", email: "a@example.com", organization: "Personal", isActive: true, isSaved: true }, diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 31634bd910..f30ece6a76 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -24,6 +24,10 @@ import { SimpleBarChart, StackedBarChart } from "./MiniBarChart"; import { InventoryItemRow } from "./InventoryRows"; import { QuotaWindowHistory } from "./QuotaWindowHistory"; import { getPaceBudget, type PaceBudget } from "../lib/paceBudget"; +import { + providerCostPeriodTitle, + providerDisplayDetailTitle, +} from "../lib/providerLabels"; import PaceDetailsChart from "./PaceDetailsChart"; /** Format a reserve description from raw pace data at render time. */ @@ -541,6 +545,11 @@ export default function MenuCardDetails({ ); const localCostHistory = chartData?.costHistory ?? []; const costStyle = display.costSummaryDisplayStyle ?? "detailed"; + const costPeriod = providerCostPeriodTitle( + provider.providerId, + provider.cost?.period ?? "", + t, + ); const { hasMetrics, @@ -598,15 +607,7 @@ export default function MenuCardDetails({ {!provider.error && hasDisplayDetails && !compactOverview && (
{provider.displayDetails?.map((detail, index) => ( - - ))} -
- )} - - {!provider.error && hasDisplayDetails && ( -
- {provider.displayDetails?.map((detail, index) => ( - + ))}
)} @@ -621,8 +622,8 @@ export default function MenuCardDetails({ {provider.cost.alwaysVisible === true && (provider.cost.limit ?? 0) <= 0 ? t("ApiSpendTitle") : provider.cost.balance != null && provider.cost.limit == null - ? provider.cost.period || t("CreditsLabel") - : `${t("DetailCostTitle")} — ${provider.cost.period}`} + ? costPeriod || t("CreditsLabel") + : `${t("DetailCostTitle")} — ${costPeriod}`}
{provider.cost.balance != null && provider.cost.limit == null ? (
@@ -787,7 +788,14 @@ export default function MenuCardDetails({ ); } -function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { +function DisplayDetailRow({ + detail, + t, +}: { + detail: ProviderDisplayDetail; + t: (key: LocaleKey) => string; +}) { + const title = providerDisplayDetailTitle(detail, t); const progress = detail.progress; const progressPercent = progress && Number.isFinite(progress.used) && Number.isFinite(progress.total) && progress.total > 0 ? Math.max(0, Math.min(100, (progress.used / progress.total) * 100)) @@ -796,13 +804,13 @@ function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { return (
- {detail.title}: {detail.value} + {title}: {detail.value} {detail.secondaryValue && ( {detail.secondaryValue} )}
{progressPercent != null && ( -
+
)} diff --git a/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx b/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx index 86982b66e4..ddc3535d37 100644 --- a/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx +++ b/apps/desktop-tauri/src/components/ProviderDisplayRow.tsx @@ -1,4 +1,6 @@ import type { ProviderDisplayDetail } from "../types/bridge"; +import type { LocaleKey } from "../i18n/keys"; +import { providerDisplayDetailTitle } from "../lib/providerLabels"; /** * One transient provider detail line: "{title}: {value} [secondary]" @@ -13,13 +15,16 @@ export function ProviderDisplayRow({ secondaryClassName, trackClassName, fillClassName, + t, }: { detail: ProviderDisplayDetail; lineClassName: string; secondaryClassName?: string; trackClassName: string; fillClassName: string; + t: (key: LocaleKey) => string; }) { + const title = providerDisplayDetailTitle(detail, t); const progress = detail.progress; const progressPercent = progress && @@ -32,13 +37,13 @@ export function ProviderDisplayRow({ return (
- {detail.title}: {detail.value} + {title}: {detail.value} {detail.secondaryValue && secondaryClassName && ( {detail.secondaryValue} )}
{progressPercent != null && ( -
+
)} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 1c07040562..f066db2e92 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -76,6 +76,8 @@ export const ALL_LOCALE_KEYS = [ "ProviderDisabled", "ProviderInfo", "ProviderUsage", + "AtlasCloudAvailableBalance", + "AtlasCloudBalance", "PanelUsageDetails", "AuthType", "DataSource", diff --git a/apps/desktop-tauri/src/lib/providerLabels.test.ts b/apps/desktop-tauri/src/lib/providerLabels.test.ts new file mode 100644 index 0000000000..80a84346b5 --- /dev/null +++ b/apps/desktop-tauri/src/lib/providerLabels.test.ts @@ -0,0 +1,38 @@ +import { describe, expect, it } from "vitest"; +import type { LocaleKey } from "../i18n/keys"; +import type { ProviderDisplayDetail } from "../types/bridge"; +import { + providerCostPeriodTitle, + providerDisplayDetailTitle, +} from "./providerLabels"; + +const translate = (key: LocaleKey) => `translated:${key}`; + +describe("provider labels", () => { + it("localizes Atlas Cloud display labels and preserves other provider titles", () => { + const atlasBalance: ProviderDisplayDetail = { + id: "atlascloud-available", + title: "Available balance", + value: "$95.50", + secondaryValue: null, + progress: null, + }; + const other: ProviderDisplayDetail = { + ...atlasBalance, + id: "other", + title: "Credits", + }; + + expect(providerDisplayDetailTitle(atlasBalance, translate)).toBe( + "translated:AtlasCloudAvailableBalance", + ); + expect(providerDisplayDetailTitle(other, translate)).toBe("Credits"); + }); + + it("localizes Atlas Cloud balance period and preserves other periods", () => { + expect( + providerCostPeriodTitle("atlascloud", "Atlas Cloud balance", translate), + ).toBe("translated:AtlasCloudBalance"); + expect(providerCostPeriodTitle("other", "This month", translate)).toBe("This month"); + }); +}); diff --git a/apps/desktop-tauri/src/lib/providerLabels.ts b/apps/desktop-tauri/src/lib/providerLabels.ts new file mode 100644 index 0000000000..7dd97d6778 --- /dev/null +++ b/apps/desktop-tauri/src/lib/providerLabels.ts @@ -0,0 +1,29 @@ +import type { LocaleKey } from "../i18n/keys"; +import type { ProviderDisplayDetail } from "../types/bridge"; + +type Translate = (key: LocaleKey) => string; + +const DISPLAY_DETAIL_TITLE_KEYS: ReadonlyMap = new Map([ + ["atlascloud-available", "AtlasCloudAvailableBalance"], +]); + +const COST_PERIOD_KEYS: ReadonlyMap = new Map([ + ["atlascloud:Atlas Cloud balance", "AtlasCloudBalance"], +]); + +export function providerDisplayDetailTitle( + detail: ProviderDisplayDetail, + t: Translate, +): string { + const key = DISPLAY_DETAIL_TITLE_KEYS.get(detail.id); + return key ? t(key) : detail.title; +} + +export function providerCostPeriodTitle( + providerId: string, + period: string, + t: Translate, +): string { + const key = COST_PERIOD_KEYS.get(`${providerId}:${period}`); + return key ? t(key) : period; +} diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.test.tsx new file mode 100644 index 0000000000..36189a47a8 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.test.tsx @@ -0,0 +1,30 @@ +import { describe, expect, it } from "vitest"; +import { render, screen } from "@testing-library/react"; +import type { CostSnapshotBridge } from "../../../../types/bridge"; +import type { LocaleKey } from "../../../../i18n/keys"; +import { CostSection } from "./CostSection"; + +describe("CostSection", () => { + it("does not present the zero usage carrier as spend for balance-only providers", () => { + const cost: CostSnapshotBridge = { + used: 0, + limit: null, + remaining: null, + currencyCode: "USD", + period: "Atlas Cloud balance", + resetsAt: null, + formattedUsed: "$0.00", + formattedLimit: null, + balance: 95.5, + formattedBalance: "$95.50", + }; + + const { container } = render( + key} />, + ); + + expect(container.firstChild).toBeNull(); + expect(screen.queryByText("DetailCostUsed")).toBeNull(); + expect(screen.queryByText("$0.00")).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.tsx index e1a078f29a..8d6e89b9c3 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CostSection.tsx @@ -13,6 +13,13 @@ interface Props { export function CostSection({ cost, t }: Props) { if (!cost) return null; + const balanceOnly = + cost.balance != null && + cost.limit == null && + cost.used === 0 && + (cost.daily?.length ?? 0) === 0; + if (balanceOnly) return null; + const rows: { label: string; value: string | null }[] = [ { label: t("DetailCostUsed"), value: cost.formattedUsed }, { label: t("DetailCostLimit"), value: cost.formattedLimit }, diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index b2d24dbdfc..8026099d83 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -98,6 +98,7 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { ); -} \ No newline at end of file +} diff --git a/rust/src/locale.rs b/rust/src/locale.rs index 276a31160e..41ec464786 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -266,6 +266,8 @@ locale_keys! { ProviderDisabled, ProviderInfo, ProviderUsage, + AtlasCloudAvailableBalance, + AtlasCloudBalance, PanelUsageDetails, AuthType, DataSource, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 53b4453f6d..9b6db82324 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -69,6 +69,8 @@ ProviderEnabled = Enabled ProviderDisabled = Disabled ProviderInfo = Info ProviderUsage = Usage +AtlasCloudAvailableBalance = Available balance +AtlasCloudBalance = Atlas Cloud balance PanelUsageDetails = Usage details AuthType = Authentication DataSource = Data Source diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index bae2c814c3..c8d79fc7a5 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -61,6 +61,8 @@ ProviderEnabled = Habilitado ProviderDisabled = Deshabilitado ProviderInfo = Información ProviderUsage = Uso +AtlasCloudAvailableBalance = Saldo disponible +AtlasCloudBalance = Saldo de Atlas Cloud AuthType = Autenticación DataSource = Fuente de datos ProviderNotDetected = no detectado diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 92618f72f7..f523f1a61d 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -61,6 +61,8 @@ ProviderEnabled = 有効 ProviderDisabled = 無効 ProviderInfo = 情報 ProviderUsage = 使用量 +AtlasCloudAvailableBalance = 利用可能残高 +AtlasCloudBalance = Atlas Cloud の残高 AuthType = 認証 DataSource = データソース ProviderNotDetected = 未検出 diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 7f59c6fc29..9b80f771f7 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -61,6 +61,8 @@ ProviderEnabled = 활성화됨 ProviderDisabled = 비활성화됨 ProviderInfo = 정보 ProviderUsage = 사용량 +AtlasCloudAvailableBalance = 사용 가능 잔액 +AtlasCloudBalance = Atlas Cloud 잔액 AuthType = 인증 DataSource = 데이터 소스 ProviderNotDetected = 감지되지 않음 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index 855f7a9eb6..de7c6a1f57 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -45,6 +45,8 @@ ProviderEnabled = Включено ProviderDisabled = Отключено ProviderInfo = Информация ProviderUsage = Использование +AtlasCloudAvailableBalance = Доступный баланс +AtlasCloudBalance = Баланс Atlas Cloud AuthType = Аутентификация DataSource = Источник данных ProviderNotDetected = не обнаружено diff --git a/rust/src/locale/tests.rs b/rust/src/locale/tests.rs index 9daf93d809..524d3fde35 100644 --- a/rust/src/locale/tests.rs +++ b/rust/src/locale/tests.rs @@ -80,6 +80,8 @@ fn test_japanese_menu_card_locale_values_are_translated() { (LocaleKey::DetailCostLimit, "上限"), (LocaleKey::DetailCostRemaining, "残り"), (LocaleKey::DetailCostBalance, "残高"), + (LocaleKey::AtlasCloudAvailableBalance, "利用可能残高"), + (LocaleKey::AtlasCloudBalance, "Atlas Cloud の残高"), (LocaleKey::DetailCostResets, "リセット"), (LocaleKey::DetailChartCost, "コスト（30日間）"), (LocaleKey::DetailChartCredits, "使用クレジット（30日間）"), diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index ca4fabe791..d57153f838 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -67,6 +67,8 @@ ProviderEnabled = Etkin ProviderDisabled = Devre dışı ProviderInfo = Bilgi ProviderUsage = Kullanım +AtlasCloudAvailableBalance = Kullanılabilir bakiye +AtlasCloudBalance = Atlas Cloud bakiyesi AuthType = Kimlik Doğrulama DataSource = Veri Kaynağı ProviderNotDetected = algılanmadı diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 98dd870db8..67cdedcc23 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -61,6 +61,8 @@ ProviderEnabled = 已启用 ProviderDisabled = 已禁用 ProviderInfo = 信息 ProviderUsage = 用量 +AtlasCloudAvailableBalance = 可用余额 +AtlasCloudBalance = Atlas Cloud 余额 AuthType = 认证方式 DataSource = 数据来源 ProviderNotDetected = 未检测到 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index f893d5b423..d8b8f6d968 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -61,6 +61,8 @@ ProviderEnabled = 已啟用 ProviderDisabled = 已停用 ProviderInfo = 資訊 ProviderUsage = 用量 +AtlasCloudAvailableBalance = 可用餘額 +AtlasCloudBalance = Atlas Cloud 餘額 AuthType = 認證方式 DataSource = 資料來源 ProviderNotDetected = 未檢測到 From 8da64f9d2522fb1a308e26cb5149ab4f2f72d59a Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 21:01:13 +0700 Subject: [PATCH 141/259] Address thermo review --- rust/src/cost_scanner.rs | 31 ++-- rust/src/cost_scanner/claude_roots.rs | 173 +++++++++++++++------ rust/src/cost_scanner/tests/claude_swap.rs | 7 +- 3 files changed, 151 insertions(+), 60 deletions(-) diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index b7db328e5b..c83e2d8957 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -597,7 +597,7 @@ impl CostScanner { // Walk through projects directory, de-duplicating usage records // that appear across multiple files. let mut claude_scan = ClaudeFileScanResult::default(); - if !roots.is_empty() { + if !roots.paths.is_empty() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let traversal_read_failures = { @@ -624,10 +624,11 @@ impl CostScanner { } claude_scan.absorb(file_result); }; - self.walk_claude_roots(&roots, &cutoff, cancel, &mut handle_file) + self.walk_claude_roots(&roots.paths, &cutoff, cancel, &mut handle_file) }; claude_scan.read_failures = claude_scan .read_failures + .saturating_add(roots.read_failures) .saturating_add(traversal_read_failures); } @@ -650,7 +651,7 @@ impl CostScanner { // than turning a partial zero into a known zero. finalize_claude_summary( &mut summary, - !roots.is_empty(), + !roots.paths.is_empty(), claude_scan, is_cancelled(cancel), ); @@ -684,11 +685,11 @@ impl CostScanner { let mut quota_records = Vec::new(); let mut scan_result = ClaudeFileScanResult::default(); - if !roots.is_empty() { + if !roots.paths.is_empty() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let traversal_read_failures = - self.walk_claude_roots(&roots, &cutoff, cancel, &mut |path| { + self.walk_claude_roots(&roots.paths, &cutoff, cancel, &mut |path| { let mut file_has_usage = false; let mut aggregation_complete = true; let mut file_result = scan_claude_file_with_pricing( @@ -725,6 +726,7 @@ impl CostScanner { }); scan_result.read_failures = scan_result .read_failures + .saturating_add(roots.read_failures) .saturating_add(traversal_read_failures); } @@ -736,10 +738,11 @@ impl CostScanner { &mut HashSet::new(), ); - let complete = !roots.is_empty() && !is_cancelled(cancel) && scan_result.is_complete(); + let complete = + !roots.paths.is_empty() && !is_cancelled(cancel) && scan_result.is_complete(); finalize_claude_summary( &mut summary, - !roots.is_empty(), + !roots.paths.is_empty(), scan_result, is_cancelled(cancel), ); @@ -807,7 +810,7 @@ impl CostScanner { /// Existing Claude transcript roots: the profile root plus claude-swap /// session homes, de-duplicated by resolved path. - fn claude_projects_roots(&self) -> Vec { + fn claude_projects_roots(&self) -> claude_roots::ClaudeProjectsRoots { claude_roots::claude_projects_roots( std::env::var("CLAUDE_CONFIG_DIR").ok().as_deref(), dirs::home_dir().as_deref(), @@ -1212,7 +1215,7 @@ pub fn has_cost_usage_sources() -> bool { .get_codex_sessions_dirs() .iter() .any(|dir| dir.exists()) - || !scanner.claude_projects_roots().is_empty() + || scanner.claude_projects_roots().has_possible_roots() || crate::pi_session_cost::pi_compatible_session_roots(dirs::home_dir()) .iter() .any(|dir| dir.exists()) @@ -1276,7 +1279,7 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< // Real per-day breakdown: walk the project logs once, // de-duplicating records across files. let roots = scanner.claude_projects_roots(); - if !roots.is_empty() { + if !roots.paths.is_empty() { let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); @@ -1306,10 +1309,11 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< } claude_scan.absorb(file_result); }; - scanner.walk_claude_roots(&roots, &cutoff, None, &mut handle_file) + scanner.walk_claude_roots(&roots.paths, &cutoff, None, &mut handle_file) }; claude_scan.read_failures = claude_scan .read_failures + .saturating_add(roots.read_failures) .saturating_add(traversal_read_failures); if claude_scan.is_complete() { zero_fill_uninitialized_claude_daily_costs( @@ -1398,7 +1402,7 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> // as the cost chart. Only a complete valid walk establishes // authoritative coverage of the requested history window. let roots = scanner.claude_projects_roots(); - if !roots.is_empty() { + if !roots.paths.is_empty() { let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); @@ -1414,10 +1418,11 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> ); claude_scan.absorb(file_result); }; - scanner.walk_claude_roots(&roots, &cutoff, None, &mut handle_file) + scanner.walk_claude_roots(&roots.paths, &cutoff, None, &mut handle_file) }; claude_scan.read_failures = claude_scan .read_failures + .saturating_add(roots.read_failures) .saturating_add(traversal_read_failures); mark_claude_daily_token_coverage(&mut covered_days, &daily_tokens, claude_scan); } diff --git a/rust/src/cost_scanner/claude_roots.rs b/rust/src/cost_scanner/claude_roots.rs index b127044d78..c1b983bfa5 100644 --- a/rust/src/cost_scanner/claude_roots.rs +++ b/rust/src/cost_scanner/claude_roots.rs @@ -11,37 +11,68 @@ //! claude-swap uses the legacy `~/.claude-swap-backup` root on Windows and //! macOS. The XDG data root applies to Linux/WSL only and is not needed here. +use std::collections::HashSet; use std::fs; use std::path::{Path, PathBuf}; -const SWAP_SESSIONS_DIR: [&str; 2] = [".claude-swap-backup", "sessions"]; +#[derive(Debug, Default)] +pub(super) struct ClaudeProjectsRoots { + pub(super) paths: Vec, + pub(super) read_failures: u32, +} + +impl ClaudeProjectsRoots { + pub(super) fn has_possible_roots(&self) -> bool { + !self.paths.is_empty() || self.read_failures > 0 + } +} /// Existing Claude transcript roots, de-duplicated by resolved path. /// /// `config_dir` is the raw `CLAUDE_CONFIG_DIR` value: one literal directory, -/// with empty meaning "unset". A root that does not exist is omitted, so an -/// empty result means there is no Claude history to scan. -pub(super) fn claude_projects_roots(config_dir: Option<&str>, home: Option<&Path>) -> Vec { - let mut candidates = vec![base_projects_dir(config_dir, home)]; +/// with empty meaning "unset". A root that does not exist is omitted; failures +/// that prevent discovering existing roots are retained for scan coverage. +pub(super) fn claude_projects_roots( + config_dir: Option<&str>, + home: Option<&Path>, +) -> ClaudeProjectsRoots { + let mut result = ClaudeProjectsRoots { + paths: vec![base_projects_dir(config_dir, home)], + ..ClaudeProjectsRoots::default() + }; if let Some(home) = home { - candidates.extend(swap_projects_roots(home)); + let swaps = swap_projects_roots(home); + result.paths.extend(swaps.paths); + result.read_failures = result.read_failures.saturating_add(swaps.read_failures); } - let mut seen = Vec::new(); - let mut roots = Vec::new(); - for root in candidates { - if !root.exists() { - continue; + let mut seen = HashSet::new(); + let mut paths = Vec::new(); + for root in result.paths.drain(..) { + match fs::metadata(&root) { + Ok(_) => {} + Err(error) if is_missing_path(&error) => continue, + Err(_) => { + result.read_failures = result.read_failures.saturating_add(1); + continue; + } } + // Shared-history symlinks or junctions resolve to the same directory // as the profile they point at; scan that directory once. - let resolved = fs::canonicalize(&root).unwrap_or_else(|_| root.clone()); - if !seen.contains(&resolved) { - seen.push(resolved); - roots.push(root); + let resolved = match fs::canonicalize(&root) { + Ok(resolved) => resolved, + Err(_) => { + result.read_failures = result.read_failures.saturating_add(1); + root.clone() + } + }; + if seen.insert(resolved) { + paths.push(root); } } - roots + result.paths = paths; + result } fn base_projects_dir(config_dir: Option<&str>, home: Option<&Path>) -> PathBuf { @@ -60,25 +91,54 @@ fn base_projects_dir(config_dir: Option<&str>, home: Option<&Path>) -> PathBuf { /// `/.claude-swap-backup/sessions/-