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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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/407] 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")}
+
+ )} {error &&

{error}

} ); diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 88754baa7d..bc09569e7e 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -994,6 +994,9 @@ export interface ProviderDetail { /** Phase 6c — currently-persisted cookie source value ("auto" | "manual" | "off" | …). * `null` for providers that do not expose a cookie-source picker. */ cookieSource: string | null; + /** Manual cookie source is selected with no usable header, and the provider + * fails closed instead of importing browser cookies. */ + manualCookieMissing?: boolean; /** Phase 6c — currently-persisted region value. `null` for non-regional providers. */ region: string | null; } diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..deca556214 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -539,6 +539,8 @@ locale_keys! { ProviderCodexNotRunningHelp, ProviderCookieSource, CookieSourceManual, + ProviderManualCookieMissing, + ProviderUseAutomaticCookies, ProviderRegion, ProviderClaudeCookies, ProviderClaudeCookiesHelp, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..a9f4b0b54a 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -284,6 +284,8 @@ ProviderCodexLastFetchFailedTitle = Last Codex fetch failed: ProviderCodexNotRunningHelp = Codex not running. Try running a Codex command first. ProviderCookieSource = Cookie source CookieSourceManual = Manual +ProviderManualCookieMissing = No cookie header pasted. +ProviderUseAutomaticCookies = Use automatic cookies ProviderRegion = Region ProviderClaudeCookies = Claude cookies ProviderClaudeCookiesHelp = Browser cookies/sessionKey are preferred because they match Claude's settings usage page. diff --git a/rust/src/providers/ollama/cookies.rs b/rust/src/providers/ollama/cookies.rs index e6130aef9a..9ec405149e 100644 --- a/rust/src/providers/ollama/cookies.rs +++ b/rust/src/providers/ollama/cookies.rs @@ -97,12 +97,28 @@ pub(super) fn resolve_cookie_source( .ok_or(ProviderError::NoCookies); } + // Manual source with nothing pasted: fail closed before any browser import + // (upstream `OllamaUsageError.manualCookieHeaderEmpty`, #3891). + if ctx.manual_cookie_missing { + return Err(manual_cookie_header_empty_error()); + } + match resolve_browser_cookie_header(false)? { Some(header) => Ok(OllamaCookieSource::Manual(header)), None => Err(ProviderError::NoCookies), } } +pub(super) fn manual_cookie_header_empty_error() -> ProviderError { + ProviderError::Other( + concat!( + "Ollama cookie source is Manual, but no cookie header is configured. ", + "Paste a Cookie header from https://ollama.com/settings, or set Cookie source to Auto." + ) + .to_string(), + ) +} + /// After a successful web fetch, cache the validated browser/manual session header. pub(super) fn cache_validated_session_cookie(source: &OllamaCookieSource) { use crate::browser::cookie_cache::CookieHeaderCache; @@ -276,6 +292,49 @@ pub(super) fn is_ollama_sign_in_redirect(url: &Url) -> bool { mod tests { use super::*; + #[test] + fn empty_manual_cookie_source_fails_closed_with_upstream_message() { + // Upstream #3891 `OllamaUsageError.manualCookieHeaderEmpty`. The shell + // sets `manual_cookie_missing`; no browser import may be attempted. + let ctx = FetchContext { + manual_cookie_missing: true, + ..FetchContext::default() + }; + let Err(error) = resolve_cookie_source(&ctx) else { + panic!("empty Manual source must not resolve a cookie"); + }; + assert_eq!( + error.to_string(), + "Ollama cookie source is Manual, but no cookie header is configured. Paste a Cookie header from https://ollama.com/settings, or set Cookie source to Auto." + ); + } + + #[test] + fn blank_manual_header_with_missing_flag_fails_closed() { + let ctx = FetchContext { + manual_cookie_header: Some(" ".to_string()), + manual_cookie_missing: true, + ..FetchContext::default() + }; + assert!(matches!( + resolve_cookie_source(&ctx), + Err(ProviderError::Other(message)) if message.contains("no cookie header is configured") + )); + } + + #[test] + fn non_empty_header_without_session_cookie_keeps_existing_error() { + let ctx = FetchContext { + manual_cookie_header: Some("theme=dark".to_string()), + manual_cookie_missing: false, + ..FetchContext::default() + }; + assert!(matches!( + resolve_cookie_source(&ctx), + Err(ProviderError::NoCookies) + )); + } + #[test] fn normalizes_raw_ollama_session_cookie_value() { assert_eq!( diff --git a/rust/src/providers/ollama/mod.rs b/rust/src/providers/ollama/mod.rs index 561330100c..d042f5ca40 100755 --- a/rust/src/providers/ollama/mod.rs +++ b/rust/src/providers/ollama/mod.rs @@ -14,8 +14,8 @@ use reqwest::Url; use serde::Deserialize; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + FetchContext, ManualEmptyCookiePolicy, Provider, ProviderError, ProviderFetchResult, + ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; use crate::settings::ApiKeys; @@ -416,6 +416,10 @@ impl Provider for OllamaProvider { fn retains_last_good_on_transport_failure(&self) -> bool { true } + + fn manual_empty_cookie_policy(&self) -> ManualEmptyCookiePolicy { + ManualEmptyCookiePolicy::FailClosedWeb + } } fn clean_secret(raw: Option<&str>) -> Option { From 7e942412906fc6447f112458db752701faad16f2 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 16:23:27 +0700 Subject: [PATCH 108/407] Port upstream 0.66.0: Chutes quota classification and enrichment --- rust/src/providers/chutes/mod.rs | 646 ++++++------------------ rust/src/providers/chutes/parse.rs | 667 +++++++++++++++++++++++++ rust/src/providers/chutes/tests.rs | 754 +++++++++++++++++++++++++++++ 3 files changed, 1562 insertions(+), 505 deletions(-) create mode 100644 rust/src/providers/chutes/parse.rs create mode 100644 rust/src/providers/chutes/tests.rs diff --git a/rust/src/providers/chutes/mod.rs b/rust/src/providers/chutes/mod.rs index c6c75c4727..c2b1d45f92 100644 --- a/rust/src/providers/chutes/mod.rs +++ b/rust/src/providers/chutes/mod.rs @@ -1,14 +1,21 @@ use async_trait::async_trait; -use chrono::{DateTime, TimeZone, Utc}; -use reqwest::Client; -use serde_json::Value; +use reqwest::header::ACCEPT; +use reqwest::{Client, StatusCode, Url}; +use serde_json::{Map, Value, json}; use crate::core::{ FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + SourceMode, UsageSnapshot, }; +mod parse; +#[cfg(test)] +mod tests; + +use parse::{ParsedQuotas, parse, text}; + const CREDENTIAL_TARGET: &str = "codexbar-chutes"; +const DEFAULT_API_URL: &str = "https://api.chutes.ai"; pub struct ChutesProvider { metadata: ProviderMetadata, @@ -64,26 +71,13 @@ impl Provider for ChutesProvider { &["CHUTES_API_KEY"], )?; let base = std::env::var("CHUTES_API_URL") - .unwrap_or_else(|_| "https://api.chutes.ai".into()); - let url = crate::providers::validated_https_url(&base, "Chutes API")? - .join("users/me/subscription_usage") - .map_err(|e| ProviderError::Other(format!("Invalid Chutes API URL: {e}")))?; - let response = self.client.get(url).bearer_auth(key).send().await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); - } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Chutes usage returned status {}", - response.status() - ))); - } - let value: Value = response.json().await.map_err(|e| { - ProviderError::Parse(format!("Failed to parse Chutes usage: {e}")) - })?; - Ok(ProviderFetchResult::new(snapshot_from_usage(&value), "api")) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| DEFAULT_API_URL.to_string()); + let base = crate::providers::validated_https_url(&base, "Chutes API")?; + let usage = fetch_usage_snapshot(&self.client, &base, &key).await?; + Ok(ProviderFetchResult::new(usage, "api")) } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -96,494 +90,136 @@ impl Provider for ChutesProvider { } } -fn snapshot_from_usage(value: &Value) -> UsageSnapshot { - let windows = quota_windows(value); - let primary = windows - .first() - .cloned() - .unwrap_or_else(|| RateWindow::new(0.0)); - let mut snapshot = UsageSnapshot::new(primary); - if let Some(second) = windows.get(1).cloned() { - snapshot = snapshot.with_secondary(second); - } - snapshot.with_login_method("Chutes API") -} - -fn quota_windows(value: &Value) -> Vec { - let mut out = Vec::new(); - collect_windows(value, &mut out); - out -} - -fn collect_windows(value: &Value, out: &mut Vec) { - match value { - Value::Object(map) => { - if let Some(window) = window_from_object(map) { - out.push(window); - } - for value in map.values() { - collect_windows(value, out); - } - } - Value::Array(items) => { - for value in items { - collect_windows(value, out); - } - } - _ => {} - } -} - -/// Build a rate window from a quota object. +/// Read `users/me/subscription_usage`, then enrich it from `quotas` and +/// `quota_usage/{id}` only when the rolling or monthly lane is still missing. /// -/// Quota counts (`used`/`limit`) go into `reset_description` as detail text -/// (e.g. `"25/100 credits"`). They are never treated as reset schedules. -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, - window_minutes, - resets_at, - detail, - )) -} - -fn percent_from_object(map: &serde_json::Map) -> Option { - for key in [ - "usage_percent", - "usagePercent", - "percent_used", - "percentUsed", - ] { - let Some(v) = map.get(key).and_then(Value::as_f64) else { - continue; - }; - // Chutes usage/quota payloads (GET /users/me/subscription_usage, and - // whatever `collect_windows` recurses into) carry whole percent values - // in 0..=100 — no documented field is a 0..=1 fraction. Rescaling - // 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). - 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( - map, - &["limit", "quota", "quota_limit", "quotaLimit", "total"], - ); - let remaining = first_f64(map, &["remaining", "remaining_quota", "remainingQuota"]); - match (used, limit, remaining) { - (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 => { - 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; - if limit <= 0.0 { - return None; - } - let percent = (used / limit) * 100.0; - percent.is_finite().then_some(percent.clamp(0.0, 100.0)) - } - _ => None, - } -} - -fn quota_count_description(map: &serde_json::Map) -> Option { - let used = first_f64(map, &["used", "usage", "current_usage", "currentUsage"]); - let mut limit = first_f64( - map, - &["limit", "quota", "quota_limit", "quotaLimit", "total"], - ); - let remaining = first_f64(map, &["remaining", "remaining_quota", "remainingQuota"]); - - if limit.is_none() - && let (Some(used), Some(remaining)) = (used, remaining) - { - limit = Some(used + remaining); - } - 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))?, - }; - let unit = first_str(map, &["unit", "units", "quota_unit", "quotaUnit"]).unwrap_or("credits"); - Some(format!( - "{}/{} {}", - format_quota_amount(used), - format_quota_amount(limit), - unit - )) -} - -fn window_minutes_from_object(map: &serde_json::Map) -> 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); +/// Enrichment is best effort: a rejected key stays fatal, every other failure +/// leaves the subscription result as it was. A cancelled refresh drops this +/// future, so cancellation needs no separate handling here. +async fn fetch_usage_snapshot( + client: &Client, + base: &Url, + key: &str, +) -> Result { + let mut subscription = parse(&get_json(client, base, key, &["subscription_usage"]).await?); + if !subscription.has_rolling_and_monthly() { + match fetch_quota_details(client, base, key).await { + Ok(quotas) if quotas.lanes().has_windows() => subscription.fill_missing_from(quotas), + Ok(_) => {} + Err(ProviderError::AuthRequired) => return Err(ProviderError::AuthRequired), + Err(error) => tracing::debug!(%error, "Chutes quota enrichment failed"), } } - - ["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; + Ok(subscription.lanes().into_usage()) +} + +/// `GET users/me/quotas`, then `GET users/me/quota_usage/{id}` for each quota +/// definition, merged with the usage fields taking precedence. +async fn fetch_quota_details( + client: &Client, + base: &Url, + key: &str, +) -> Result { + let raw = get_json(client, base, key, &["quotas"]).await?; + let mut quotas = parse(&raw); + let Some(definitions) = + quota_definitions(&raw).filter(|list| list.iter().any(Value::is_object)) + else { + return Ok(quotas); }; - rounded_window_minutes(value * multiplier) -} - -fn first_reset_timestamp(map: &serde_json::Map) -> Option> { - for key in [ - "resets_at", - "resetsAt", - "reset_at", - "resetAt", - "reset_time", - "resetTime", - "renews_at", - "renewsAt", - ] { - if let Some(dt) = map.get(key).and_then(parse_timestamp_value) { - return Some(dt); - } - } - None -} - -fn parse_timestamp_value(value: &Value) -> Option> { - match value { - Value::String(raw) => { - let text = raw.trim(); - if text.is_empty() { - return None; - } - // ISO-8601 / RFC3339 only for strings that look like dates. - // Bare small numbers as strings are quota counts, not schedules. - if let Ok(dt) = DateTime::parse_from_rfc3339(text) { - return Some(dt.with_timezone(&Utc)); - } - if let Ok(number) = text.parse::() { - return epoch_to_datetime(number); + let mut enriched = Vec::new(); + for definition in definitions.iter().filter_map(Value::as_object) { + let mut merged = definition.clone(); + if let Some(id) = quota_id(definition) { + match get_json(client, base, key, &["quota_usage", &id]).await { + Ok(usage) => merged.extend(usage_fields(&usage)), + Err(ProviderError::AuthRequired) => return Err(ProviderError::AuthRequired), + Err(error) => tracing::debug!(%error, "Chutes quota usage request failed"), } - None } - Value::Number(n) => n.as_f64().and_then(epoch_to_datetime), - _ => None, - } + enriched.push(Value::Object(merged)); + } + let detailed = parse(&json!({ "quotas": enriched })); + if detailed.lanes().has_windows() { + quotas = detailed; + } + Ok(quotas) +} + +/// The quota definition list: the payload itself, `quotas`, `data`, or +/// `data.quotas`, whichever is an array first. +fn quota_definitions(raw: &Value) -> Option<&Vec> { + let root = raw.as_object(); + let data = root.and_then(|root| root.get("data")); + [ + Some(raw), + root.and_then(|root| root.get("quotas")), + data, + data.and_then(Value::as_object) + .and_then(|data| data.get("quotas")), + ] + .into_iter() + .flatten() + .find_map(Value::as_array) +} + +fn quota_id(definition: &Map) -> Option { + ["chute_id", "chuteId", "id"] + .iter() + .find_map(|key| text(definition.get(*key))) } -fn epoch_to_datetime(value: f64) -> Option> { - if !value.is_finite() || value <= 0.0 { - return None; - } - // Reject small integers that are clearly quota counts, not unix epochs. - // Unix seconds ~1e9; ms ~1e12. Quota counts are typically << 1e8. - if value < 1_000_000_000.0 { - return None; - } - let seconds = if value > 10_000_000_000.0 { - value / 1000.0 - } else { - value +/// Usage fields of a `quota_usage/{id}` response, unwrapping `data` or `result`. +fn usage_fields(response: &Value) -> Map { + let Some(result) = response.as_object() else { + return Map::new(); }; - // Epoch seconds; the sub-second fraction is below timestamp resolution. - #[expect( - clippy::cast_possible_truncation, - reason = "epoch seconds; sub-second fraction below timestamp resolution" - )] - let whole_seconds = seconds as i64; - Utc.timestamp_opt(whole_seconds, 0).single() -} - -fn first_f64(map: &serde_json::Map, keys: &[&str]) -> Option { - 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> { - keys.iter() - .find_map(|k| map.get(*k).and_then(Value::as_str)) - .map(str::trim) - .filter(|s| !s.is_empty()) -} - -fn format_quota_amount(value: f64) -> String { - 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 and the rounded value fits in i64. - #[allow( - clippy::cast_possible_truncation, - reason = "finite rounded value is bounded to the i64 range above" - )] - let whole = rounded as i64; - format!("{}", whole) - } else { - let mut text = format!("{value:.2}"); - while text.contains('.') && text.ends_with('0') { - text.pop(); - } - if text.ends_with('.') { - text.pop(); - } - text - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_ratio_window() { - let snapshot = - snapshot_from_usage(&serde_json::json!({"quotas":[{"used":25,"limit":100}]})); - assert_eq!(snapshot.primary.used_percent, 25.0); - } - - #[test] - fn quota_counts_go_to_reset_description_not_schedule() { - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{ - "used": 25, - "limit": 100, - "unit": "credits", - "remaining": 75 - }] - })); - assert_eq!(snapshot.primary.used_percent, 25.0); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("25/100 credits") - ); - assert!(snapshot.primary.resets_at.is_none()); - } - - #[test] - fn small_numeric_fields_are_not_reset_schedules() { - // A mis-keyed payload where "reset_time" is accidentally a remaining count. - // Values below unix-epoch range must not become resets_at. - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{ - "used": 10, - "limit": 50, - "reset_time": 40 - }] - })); - assert_eq!(snapshot.primary.used_percent, 20.0); - assert!(snapshot.primary.resets_at.is_none()); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("10/50 credits") - ); - } - - #[test] - fn iso_reset_timestamps_still_parse() { - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{ - "used": 1, - "limit": 2, - "resets_at": "2026-08-01T00:00:00Z" - }] - })); - assert_eq!( - snapshot.primary.resets_at, - Some(Utc.with_ymd_and_hms(2026, 8, 1, 0, 0, 0).unwrap()) - ); - } - - #[test] - fn usage_percent_one_stays_one() { - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{"usage_percent": 1, "limit": 100}] - })); - assert_eq!(snapshot.primary.used_percent, 1.0); - } - - #[test] - fn usage_percent_half_stays_half() { - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{"usage_percent": 0.5, "limit": 100}] - })); - assert_eq!(snapshot.primary.used_percent, 0.5); - } - - #[test] - fn usage_percent_hundred_stays_hundred() { - let snapshot = snapshot_from_usage(&serde_json::json!({ - "quotas": [{"usage_percent": 100, "limit": 100}] - })); - 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"); - } + ["data", "result"] + .iter() + .find_map(|key| result.get(*key).and_then(Value::as_object)) + .unwrap_or(result) + .clone() +} + +/// `{base}/users/me/{segments...}`; the configured path and query are kept. +fn endpoint(base: &Url, segments: &[&str]) -> Result { + let mut url = base.clone(); + url.set_fragment(None); + url.path_segments_mut() + .map_err(|()| ProviderError::Other("Invalid Chutes API URL: cannot be a base".into()))? + .pop_if_empty() + .push("users") + .push("me") + .extend(segments); + Ok(url) +} + +async fn get_json( + client: &Client, + base: &Url, + key: &str, + segments: &[&str], +) -> Result { + let response = client + .get(endpoint(base, segments)?) + .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.is_success() { + return Err(ProviderError::Other(format!( + "Chutes usage API error: HTTP {}", + status.as_u16() + ))); + } + response + .json::() + .await + .ok() + .filter(|value| value.is_object() || value.is_array()) + .ok_or_else(|| ProviderError::Parse("Chutes usage response is not valid JSON".into())) } diff --git a/rust/src/providers/chutes/parse.rs b/rust/src/providers/chutes/parse.rs new file mode 100644 index 0000000000..e8ac159a27 --- /dev/null +++ b/rust/src/providers/chutes/parse.rs @@ -0,0 +1,667 @@ +//! Chutes quota payload parsing. +//! +//! Mirrors upstream `chutes.js` (`chutesParse`, `chutesQuota`, `chutesSnapshot`) +//! at v0.66.0: payload keys are matched case- and punctuation-insensitively +//! against fixed key lists, quota-shaped objects are collected in traversal +//! order, and the rolling (4-hour) and monthly quotas are picked by payload +//! key, label, or window length. + +use std::collections::HashSet; + +use chrono::{DateTime, TimeZone, Utc}; +use serde_json::{Map, Value}; + +use crate::core::{RateWindow, SubscriptionMetadata, UsageSnapshot}; + +type Object = Map; + +const ROLLING_MINUTES: u32 = 240; +const MONTHLY_MINUTES: u32 = 43_200; +/// Windows of at least four weeks are classified as monthly. +const MONTHLY_MIN_WINDOW_MINUTES: u32 = 40_320; + +const ROLLING_PAYLOAD_KEYS: &[&str] = &[ + "rolling", + "rollingwindow", + "rolling4h", + "fourhour", + "fourhourusage", + "window4h", +]; +const MONTHLY_PAYLOAD_KEYS: &[&str] = &[ + "monthly", + "monthlyusage", + "subscription", + "subscriptionusage", + "billingperiod", +]; +const QUOTA_CONTAINER_KEYS: &[&str] = &[ + "quotas", + "quota", + "quotausage", + "limits", + "usage", + "entries", + "subscriptionusage", +]; +const LABEL_KEYS: &[&str] = &[ + "label", + "name", + "title", + "type", + "quotatype", + "period", + "window", + "windowname", + "chuteid", +]; +const LIMIT_KEYS: &[&str] = &[ + "limit", + "cap", + "max", + "maximum", + "quota", + "quotalimit", + "monthlycap", + "monthlylimit", + "requestlimit", + "tokenlimit", + "hardlimit", + "total", +]; +const USED_KEYS: &[&str] = &[ + "used", + "usage", + "usedamount", + "consumed", + "consumedamount", + "current", + "currentusage", + "requests", + "requestcount", + "tokens", + "tokenusage", + "monthlyusage", +]; +const REMAINING_KEYS: &[&str] = &[ + "remaining", + "available", + "balance", + "left", + "remainingamount", + "availableamount", +]; +const PERCENT_USED_KEYS: &[&str] = &[ + "percentused", + "usagepercent", + "usedpercent", + "utilization", + "utilizationpercent", +]; +const PERCENT_REMAINING_KEYS: &[&str] = &["percentremaining", "remainingpercent"]; +const RESET_KEYS: &[&str] = &[ + "resetat", + "resetsat", + "resettime", + "nextresetat", + "renewsat", + "renewalat", + "periodend", + "currentperiodend", + "expiresat", + "windowend", + "endtime", +]; +const UNIT_KEYS: &[&str] = &["unit", "units", "currency", "quotaunit"]; +const ACTIVE_KEYS: &[&str] = &[ + "active", + "isactive", + "subscriptionactive", + "hassubscription", +]; +const STATUS_KEYS: &[&str] = &["status", "state", "subscriptionstatus"]; +const PLAN_KEYS: &[&str] = &[ + "planname", + "plan", + "tier", + "subscriptionplan", + "subscriptiontier", +]; +const SUBSCRIPTION_OBJECT_KEYS: &[&str] = &[ + "subscription", + "subscriptionusage", + "currentsubscription", + "plan", +]; +const WRAPPER_KEYS: &[&str] = &["data", "result"]; +const WINDOW_MINUTE_KEYS: &[&str] = &["windowminutes", "periodminutes", "durationminutes"]; +const WINDOW_HOUR_KEYS: &[&str] = &["windowhours", "periodhours", "durationhours"]; +const WINDOW_DAY_KEYS: &[&str] = &["windowdays", "perioddays", "durationdays"]; +const WINDOW_SECOND_KEYS: &[&str] = &["windowseconds", "periodseconds", "durationseconds"]; +const WINDOW_TEXT_KEYS: &[&str] = &["window", "period", "interval", "duration"]; + +/// A quota-shaped object projected into a display window. +#[derive(Debug, Clone)] +pub(super) struct QuotaWindow { + label: Option, + unit: String, + /// Canonical form of the source object; identifies the same quota when it + /// is reachable both by payload key and by traversal. + raw_key: String, + rate: RateWindow, +} + +/// Quotas selected from one Chutes payload. +#[derive(Debug, Default)] +pub(super) struct ParsedQuotas { + rolling: Option, + monthly: Option, + fallback: Vec, + plan: Option, + active: Option, + renewal: Option>, +} + +impl ParsedQuotas { + /// Keep lanes this payload already has and take the rest from `other` + /// (upstream `??=` merge of `quotas` enrichment). + pub(super) fn fill_missing_from(&mut self, other: ParsedQuotas) { + if self.rolling.is_none() { + self.rolling = other.rolling; + } + if self.monthly.is_none() { + self.monthly = other.monthly; + } + self.fallback.extend(other.fallback); + } + + pub(super) fn has_rolling_and_monthly(&self) -> bool { + self.rolling.is_some() && self.monthly.is_some() + } + + pub(super) fn lanes(&self) -> Lanes { + let with_default_minutes = |window: &QuotaWindow, default: u32| { + let mut rate = window.rate.clone(); + rate.window_minutes = rate.window_minutes.or(Some(default)); + rate + }; + let rolling = self + .rolling + .as_ref() + .map(|window| with_default_minutes(window, ROLLING_MINUTES)); + let monthly = self + .monthly + .as_ref() + .map(|window| with_default_minutes(window, MONTHLY_MINUTES)); + let fallback: Vec<&RateWindow> = self.fallback.iter().map(|window| &window.rate).collect(); + let primary = rolling.clone().or_else(|| { + if monthly.is_none() { + fallback.first().map(|window| (*window).clone()) + } else { + None + } + }); + // Without a rolling lane the first fallback already became the + // primary, so the secondary takes the next one. + let secondary = monthly.or_else(|| { + let index = usize::from(rolling.is_none()); + fallback.get(index).map(|window| (*window).clone()) + }); + let login_method = self.plan.clone().or_else(|| { + if self.active == Some(false) { + Some("No active subscription".to_string()) + } else if self.active.is_none() && primary.is_none() && secondary.is_none() { + Some("No usage data".to_string()) + } else { + None + } + }); + Lanes { + primary, + secondary, + renewal: self.renewal, + login_method, + } + } +} + +/// Primary and secondary windows plus identity resolved from [`ParsedQuotas`] +/// (upstream `chutesSnapshot`). +#[derive(Debug)] +pub(super) struct Lanes { + primary: Option, + secondary: Option, + renewal: Option>, + login_method: Option, +} + +impl Lanes { + pub(super) fn has_windows(&self) -> bool { + self.primary.is_some() || self.secondary.is_some() + } + + pub(super) fn into_usage(self) -> UsageSnapshot { + // UsageSnapshot requires a primary window; an absent 4-hour lane is + // shown as informational rather than a fake 0% quota. + let primary = self + .primary + .unwrap_or_else(|| RateWindow::informational("No quota reported")); + let mut usage = UsageSnapshot::new(primary); + if let Some(secondary) = self.secondary { + usage = usage.with_secondary(secondary); + } + if let Some(login_method) = self.login_method { + usage = usage.with_login_method(login_method); + } + // Only an explicit subscription date counts; a quota reset is not one. + if let Some(renewal) = self.renewal { + usage = + usage.with_subscription(Some(SubscriptionMetadata::new(None, None, Some(renewal)))); + } + usage + } +} + +/// Selected quotas as a snapshot, for callers that only need the result. +#[cfg(test)] +pub(super) fn snapshot_from_usage(value: &Value) -> UsageSnapshot { + parse(value).lanes().into_usage() +} + +fn normalized_key(key: &str) -> String { + key.chars() + .filter(char::is_ascii_alphanumeric) + .map(|c| c.to_ascii_lowercase()) + .collect() +} + +/// Value of the first of `keys` that is present, matching payload keys after +/// lowercasing and dropping non-alphanumerics (upstream `chutesValue`). +fn lookup<'a>(object: &'a Object, keys: &[&str]) -> Option<&'a Value> { + keys.iter().find_map(|key| { + object + .iter() + .find(|(candidate, _)| normalized_key(candidate) == *key) + .map(|(_, value)| value) + }) +} + +fn as_object(value: Option<&Value>) -> Option<&Object> { + value.and_then(Value::as_object) +} + +fn number(value: Option<&Value>) -> Option { + let parsed = match value? { + Value::Number(number) => number.as_f64(), + Value::String(text) => { + let cleaned: String = text + .trim() + .chars() + .filter(|c| !matches!(c, ',' | '$' | '%')) + .collect(); + cleaned.trim().parse::().ok() + } + Value::Bool(flag) => Some(f64::from(u8::from(*flag))), + _ => None, + }; + parsed.filter(|number| number.is_finite()) +} + +pub(super) fn text(value: Option<&Value>) -> Option { + let text = match value? { + Value::Bool(flag) => (if *flag { "1" } else { "0" }).to_string(), + Value::String(text) => text.trim().to_string(), + Value::Number(number) => number.to_string(), + _ => return None, + }; + (!text.is_empty()).then_some(text) +} + +fn flag(value: Option<&Value>) -> Option { + match value? { + Value::Bool(flag) => Some(*flag), + Value::Number(number) => number.as_f64().map(|number| number != 0.0), + Value::String(text) => match text.trim().to_ascii_lowercase().as_str() { + "true" | "1" | "yes" | "active" => Some(true), + "false" | "0" | "no" | "inactive" | "none" => Some(false), + _ => None, + }, + _ => None, + } +} + +fn date(value: Option<&Value>) -> Option> { + match value? { + Value::String(raw) => { + let raw = raw.trim(); + // ISO-8601 / RFC3339 only for strings that look like dates. + if let Ok(parsed) = DateTime::parse_from_rfc3339(raw) { + return Some(parsed.with_timezone(&Utc)); + } + epoch_to_datetime(raw.parse::().ok()?) + } + Value::Number(number) => number.as_f64().and_then(epoch_to_datetime), + _ => None, + } +} + +fn epoch_to_datetime(value: f64) -> Option> { + if !value.is_finite() || value <= 0.0 { + return None; + } + // Reject small integers that are clearly quota counts, not unix epochs. + // Unix seconds ~1e9; ms ~1e12. Quota counts are typically << 1e8. + if value < 1_000_000_000.0 { + return None; + } + let seconds = if value > 10_000_000_000.0 { + value / 1000.0 + } else { + value + }; + // Epoch seconds; the sub-second fraction is below timestamp resolution. + #[expect( + clippy::cast_possible_truncation, + reason = "epoch seconds; sub-second fraction below timestamp resolution" + )] + let whole_seconds = seconds as i64; + Utc.timestamp_opt(whole_seconds, 0).single() +} + +fn window_minutes(payload: &Object) -> Option { + for (keys, multiplier) in [ + (WINDOW_MINUTE_KEYS, 1.0), + (WINDOW_HOUR_KEYS, 60.0), + (WINDOW_DAY_KEYS, 24.0 * 60.0), + (WINDOW_SECOND_KEYS, 1.0 / 60.0), + ] { + if let Some(value) = number(lookup(payload, keys)) { + return rounded_window_minutes(value * multiplier); + } + } + text(lookup(payload, WINDOW_TEXT_KEYS)).and_then(|raw| parse_window_duration_text(&raw)) +} + +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()) + .map(|character| character.to_ascii_lowercase()) + .collect(); + let split_at = compact.find(|character: char| { + !character.is_ascii_digit() && !matches!(character, '.' | '+' | '-' | '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" { + f64::from(MONTHLY_MINUTES) + } else { + return None; + }; + rounded_window_minutes(value * multiplier) +} + +pub(super) fn format_quota_amount(value: f64) -> String { + 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 and the rounded value fits in i64. + #[allow( + clippy::cast_possible_truncation, + reason = "finite rounded value is bounded to the i64 range above" + )] + let whole = rounded as i64; + format!("{whole}") + } else { + let mut text = format!("{value:.2}"); + while text.contains('.') && text.ends_with('0') { + text.pop(); + } + if text.ends_with('.') { + text.pop(); + } + text + } +} + +/// Build a display window from one quota object (upstream `chutesQuota`). +/// +/// Quota counts (`used`/`limit`) go into `reset_description` as detail text +/// (e.g. `"25/100 credits"`). They are never treated as reset schedules. +/// Percent fields are whole percents in 0..=100, never 0..=1 fractions +/// (#408: a real 1% must not become 100%). +fn quota(payload: &Object, label: Option<&str>, minutes: Option) -> Option { + let read = |keys| number(lookup(payload, keys)); + let limit = read(LIMIT_KEYS); + let used = read(USED_KEYS); + let remaining = read(REMAINING_KEYS); + let percent_used = read(PERCENT_USED_KEYS).map(|value| value.clamp(0.0, 100.0)); + let percent_remaining = + read(PERCENT_REMAINING_KEYS).map(|value| 100.0 - value.clamp(0.0, 100.0)); + let total = limit.or_else(|| { + used.zip(remaining) + .map(|(used, remaining)| used + remaining) + }); + let consumed = used.or_else(|| { + total + .zip(remaining) + .map(|(total, remaining)| total - remaining) + }); + let percent = percent_used.or(percent_remaining).or_else(|| { + let (total, consumed) = total.zip(consumed)?; + (total > 0.0).then(|| consumed / total * 100.0) + })?; + let unit = text(lookup(payload, UNIT_KEYS)).unwrap_or_else(|| "credits".to_string()); + let described_used = used.or_else(|| { + limit + .zip(remaining) + .map(|(limit, remaining)| (limit - remaining).max(0.0)) + }); + let detail = limit + .filter(|limit| *limit > 0.0) + .zip(described_used) + .map(|(limit, used)| { + format!( + "{}/{} {unit}", + format_quota_amount(used), + format_quota_amount(limit) + ) + }); + Some(QuotaWindow { + label: text(lookup(payload, LABEL_KEYS)).or_else(|| label.map(str::to_string)), + unit, + raw_key: Value::Object(payload.clone()).to_string(), + rate: RateWindow::with_details( + percent.clamp(0.0, 100.0), + window_minutes(payload).or(minutes), + date(lookup(payload, RESET_KEYS)), + detail, + ), + }) +} + +fn collect_windows(value: &Value, seen: &mut HashSet, windows: &mut Vec) { + match value { + Value::Array(items) => { + for item in items { + collect_windows(item, seen, windows); + } + } + Value::Object(object) => collect_object(object, seen, windows), + _ => {} + } +} + +fn collect_object(object: &Object, seen: &mut HashSet, windows: &mut Vec) { + if let Some(window) = quota(object, None, None) + && seen.insert(window.raw_key.clone()) + { + windows.push(window); + } + for value in object.values() { + collect_windows(value, seen, windows); + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum QuotaKind { + Rolling, + Monthly, +} + +fn quota_kind(window: &QuotaWindow) -> Option { + let label = format!("{} {}", window.label.as_deref().unwrap_or(""), window.unit).to_lowercase(); + let minutes = window.rate.window_minutes; + if ["rolling", "4h", "4 h", "4-hour", "four hour", "four-hour"] + .iter() + .any(|needle| label.contains(needle)) + || minutes == Some(ROLLING_MINUTES) + { + return Some(QuotaKind::Rolling); + } + if ["month", "billing", "subscription"] + .iter() + .any(|needle| label.contains(needle)) + || minutes.unwrap_or(0) >= MONTHLY_MIN_WINDOW_MINUTES + { + return Some(QuotaKind::Monthly); + } + None +} + +/// Objects searched, in order, for subscription-level fields. +struct Sources<'a> { + root: &'a Object, + data: &'a Object, + subscription: &'a Object, +} + +impl Sources<'_> { + fn context(&self, keys: &[&str], convert: fn(Option<&Value>) -> Option) -> Option { + convert(lookup(self.root, keys)) + .or_else(|| convert(lookup(self.data, keys))) + .or_else(|| convert(lookup(self.subscription, keys))) + } +} + +/// Parse a Chutes usage or quota payload (upstream `chutesParse`). +pub(super) fn parse(value: &Value) -> ParsedQuotas { + let wrapped; + let root = match value { + Value::Object(object) => object, + Value::Array(_) => { + wrapped = Object::from_iter([("quotas".to_string(), value.clone())]); + &wrapped + } + _ => { + wrapped = Object::new(); + &wrapped + } + }; + let data = as_object(lookup(root, WRAPPER_KEYS)).unwrap_or(root); + let dictionary = + |keys: &[&str]| as_object(lookup(root, keys)).or(as_object(lookup(data, keys))); + let empty = Object::new(); + let sources = Sources { + root, + data, + subscription: dictionary(SUBSCRIPTION_OBJECT_KEYS).unwrap_or(&empty), + }; + + let mut seen = HashSet::new(); + let mut windows = Vec::new(); + for container in [ + lookup(root, QUOTA_CONTAINER_KEYS), + lookup(data, QUOTA_CONTAINER_KEYS), + ] + .into_iter() + .flatten() + { + collect_windows(container, &mut seen, &mut windows); + } + collect_object(data, &mut seen, &mut windows); + collect_object(root, &mut seen, &mut windows); + + let select = |payload_keys: &[&str], label: &str, minutes: u32, kind: QuotaKind| { + dictionary(payload_keys) + .and_then(|payload| quota(payload, Some(label), Some(minutes))) + .or_else(|| { + windows + .iter() + .find(|window| quota_kind(window) == Some(kind)) + .cloned() + }) + }; + let rolling = select( + ROLLING_PAYLOAD_KEYS, + "4-hour quota", + ROLLING_MINUTES, + QuotaKind::Rolling, + ); + let monthly = select( + MONTHLY_PAYLOAD_KEYS, + "Monthly quota", + MONTHLY_MINUTES, + QuotaKind::Monthly, + ); + + let mut active = sources.context(ACTIVE_KEYS, flag); + if active.is_none() { + let status = sources + .context(STATUS_KEYS, text) + .unwrap_or_default() + .to_lowercase(); + if status.contains("active") && !status.contains("inactive") { + active = Some(true); + } else if ["free", "inactive", "cancel", "none", "expired"] + .iter() + .any(|needle| status.contains(needle)) + { + active = Some(false); + } + } + + let is_selected = |window: &QuotaWindow| { + [&rolling, &monthly] + .into_iter() + .flatten() + .any(|selected| selected.raw_key == window.raw_key) + }; + let fallback = windows + .iter() + .filter(|window| !is_selected(window)) + .cloned() + .collect(); + ParsedQuotas { + plan: sources.context(PLAN_KEYS, text), + renewal: sources.context(RESET_KEYS, date), + rolling, + monthly, + fallback, + active, + } +} diff --git a/rust/src/providers/chutes/tests.rs b/rust/src/providers/chutes/tests.rs new file mode 100644 index 0000000000..592eb1c316 --- /dev/null +++ b/rust/src/providers/chutes/tests.rs @@ -0,0 +1,754 @@ +use chrono::{TimeZone, Utc}; +use mockito::{Matcher, Mock, ServerGuard}; +use serde_json::json; + +use super::parse::{format_quota_amount, snapshot_from_usage}; +use super::*; + +fn utc(year: i32, month: u32, day: u32, hour: u32) -> chrono::DateTime { + Utc.with_ymd_and_hms(year, month, day, hour, 0, 0).unwrap() +} + +fn secondary(snapshot: &UsageSnapshot) -> &crate::core::RateWindow { + snapshot + .secondary + .as_ref() + .expect("secondary window should be present") +} + +// --- parsing --------------------------------------------------------------- + +#[test] +fn parses_ratio_window() { + let snapshot = snapshot_from_usage(&json!({"quotas":[{"used":25,"limit":100}]})); + assert_eq!(snapshot.primary.used_percent, 25.0); +} + +#[test] +fn quota_counts_go_to_reset_description_not_schedule() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [{"used": 25, "limit": 100, "unit": "credits", "remaining": 75}] + })); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("25/100 credits") + ); + assert!(snapshot.primary.resets_at.is_none()); +} + +#[test] +fn small_numeric_fields_are_not_reset_schedules() { + // A mis-keyed payload where "reset_time" is accidentally a remaining count. + // Values below unix-epoch range must not become resets_at. + let snapshot = snapshot_from_usage(&json!({ + "quotas": [{"used": 10, "limit": 50, "reset_time": 40}] + })); + assert_eq!(snapshot.primary.used_percent, 20.0); + assert!(snapshot.primary.resets_at.is_none()); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("10/50 credits") + ); +} + +#[test] +fn iso_reset_timestamps_still_parse() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [{"used": 1, "limit": 2, "resets_at": "2026-08-01T00:00:00Z"}] + })); + assert_eq!(snapshot.primary.resets_at, Some(utc(2026, 8, 1, 0))); +} + +#[test] +fn whole_percent_values_are_never_rescaled_as_fractions() { + for (percent, expected) in [(1.0, 1.0), (0.5, 0.5), (100.0, 100.0)] { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [{"usage_percent": percent, "limit": 100}] + })); + assert_eq!(snapshot.primary.used_percent, expected); + } +} + +#[test] +fn exact_percent_of_one_stays_one_percent() { + let used = snapshot_from_usage(&json!({"rolling_window": {"usage_percent": 1}})); + let remaining = snapshot_from_usage(&json!({"rolling_window": {"percent_remaining": 1}})); + assert_eq!(used.primary.used_percent, 1.0); + assert_eq!(remaining.primary.used_percent, 99.0); +} + +#[test] +fn large_quota_amounts_keep_their_description() { + let snapshot = snapshot_from_usage(&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 very_large_amounts_keep_the_represented_integer_digits() { + let snapshot = snapshot_from_usage(&json!({"rolling": {"used": 1e25, "limit": 2e25}})); + assert_eq!(snapshot.primary.used_percent, 50.0); + assert_eq!( + snapshot.primary.reset_description, + Some(format!("{:.0}/{:.0} credits", 1e25, 2e25)) + ); +} + +#[test] +fn large_integer_amounts_preserve_native_digits() { + let body: serde_json::Value = serde_json::from_str( + r#"{"rolling":{"used":1234567890123456789,"limit":2469135780246913578}}"#, + ) + .unwrap(); + let snapshot = snapshot_from_usage(&body); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("1234567890123456768/2469135780246913536 credits") + ); +} + +#[test] +fn fractional_amounts_round_like_native_printf() { + for (value, expected) in [ + (1.125, "1.12"), + (1.375, "1.38"), + (1.625, "1.62"), + (-1.125, "-1.12"), + (49.585, "49.59"), + ] { + let snapshot = snapshot_from_usage(&json!({"rolling": {"used": value, "limit": 100}})); + assert_eq!( + snapshot.primary.reset_description, + Some(format!("{expected}/100 credits")) + ); + } +} + +#[test] +fn duration_fields_populate_window_minutes() { + let snapshot = snapshot_from_usage(&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!(secondary(&snapshot).window_minutes, Some(30)); +} + +#[test] +fn unrepresentable_durations_keep_usage_with_the_known_window_default() { + for (key, value) in [ + ("window_minutes", json!(9_223_372_036_854_775_808_u64)), + ("window_hours", json!("1e308")), + ("window_days", json!("1e308")), + ("window_seconds", json!("1e308")), + ("window", json!("1e308 minutes")), + ("window", json!("1e308 hours")), + ("window", json!("1e308 days")), + ("window", json!("1e308 months")), + ] { + let snapshot = snapshot_from_usage(&json!({ + "rolling_window": {"used": 25, "limit": 100, key: value} + })); + assert_eq!(snapshot.primary.used_percent, 25.0, "{key}"); + assert_eq!(snapshot.primary.window_minutes, Some(240), "{key}"); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("25/100 credits") + ); + } +} + +#[test] +fn unrepresentable_generic_quota_duration_remains_unknown() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [{"used": 25, "limit": 100, "window_minutes": 9_223_372_036_854_775_808_u64}] + })); + 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() { + assert_eq!(format_quota_amount(2_f64.powi(63)), "9223372036854775808"); +} + +#[test] +fn identical_usage_values_keep_distinct_quota_windows() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [ + {"used": 0, "limit": 100, "window_minutes": 240}, + {"used": 0, "limit": 100, "window_minutes": 43200} + ] + })); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.primary.window_minutes, Some(240)); + assert_eq!(secondary(&snapshot).used_percent, 0.0); + assert_eq!(secondary(&snapshot).window_minutes, Some(43_200)); +} + +#[test] +fn distinct_raw_quotas_with_equal_percentages_remain_separate() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [ + {"used": 25, "remaining": 75, "window_minutes": 240}, + {"used": 50, "remaining": 150, "window_minutes": 240} + ] + })); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!(secondary(&snapshot).used_percent, 25.0); + assert_eq!(snapshot.primary.reset_description, None); +} + +#[test] +fn active_subscription_maps_monthly_and_rolling_windows() { + let snapshot = snapshot_from_usage(&json!({ + "subscription": { + "active": true, + "plan_name": "Pro", + "current_period_end": "2026-07-01T00:00:00Z" + }, + "monthly": { + "used": 250, + "limit": 1000, + "resets_at": "2026-07-01T00:00:00Z", + "unit": "credits" + }, + "rolling_window": { + "requests": 40, + "limit": 100, + "window_minutes": 240, + "reset_at": "2026-06-13T18:00:00Z", + "unit": "requests" + } + })); + assert_eq!(snapshot.primary.used_percent, 40.0); + assert_eq!(snapshot.primary.window_minutes, Some(240)); + assert_eq!(snapshot.primary.resets_at, Some(utc(2026, 6, 13, 18))); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("40/100 requests") + ); + let monthly = secondary(&snapshot); + assert_eq!(monthly.used_percent, 25.0); + assert_eq!(monthly.window_minutes, Some(43_200)); + assert_eq!(monthly.resets_at, Some(utc(2026, 7, 1, 0))); + assert_eq!( + monthly.reset_description.as_deref(), + Some("250/1000 credits") + ); + assert_eq!( + snapshot.subscription.and_then(|s| s.renews_at), + Some(utc(2026, 7, 1, 0)) + ); + assert_eq!(snapshot.login_method.as_deref(), Some("Pro")); +} + +#[test] +fn monthly_reset_is_not_reported_as_a_subscription_renewal() { + let snapshot = snapshot_from_usage(&json!({ + "monthly": {"used": 1, "limit": 10, "resets_at": "2026-07-01T00:00:00Z"} + })); + assert!(snapshot.subscription.is_none()); + assert_eq!(secondary(&snapshot).resets_at, Some(utc(2026, 7, 1, 0))); +} + +#[test] +fn normalized_aliases_and_epoch_resets_preserve_quota_projection() { + let snapshot = snapshot_from_usage(&json!({"result": { + "plan_name": "Fixture", + "rolling_4h": {"used_percent": "25", "reset_at": "1800000000000"}, + "billing_period": {"remaining": "$1,500", "cap": "2,000", "duration": "1 month"} + }})); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!( + snapshot.primary.resets_at, + Utc.timestamp_opt(1_800_000_000, 0).single() + ); + let monthly = secondary(&snapshot); + assert_eq!(monthly.used_percent, 25.0); + assert_eq!( + monthly.reset_description.as_deref(), + Some("500/2000 credits") + ); + assert_eq!(monthly.window_minutes, Some(43_200)); + assert_eq!(snapshot.login_method.as_deref(), Some("Fixture")); +} + +#[test] +fn rolling_and_monthly_are_classified_by_label_or_window_length() { + for quotas in [ + json!([ + {"name": "Monthly credits", "used": 30, "limit": 100}, + {"name": "4-hour usage", "used": 10, "limit": 100} + ]), + json!([ + {"used": 30, "limit": 100, "window_minutes": 43200}, + {"used": 10, "limit": 100, "window_minutes": 240} + ]), + json!([ + {"used": 30, "limit": 100, "window_days": 28}, + {"used": 10, "limit": 100, "window": "4h"} + ]), + ] { + let snapshot = snapshot_from_usage(&json!({ "quotas": quotas })); + assert_eq!(snapshot.primary.used_percent, 10.0, "{quotas}"); + assert_eq!(snapshot.primary.window_minutes, Some(240), "{quotas}"); + assert_eq!(secondary(&snapshot).used_percent, 30.0, "{quotas}"); + assert!( + secondary(&snapshot).window_minutes >= Some(40_320), + "{quotas}" + ); + } +} + +#[test] +fn payload_lanes_default_their_window_minutes() { + let snapshot = snapshot_from_usage(&json!({ + "fourhour": {"used": 1, "limit": 10}, + "monthly": {"used": 2, "limit": 10} + })); + assert_eq!(snapshot.primary.window_minutes, Some(240)); + assert_eq!(secondary(&snapshot).window_minutes, Some(43_200)); +} + +#[test] +fn a_single_rolling_lane_is_not_repeated_as_the_secondary() { + let snapshot = snapshot_from_usage(&json!({"rolling": {"used": 1, "limit": 10}})); + assert_eq!(snapshot.primary.used_percent, 10.0); + assert!(snapshot.secondary.is_none()); +} + +#[test] +fn unclassified_windows_fill_primary_then_secondary() { + let snapshot = snapshot_from_usage(&json!({ + "quotas": [ + {"used": 1, "limit": 10}, + {"used": 2, "limit": 10}, + {"used": 3, "limit": 10} + ] + })); + assert_eq!(snapshot.primary.used_percent, 10.0); + assert_eq!(secondary(&snapshot).used_percent, 20.0); + + let with_rolling = snapshot_from_usage(&json!({ + "rolling": {"used": 5, "limit": 10}, + "other": {"used": 2, "limit": 10, "name": "extra"} + })); + assert_eq!(with_rolling.primary.used_percent, 50.0); + assert_eq!(secondary(&with_rolling).used_percent, 20.0); + + let with_monthly = snapshot_from_usage(&json!({ + "monthly": {"used": 5, "limit": 10}, + "other": {"used": 2, "limit": 10, "name": "extra"} + })); + assert!(with_monthly.primary.is_informational); + assert_eq!(secondary(&with_monthly).used_percent, 50.0); +} + +#[test] +fn missing_usage_fields_return_a_no_data_snapshot() { + let snapshot = snapshot_from_usage(&json!({ + "subscription": {"active": true}, + "unexpected": {"nested": true} + })); + assert!(snapshot.primary.is_informational); + assert!(snapshot.secondary.is_none()); + assert_eq!(snapshot.login_method, None); + + let unknown = snapshot_from_usage(&json!({})); + assert!(unknown.primary.is_informational); + assert_eq!(unknown.login_method.as_deref(), Some("No usage data")); +} + +#[test] +fn inactive_subscription_is_reported_and_plan_name_wins() { + for payload in [ + json!({"subscription": {"active": false}}), + json!({"subscription": {"active": "no"}}), + json!({"active": 0}), + json!({"status": "free"}), + json!({"data": {"subscription": {"status": "Cancelled"}}}), + json!({"state": "expired"}), + ] { + let snapshot = snapshot_from_usage(&payload); + assert_eq!( + snapshot.login_method.as_deref(), + Some("No active subscription"), + "{payload}" + ); + } + let active = snapshot_from_usage(&json!({"status": "Active"})); + assert_eq!(active.login_method, None); + let inactive_plan = snapshot_from_usage(&json!({"active": false, "tier": "Basic"})); + assert_eq!(inactive_plan.login_method.as_deref(), Some("Basic")); +} + +// --- endpoint construction and HTTP flow ------------------------------------- + +fn client() -> Client { + Client::builder() + .no_proxy() + .build() + .expect("the test client should build") +} + +async fn fetch(server: &ServerGuard, path_and_query: &str) -> Result { + let base = Url::parse(&format!("{}{path_and_query}", server.url())) + .expect("the mock server URL should be valid"); + fetch_usage_snapshot(&client(), &base, "chutes-key").await +} + +async fn respond( + server: &mut ServerGuard, + path: &str, + status: usize, + body: &str, + hits: usize, +) -> Mock { + server + .mock("GET", path) + .match_header("authorization", "Bearer chutes-key") + .match_header("accept", "application/json") + .with_status(status) + .with_body(body) + .expect(hits) + .create_async() + .await +} + +#[test] +fn endpoint_keeps_the_configured_path_and_query() { + let base = Url::parse("https://chutes.test/proxy/?tenant=fixture#ignored").unwrap(); + assert_eq!( + endpoint(&base, &["quota_usage", "a/b c"]).unwrap().as_str(), + "https://chutes.test/proxy/users/me/quota_usage/a%2Fb%20c?tenant=fixture" + ); + let bare = Url::parse("https://api.chutes.ai").unwrap(); + assert_eq!( + endpoint(&bare, &["subscription_usage"]).unwrap().as_str(), + "https://api.chutes.ai/users/me/subscription_usage" + ); +} + +#[test] +fn endpoint_override_must_be_https() { + assert!(crate::providers::validated_https_url("http://chutes.test", "Chutes API").is_err()); + let with_query = crate::providers::validated_https_url("chutes.test/p?t=1", "Chutes API") + .expect("a bare host is accepted as HTTPS"); + assert_eq!(with_query.as_str(), "https://chutes.test/p?t=1"); +} + +#[tokio::test] +async fn complete_subscription_usage_needs_a_single_request() { + let mut server = mockito::Server::new_async().await; + let usage = respond( + &mut server, + "/users/me/subscription_usage", + 200, + r#"{"subscription":{"active":true,"plan_name":"Pro"}, + "monthly":{"used":250,"limit":1000}, + "rolling_window":{"requests":40,"limit":100,"window_minutes":240}}"#, + 1, + ) + .await; + let quotas = respond(&mut server, "/users/me/quotas", 200, "{}", 0).await; + + let snapshot = fetch(&server, "").await.unwrap(); + + assert_eq!(snapshot.primary.used_percent, 40.0); + assert_eq!(secondary(&snapshot).used_percent, 25.0); + assert_eq!(snapshot.login_method.as_deref(), Some("Pro")); + usage.assert_async().await; + quotas.assert_async().await; +} + +#[tokio::test] +async fn inactive_subscription_falls_back_to_the_quotas_endpoint() { + let mut server = mockito::Server::new_async().await; + let mocks = [ + respond( + &mut server, + "/users/me/subscription_usage", + 200, + r#"{"subscription":{"active":false,"status":"free"}}"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quotas", + 200, + r#"[{"chute_id":"0","is_default":true,"quota":100}]"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quota_usage/0", + 200, + r#"{"quota":100,"used":10}"#, + 1, + ) + .await, + ]; + + let snapshot = fetch(&server, "").await.unwrap(); + + assert_eq!(snapshot.primary.used_percent, 10.0); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("10/100 credits") + ); + assert!(snapshot.secondary.is_none()); + assert_eq!( + snapshot.login_method.as_deref(), + Some("No active subscription") + ); + for mock in mocks { + mock.assert_async().await; + } +} + +#[tokio::test] +async fn wrapped_quota_list_fetches_per_quota_usage() { + let mut server = mockito::Server::new_async().await; + let mocks = [ + respond( + &mut server, + "/users/me/subscription_usage", + 200, + r#"{"subscription":{"active":false}}"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quotas", + 200, + r#"{"quotas":{"metadata":true},"data":[{"chute_id":"wrapped","quota":200}]}"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quota_usage/wrapped", + 200, + r#"{"data":{"quota":200,"used":50}}"#, + 1, + ) + .await, + ]; + + let snapshot = fetch(&server, "").await.unwrap(); + + assert_eq!(snapshot.primary.used_percent, 25.0); + for mock in mocks { + mock.assert_async().await; + } +} + +#[tokio::test] +async fn partial_subscription_usage_fills_the_missing_rolling_window() { + let mut server = mockito::Server::new_async().await; + let mocks = [ + respond( + &mut server, + "/users/me/subscription_usage", + 200, + r#"{"subscription":{"active":true,"plan_name":"Pro", + "current_period_end":"2026-07-01T00:00:00Z"}, + "monthly":{"used":250,"limit":1000,"unit":"credits"}}"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quotas", + 200, + r#"{"rolling_window":{"requests":40,"limit":100,"window_minutes":240, + "unit":"requests"}}"#, + 1, + ) + .await, + ]; + + let snapshot = fetch(&server, "").await.unwrap(); + + assert_eq!(snapshot.primary.used_percent, 40.0); + assert_eq!(snapshot.primary.window_minutes, Some(240)); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("40/100 requests") + ); + assert_eq!(secondary(&snapshot).used_percent, 25.0); + assert_eq!( + secondary(&snapshot).reset_description.as_deref(), + Some("250/1000 credits") + ); + assert_eq!(snapshot.login_method.as_deref(), Some("Pro")); + for mock in mocks { + mock.assert_async().await; + } +} + +#[tokio::test] +async fn quota_enrichment_is_best_effort_except_for_rejected_credentials() { + for failing in ["quotas", "quota_usage/fixture"] { + for status in [403_usize, 500] { + let mut server = mockito::Server::new_async().await; + let failing_path = format!("/users/me/{failing}"); + let quota_usage_hits = usize::from(failing != "quotas"); + let mocks = [ + respond( + &mut server, + "/users/me/subscription_usage", + 200, + r#"{"subscription":{"active":false}}"#, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quotas", + if failing == "quotas" { status } else { 200 }, + if failing == "quotas" { + "upstream error" + } else { + r#"{"quotas":[{"chute_id":"fixture","limit":100}]}"# + }, + 1, + ) + .await, + respond( + &mut server, + "/users/me/quota_usage/fixture", + status, + "upstream error", + quota_usage_hits, + ) + .await, + ]; + + let result = fetch(&server, "").await; + + if status == 403 { + assert!( + matches!(result, Err(ProviderError::AuthRequired)), + "{failing_path}" + ); + } else { + let snapshot = result.unwrap(); + assert!(snapshot.primary.is_informational, "{failing_path}"); + assert_eq!( + snapshot.login_method.as_deref(), + Some("No active subscription") + ); + } + for mock in mocks { + mock.assert_async().await; + } + } + } +} + +#[tokio::test] +async fn configured_base_path_and_query_survive_every_request() { + let mut server = mockito::Server::new_async().await; + let query = Matcher::UrlEncoded("tenant".into(), "fixture".into()); + let usage = server + .mock("GET", "/proxy/users/me/subscription_usage") + .match_query(query.clone()) + .with_status(200) + .with_body("{}") + .expect(1) + .create_async() + .await; + let quotas = server + .mock("GET", "/proxy/users/me/quotas") + .match_query(query) + .with_status(200) + .with_body("{}") + .expect(1) + .create_async() + .await; + + let snapshot = fetch(&server, "/proxy?tenant=fixture").await.unwrap(); + + assert_eq!(snapshot.login_method.as_deref(), Some("No usage data")); + usage.assert_async().await; + quotas.assert_async().await; +} + +#[tokio::test] +async fn invalid_subscription_json_fails_but_optional_invalid_json_does_not() { + let mut server = mockito::Server::new_async().await; + let _usage = respond( + &mut server, + "/users/me/subscription_usage", + 200, + "not JSON", + 1, + ) + .await; + assert!(matches!( + fetch(&server, "").await, + Err(ProviderError::Parse(_)) + )); + + let mut server = mockito::Server::new_async().await; + let _usage = respond(&mut server, "/users/me/subscription_usage", 200, "{}", 1).await; + let _quotas = respond(&mut server, "/users/me/quotas", 200, "not JSON", 1).await; + let snapshot = fetch(&server, "").await.unwrap(); + assert!(snapshot.primary.is_informational); +} + +#[tokio::test] +async fn non_container_json_and_http_errors_are_reported() { + let mut server = mockito::Server::new_async().await; + let _usage = respond(&mut server, "/users/me/subscription_usage", 200, "42", 1).await; + assert!(matches!( + fetch(&server, "").await, + Err(ProviderError::Parse(_)) + )); + + let mut server = mockito::Server::new_async().await; + let _usage = respond(&mut server, "/users/me/subscription_usage", 503, "", 1).await; + let error = fetch(&server, "").await.unwrap_err(); + assert!(matches!( + error, + ProviderError::Other(message) if message == "Chutes usage API error: HTTP 503" + )); +} + +#[tokio::test] +async fn rejected_key_on_subscription_usage_requires_auth() { + for status in [401, 403] { + let mut server = mockito::Server::new_async().await; + let _usage = respond( + &mut server, + "/users/me/subscription_usage", + status, + r#"{"detail":"unauthorized"}"#, + 1, + ) + .await; + assert!(matches!( + fetch(&server, "").await, + Err(ProviderError::AuthRequired) + )); + } +} From 0f51ad7afce48aad521948b2d123408164cfebdf Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:40:00 +0700 Subject: [PATCH 109/407] Port upstream 0.66.0: Perplexity session cookie retries and credit lane parity --- rust/src/providers/perplexity/cookies.rs | 128 ++++++ rust/src/providers/perplexity/mod.rs | 444 ++++++++++-------- rust/src/providers/perplexity/tests.rs | 556 +++++++++++++++++++++++ 3 files changed, 932 insertions(+), 196 deletions(-) create mode 100644 rust/src/providers/perplexity/cookies.rs create mode 100644 rust/src/providers/perplexity/tests.rs diff --git a/rust/src/providers/perplexity/cookies.rs b/rust/src/providers/perplexity/cookies.rs new file mode 100644 index 0000000000..b17ed0fe23 --- /dev/null +++ b/rust/src/providers/perplexity/cookies.rs @@ -0,0 +1,128 @@ +//! Perplexity session-cookie resolution. +//! +//! Mirrors upstream `perplexity.js` (`cookies()`) and `PerplexityCookieHeader`: +//! a raw value is turned into the exact `Cookie` header values worth sending, +//! in the order they should be tried. + +use std::collections::BTreeMap; + +use crate::providers::normalize_cookie_header; + +/// Session cookie names Perplexity uses, in upstream's preference order. +pub(super) const SESSION_COOKIE_NAMES: [&str; 4] = [ + "__Secure-authjs.session-token", + "authjs.session-token", + "__Secure-next-auth.session-token", + "next-auth.session-token", +]; + +/// Largest integer JavaScript represents exactly (`Number.isSafeInteger`). +const MAX_SAFE_INDEX: u64 = (1 << 53) - 1; + +struct Pair<'a> { + lowered_name: String, + name: &'a str, + value: &'a str, +} + +/// Request cookies to try for `raw`, most specific first. +/// +/// - A bare token (no `=` or `;`) is tried under every supported cookie name. +/// - A `Cookie` header yields at most one cookie: the first supported name +/// present, either whole or reassembled from `name.0`, `name.1`, ... chunks. +/// Every other cookie in the header is deliberately dropped. +pub(super) fn request_cookies(raw: &str) -> Vec { + let text = raw.trim(); + if text.is_empty() || text.chars().any(char::is_control) { + return Vec::new(); + } + if !text.contains('=') && !text.contains(';') { + return SESSION_COOKIE_NAMES + .iter() + .map(|name| format!("{name}={text}")) + .collect(); + } + + let Some(header) = normalize_cookie_header(text) else { + return Vec::new(); + }; + let pairs = parse_pairs(&header); + for expected in SESSION_COOKIE_NAMES { + let lowered = expected.to_lowercase(); + if let Some(direct) = pairs.iter().find(|pair| pair.lowered_name == lowered) { + return vec![format!("{}={}", direct.name, direct.value)]; + } + if let Some(cookie) = reassemble_chunks(&pairs, &lowered) { + return vec![cookie]; + } + } + Vec::new() +} + +/// `name=value` pairs with a nonempty name and value. A repeated name (case +/// insensitive) keeps its first position but takes the last value, matching +/// upstream's `Map.set`. +fn parse_pairs(header: &str) -> Vec> { + let mut pairs: Vec> = Vec::new(); + for part in header.split(';') { + let Some((name, value)) = part.split_once('=') else { + continue; + }; + let (name, value) = (name.trim(), value.trim()); + if name.is_empty() || value.is_empty() { + continue; + } + let lowered_name = name.to_lowercase(); + match pairs + .iter_mut() + .find(|pair| pair.lowered_name == lowered_name) + { + Some(existing) => { + existing.name = name; + existing.value = value; + } + None => pairs.push(Pair { + lowered_name, + name, + value, + }), + } + } + pairs +} + +/// Join `.0`, `.1`, ... into one cookie. Any gap, or a +/// missing chunk zero, rejects the whole cookie. +fn reassemble_chunks(pairs: &[Pair<'_>], lowered_expected: &str) -> Option { + let prefix = format!("{lowered_expected}."); + let mut chunks: BTreeMap> = BTreeMap::new(); + for pair in pairs { + let Some(suffix) = pair.lowered_name.strip_prefix(&prefix) else { + continue; + }; + if let Some(index) = chunk_index(suffix) { + chunks.insert(index, pair); + } + } + let count = u64::try_from(chunks.len()).ok()?; + if count == 0 || chunks.keys().copied().ne(0..count) { + return None; + } + + let first = chunks.values().next()?; + let base_name = &first.name[..first.name.rfind('.')?]; + let token: String = chunks.values().map(|chunk| chunk.value).collect(); + Some(format!("{base_name}={token}")) +} + +/// Chunk suffix as upstream accepts it: `^[+-]?\d+$` naming a nonnegative +/// safe integer (so `-0` and `+1` count, `-1` and `1e3` do not). +fn chunk_index(suffix: &str) -> Option { + let digits = suffix.strip_prefix(['+', '-']).unwrap_or(suffix); + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + let index: u64 = digits.parse().ok()?; + let negative = suffix.starts_with('-'); + ((!negative || index == 0) && index <= MAX_SAFE_INDEX).then_some(index) +} diff --git a/rust/src/providers/perplexity/mod.rs b/rust/src/providers/perplexity/mod.rs index 25057b50f6..6226c6f6cb 100644 --- a/rust/src/providers/perplexity/mod.rs +++ b/rust/src/providers/perplexity/mod.rs @@ -1,11 +1,20 @@ //! Perplexity provider implementation //! -//! Fetches credit/usage data from Perplexity's REST billing endpoint. -//! Uses browser cookies (`__Secure-next-auth.session-token`) for authentication. +//! Fetches credit/usage data from Perplexity's REST billing endpoint with a +//! web session cookie. Behavior follows upstream v0.66.0 `perplexity.js`: +//! browser sessions are tried one by one (a 401/403 rejects only that +//! session), then the environment secret, unless the cookie is manual. + +mod cookies; +#[cfg(test)] +mod tests; + +use std::collections::HashSet; +use std::fmt::Display; use async_trait::async_trait; -use chrono::{DateTime, TimeZone, Utc}; -use reqwest::Client; +use chrono::{DateTime, Local, TimeZone, Utc}; +use reqwest::{Client, StatusCode}; use serde::Deserialize; use crate::core::{ @@ -13,39 +22,52 @@ use crate::core::{ RateWindow, SourceMode, UsageSnapshot, }; +use cookies::request_cookies; + const CREDITS_URL: &str = "https://www.perplexity.ai/rest/billing/credits?version=2.18&source=default"; const ORIGIN: &str = "https://www.perplexity.ai"; const REFERER: &str = "https://www.perplexity.ai/account/usage"; const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"; - +const ENV_SESSION_TOKEN: &str = "PERPLEXITY_SESSION_TOKEN"; +const ENV_COOKIE: &str = "PERPLEXITY_COOKIE"; +/// Recurring credits below this many cents are a Pro plan, above it Max. +const MAX_PLAN_MIN_RECURRING_CENTS: f64 = 5000.0; +/// Shown in place of the recurring lane when only bonus or purchased credits +/// exist. `UsageSnapshot` always carries a primary window, so this stands in +/// for upstream's absent lane and is skipped by Automatic metric selection. +const NO_RECURRING_CREDITS: &str = "No recurring credits"; + +/// Every numeric field is required and finite (JSON cannot carry NaN), as in +/// upstream's validation. The camelCase aliases are accepted there too. #[derive(Debug, Deserialize)] struct CreditsResponse { - #[serde(default)] - balance_cents: f64, - #[serde(default)] - renewal_date_ts: Option, - #[serde(default)] + #[serde(rename = "balance_cents", alias = "balanceCents")] + _balance_cents: f64, + #[serde(alias = "renewalDateTs")] + renewal_date_ts: f64, + #[serde(alias = "currentPeriodPurchasedCents")] current_period_purchased_cents: f64, - #[serde(default)] + #[serde(alias = "creditGrants")] credit_grants: Vec, - #[serde(default)] + #[serde(alias = "totalUsageCents")] total_usage_cents: f64, } #[derive(Debug, Deserialize)] struct CreditGrant { - #[serde(default, rename = "type")] + #[serde(rename = "type")] grant_type: String, - #[serde(default)] + #[serde(alias = "amountCents")] amount_cents: f64, - #[serde(default)] + #[serde(default, alias = "expiresAtTs")] expires_at_ts: Option, } pub struct PerplexityProvider { metadata: ProviderMetadata, client: Client, + credits_url: String, } impl PerplexityProvider { @@ -68,6 +90,7 @@ impl PerplexityProvider { .timeout(std::time::Duration::from_secs(30)) .build() .unwrap_or_else(|_| Client::new()), + credits_url: CREDITS_URL.to_string(), } } @@ -84,108 +107,144 @@ impl PerplexityProvider { Utc.timestamp_opt(secs, 0).single() } - fn parse_response(resp: CreditsResponse) -> Result { - // Bucket grants by type. - let mut recurring_total = 0.0_f64; - let mut bonus_total = 0.0_f64; - let mut bonus_expiry: Option> = None; - let mut purchased_total = 0.0_f64; - - for grant in &resp.credit_grants { - match grant.grant_type.as_str() { - "recurring" => recurring_total += grant.amount_cents, - "promotional" | "bonus" => { - bonus_total += grant.amount_cents; - if let Some(ts) = grant.expires_at_ts.and_then(Self::ts_to_datetime) - && bonus_expiry.is_none_or(|cur| ts < cur) - { - bonus_expiry = Some(ts); - } - } - "purchased" => purchased_total += grant.amount_cents, - _ => {} - } - } - - if purchased_total <= 0.0 && resp.current_period_purchased_cents > 0.0 { - purchased_total = resp.current_period_purchased_cents; - } - - // Waterfall attribution: usage burns recurring first, then bonus, then purchased. - let mut remaining_usage = resp.total_usage_cents.max(0.0); - let recurring_used = remaining_usage.min(recurring_total); - remaining_usage -= recurring_used; - let bonus_used = remaining_usage.min(bonus_total); - remaining_usage -= bonus_used; - let purchased_used = remaining_usage.min(purchased_total); - - let pct = |used: f64, total: f64| -> f64 { - if !used.is_finite() || !total.is_finite() || total <= 0.0 { - 0.0 - } else { - ((used / total) * 100.0).clamp(0.0, 100.0) - } + /// Map the credit payload to usage lanes. + /// + /// Usage burns recurring credits first, then purchased, then promotional. + /// Only unexpired promotional grants count. Secondary (bonus) and tertiary + /// (purchased) lanes are always present and read 100% used when their pool + /// is empty; the primary lane is dropped when recurring credits are absent + /// but another pool exists, and reads `0/0 credits` at 100% when none does. + /// `expiry_tz` renders the promotional expiry date. + fn parse_response( + resp: &CreditsResponse, + now: DateTime, + expiry_tz: &Tz, + ) -> UsageSnapshot + where + Tz: TimeZone, + Tz::Offset: Display, + { + #[expect( + clippy::cast_precision_loss, + reason = "epoch milliseconds are far below 2^53" + )] + let now_seconds = now.timestamp_millis() as f64 / 1000.0; + let sum = |grants: &[&CreditGrant]| { + grants + .iter() + .fold(0.0, |total, grant| total + grant.amount_cents) + .max(0.0) + }; + let of_type = |kind: &str| -> Vec<&CreditGrant> { + resp.credit_grants + .iter() + .filter(|grant| grant.grant_type == kind) + .collect() }; - let renewal = resp.renewal_date_ts.and_then(Self::ts_to_datetime); - - let mut primary = RateWindow::new(pct(recurring_used, recurring_total)); - primary.resets_at = renewal; - primary.reset_description = Self::credit_description(recurring_used, recurring_total); - - let mut snapshot = UsageSnapshot::new(primary); + let recurring = sum(&of_type("recurring")); + let promo_grants: Vec<&CreditGrant> = of_type("promotional") + .into_iter() + .filter(|grant| grant.expires_at_ts.unwrap_or(f64::INFINITY) > now_seconds) + .collect(); + let promo = sum(&promo_grants); + // Purchased credits may sit in the top-level field, in `purchased` + // grants, or both; the larger figure avoids double counting. + let purchased = sum(&of_type("purchased")) + .max(resp.current_period_purchased_cents) + .max(0.0); + + let mut remaining = resp.total_usage_cents; + let recurring_used = remaining.min(recurring); + remaining -= recurring_used; + let purchased_used = remaining.min(purchased); + remaining -= purchased_used; + let promo_used = remaining.min(promo); + + let promo_expiry = promo_grants + .iter() + .filter_map(|grant| grant.expires_at_ts) + .min_by(f64::total_cmp) + .and_then(Self::ts_to_datetime); + let renewal = Self::ts_to_datetime(resp.renewal_date_ts); + + let primary = if recurring > 0.0 { + RateWindow::with_details( + Self::used_percent(recurring_used, recurring), + None, + renewal, + Self::credit_description(recurring_used, recurring, "credits"), + ) + } else if promo > 0.0 || purchased > 0.0 { + RateWindow::informational(NO_RECURRING_CREDITS) + } else { + RateWindow::with_details(100.0, None, renewal, Some("0/0 credits".to_string())) + }; - if bonus_total > 0.0 { - let mut secondary = RateWindow::new(pct(bonus_used, bonus_total)); - secondary.resets_at = bonus_expiry; - 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 = bonus_description; - snapshot = snapshot.with_secondary(secondary); + let mut promo_description = Self::credit_description(promo_used, promo, "bonus"); + if let (Some(expiry), Some(description)) = (promo_expiry, promo_description.as_mut()) { + description.push_str(&format!( + " · exp. {}", + expiry.with_timezone(expiry_tz).format("%b %-d") + )); } + let secondary = RateWindow::with_details( + Self::used_percent(promo_used, promo), + None, + None, + promo_description, + ); + let tertiary = RateWindow::with_details( + Self::used_percent(purchased_used, purchased), + None, + None, + Self::credit_description(purchased_used, purchased, "credits"), + ); - if purchased_total > 0.0 { - let mut tertiary = RateWindow::new(pct(purchased_used, purchased_total)); - tertiary.reset_description = Self::credit_description(purchased_used, purchased_total); - snapshot = snapshot.with_tertiary(tertiary); + let mut snapshot = UsageSnapshot::new(primary) + .with_secondary(secondary) + .with_tertiary(tertiary); + if recurring > 0.0 { + let plan = if recurring < MAX_PLAN_MIN_RECURRING_CENTS { + "Pro" + } else { + "Max" + }; + snapshot = snapshot.with_login_method(plan); } + snapshot + } - let plan = if recurring_total <= 0.0 { - None - } else if recurring_total < 5000.0 { - Some("Pro") + /// An empty pool reads 100% used so its bar renders empty, not full. + fn used_percent(used: f64, total: f64) -> f64 { + if total > 0.0 { + (used / total * 100.0).clamp(0.0, 100.0) } else { - Some("Max") - }; - if let Some(p) = plan { - snapshot = snapshot.with_login_method(p); + 100.0 } - - // Avoid unused-warning for balance_cents on some build configs. - let _ = resp.balance_cents; - - Ok(snapshot) } - fn credit_description(used: f64, total: f64) -> Option { + /// `used/total unit` as whole numbers: used rounds half away from zero, + /// total truncates, and a negative zero prints as `0`. Non-finite counts + /// have no description. + fn credit_description(used: f64, total: f64, unit: &str) -> Option { if !used.is_finite() || !total.is_finite() { return None; } - Some(format!("${:.2}/${:.2}", used / 100.0, total / 100.0)) + let whole = |value: f64| if value == 0.0 { 0.0 } else { value }; + let used = whole(used.round()); + let total = whole(total.trunc()); + Some(format!("{used:.0}/{total:.0} {unit}")) } - async fn fetch_with_cookies( - &self, - cookie_header: &str, - ) -> Result { + /// Request the credits payload with one session cookie. `Ok(None)` means + /// Perplexity rejected this cookie (401/403), so the next candidate may + /// still work. + async fn fetch_credits(&self, cookie: &str) -> Result, ProviderError> { let response = self .client - .get(CREDITS_URL) - .header("Cookie", cookie_header) + .get(&self.credits_url) + .header("Cookie", cookie) .header("Origin", ORIGIN) .header("Referer", REFERER) .header("User-Agent", USER_AGENT) @@ -194,24 +253,100 @@ impl PerplexityProvider { .await?; let status = response.status(); - if status.as_u16() == 401 || status.as_u16() == 403 { - return Err(ProviderError::AuthRequired); + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Ok(None); } - if !status.is_success() { + if status != StatusCode::OK { return Err(ProviderError::Other(format!( - "Perplexity API returned {}", - status + "Perplexity API error: HTTP {}", + status.as_u16() ))); } let body = response.text().await?; - let parsed: CreditsResponse = serde_json::from_str(&body) - .map_err(|e| ProviderError::Parse(format!("Failed to parse credits: {}", e)))?; + let parsed: CreditsResponse = serde_json::from_str(&body).map_err(|error| { + let reason = if error.is_data() { + "invalid credit fields" + } else { + "invalid JSON" + }; + ProviderError::Parse(format!("Failed to parse Perplexity usage data: {reason}")) + })?; - Self::parse_response(parsed) + Ok(Some(Self::parse_response(&parsed, Utc::now(), &Local))) + } + + /// Try each header's session cookies in order and return the first + /// accepted response. A rejected cookie moves on to the next one; any + /// other failure ends the attempt. + async fn fetch_first_accepted( + &self, + headers: &[String], + ) -> Result { + let mut attempted = HashSet::new(); + let mut rejected = false; + for header in headers { + for cookie in request_cookies(header) { + if !attempted.insert(cookie.clone()) { + continue; + } + match self.fetch_credits(&cookie).await? { + Some(usage) => return Ok(usage), + None => rejected = true, + } + } + } + Err(if rejected { + ProviderError::AuthRequired + } else { + ProviderError::NoCookies + }) + } + + async fn fetch_web(&self, ctx: &FetchContext) -> Result { + // A manual cookie is exclusive: no browser sessions, no environment. + if let Some(manual) = ctx.manual_cookie_header.as_deref() { + return self.fetch_first_accepted(&[manual.to_string()]).await; + } + + let (mut headers, browser_error) = + match crate::providers::browser_cookie_headers_for_domain("perplexity.ai") { + Ok(candidates) => ( + candidates.into_iter().map(|(_, header)| header).collect(), + None, + ), + Err(ProviderError::NoCookies) => (Vec::new(), None), + Err(error) => (Vec::new(), Some(error)), + }; + headers.extend(environment_cookie(|name| std::env::var(name).ok())); + + match self.fetch_first_accepted(&headers).await { + // An unreadable browser store is more useful than "no cookies" + // when nothing else produced a session. + Err(ProviderError::NoCookies) if browser_error.is_some() => { + Err(browser_error.unwrap_or(ProviderError::NoCookies)) + } + result => result, + } } } +/// Environment session, `PERPLEXITY_SESSION_TOKEN` before `PERPLEXITY_COOKIE`. +/// Blank values are ignored and one pair of surrounding quotes is removed. +fn environment_cookie(get: impl Fn(&str) -> Option) -> Option { + let cleaned = |name: &str| { + let value = get(name)?; + let value = value.trim(); + let unquoted = ['"', '\''] + .iter() + .find_map(|quote| value.strip_prefix(*quote)?.strip_suffix(*quote)) + .unwrap_or(value) + .trim(); + (!unquoted.is_empty()).then(|| unquoted.to_string()) + }; + cleaned(ENV_SESSION_TOKEN).or_else(|| cleaned(ENV_COOKIE)) +} + impl Default for PerplexityProvider { fn default() -> Self { Self::new() @@ -237,22 +372,8 @@ impl Provider for PerplexityProvider { match ctx.source_mode { SourceMode::Auto | SourceMode::Web => { - if let Some(ref cookie_header) = ctx.manual_cookie_header { - let usage = self.fetch_with_cookies(cookie_header).await?; - return Ok(ProviderFetchResult::new(usage, "web")); - } - - match crate::providers::browser_cookie_header(&["perplexity.ai"]) { - Ok(cookie_header) => match self.fetch_with_cookies(&cookie_header).await { - Ok(usage) => return Ok(ProviderFetchResult::new(usage, "web")), - Err(ProviderError::AuthRequired) => {} - Err(e) => return Err(e), - }, - Err(ProviderError::NoCookies) => {} - Err(e) => return Err(e), - } - - Err(ProviderError::AuthRequired) + let usage = self.fetch_web(ctx).await?; + Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::Cli => Err(ProviderError::UnsupportedSource(SourceMode::Cli)), SourceMode::OAuth => Err(ProviderError::UnsupportedSource(SourceMode::OAuth)), @@ -270,79 +391,10 @@ impl Provider for PerplexityProvider { fn supports_cli(&self) -> bool { false } -} -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_credits_with_waterfall_attribution() { - let resp: CreditsResponse = serde_json::from_value(serde_json::json!({ - "balance_cents": 450.0, - "renewal_date_ts": 1_750_000_000.0, - "current_period_purchased_cents": 0.0, - "credit_grants": [ - { "type": "recurring", "amount_cents": 500.0, "expires_at_ts": null }, - { "type": "promotional", "amount_cents": 100.0, "expires_at_ts": 1_750_000_000.0 }, - { "type": "purchased", "amount_cents": 0.0, "expires_at_ts": null } - ], - "total_usage_cents": 150.0 - })) - .unwrap(); - - let snap = PerplexityProvider::parse_response(resp).unwrap(); - // 150 cents of usage all attributed to recurring (500 cap). - assert!((snap.primary.used_percent - 30.0).abs() < 0.001); - let bonus = snap.secondary.expect("bonus window"); - assert!((bonus.used_percent - 0.0).abs() < 0.001); - assert!( - bonus - .reset_description - .as_deref() - .is_some_and(|description| description.contains("exp. 2025-06-15")) - ); - assert_eq!(snap.login_method.as_deref(), Some("Pro")); - } - - #[test] - fn classifies_max_plan_when_recurring_high() { - let resp: CreditsResponse = serde_json::from_value(serde_json::json!({ - "balance_cents": 5000.0, - "renewal_date_ts": null, - "current_period_purchased_cents": 0.0, - "credit_grants": [ - { "type": "recurring", "amount_cents": 6000.0, "expires_at_ts": null } - ], - "total_usage_cents": 0.0 - })) - .unwrap(); - 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()); + /// Browser sessions are tried one by one inside the provider, so the shell + /// must not pre-merge a single browser cookie header for Auto. + fn owns_browser_cookie_resolution(&self) -> bool { + true } } diff --git a/rust/src/providers/perplexity/tests.rs b/rust/src/providers/perplexity/tests.rs new file mode 100644 index 0000000000..5cf75844eb --- /dev/null +++ b/rust/src/providers/perplexity/tests.rs @@ -0,0 +1,556 @@ +use chrono::FixedOffset; +use serde_json::json; + +use super::cookies::{SESSION_COOKIE_NAMES, request_cookies}; +use super::*; + +// Fixture epochs: 2030-01-01T00:00:00Z renewal, promo expiries around it. +const NOW: i64 = 1_750_000_000; // 2025-06-15T15:06:40Z +const RENEWAL: i64 = 1_893_456_000; + +fn now() -> DateTime { + Utc.timestamp_opt(NOW, 0).single().expect("valid epoch") +} + +fn response(value: serde_json::Value) -> CreditsResponse { + serde_json::from_value(value).expect("valid credits payload") +} + +fn snapshot(value: serde_json::Value) -> UsageSnapshot { + PerplexityProvider::parse_response(&response(value), now(), &Utc) +} + +fn description(window: &RateWindow) -> &str { + window.reset_description.as_deref().unwrap_or_default() +} + +fn body(usage: f64) -> String { + json!({ + "balance_cents": 500, + "renewal_date_ts": RENEWAL, + "current_period_purchased_cents": 0, + "credit_grants": [{ "type": "recurring", "amount_cents": 1000 }], + "total_usage_cents": usage, + }) + .to_string() +} + +// ---- cookies ---- + +#[test] +fn bare_token_is_tried_under_every_session_cookie_name() { + let cookies = request_cookies(" manual-fixture "); + let expected: Vec = SESSION_COOKIE_NAMES + .iter() + .map(|name| format!("{name}=manual-fixture")) + .collect(); + assert_eq!(cookies, expected); +} + +#[test] +fn cookie_header_sends_only_the_session_cookie() { + let cookies = request_cookies("Cookie: theme=dark; next-auth.session-token=abc; other=1"); + assert_eq!(cookies, vec!["next-auth.session-token=abc"]); +} + +#[test] +fn earlier_supported_names_win_regardless_of_header_order() { + let cookies = request_cookies("next-auth.session-token=old; __Secure-authjs.session-token=new"); + assert_eq!(cookies, vec!["__Secure-authjs.session-token=new"]); +} + +#[test] +fn chunked_cookie_is_reassembled_in_index_order() { + let cookies = request_cookies( + "authjs.session-token.1=second; ignored=fixture; authjs.session-token.0=first", + ); + assert_eq!(cookies, vec!["authjs.session-token=firstsecond"]); +} + +#[test] +fn chunk_gap_or_missing_zero_yields_no_cookie() { + assert!( + request_cookies("authjs.session-token.0=first; authjs.session-token.2=third").is_empty() + ); + assert!(request_cookies("authjs.session-token.1=second").is_empty()); +} + +#[test] +fn header_without_a_session_cookie_yields_nothing() { + assert!(request_cookies("theme=dark; other=1").is_empty()); + assert!(request_cookies("").is_empty()); + assert!(request_cookies("a=b\r\nc=d").is_empty()); +} + +#[test] +fn cookie_names_match_case_insensitively_and_keep_header_casing() { + let cookies = request_cookies("Authjs.Session-Token=abc"); + assert_eq!(cookies, vec!["Authjs.Session-Token=abc"]); +} + +#[test] +fn empty_values_are_ignored_and_repeats_keep_the_last_value() { + assert!(request_cookies("authjs.session-token=").is_empty()); + let cookies = request_cookies("authjs.session-token=first; authjs.session-token=second"); + assert_eq!(cookies, vec!["authjs.session-token=second"]); +} + +#[test] +fn chunk_suffix_follows_upstream_integer_rules() { + // `+1` and `-0` count as indices 1 and 0; `-1` and `1e3` do not. + let cookies = request_cookies("authjs.session-token.-0=a; authjs.session-token.+1=b"); + assert_eq!(cookies, vec!["authjs.session-token=ab"]); + assert!(request_cookies("authjs.session-token.-1=a").is_empty()); + assert!(request_cookies("authjs.session-token.1e3=a").is_empty()); +} + +// ---- environment secret ---- + +#[test] +fn environment_session_token_wins_over_cookie_and_strips_quotes() { + let env = |name: &str| match name { + "PERPLEXITY_SESSION_TOKEN" => Some(" \"token-value\" ".to_string()), + "PERPLEXITY_COOKIE" => Some("authjs.session-token=other".to_string()), + _ => None, + }; + assert_eq!(environment_cookie(env).as_deref(), Some("token-value")); +} + +#[test] +fn environment_falls_back_to_cookie_when_token_is_blank() { + let env = |name: &str| match name { + "PERPLEXITY_SESSION_TOKEN" => Some(" ".to_string()), + "PERPLEXITY_COOKIE" => Some("'authjs.session-token=abc'".to_string()), + _ => None, + }; + assert_eq!( + environment_cookie(env).as_deref(), + Some("authjs.session-token=abc") + ); + assert_eq!(environment_cookie(|_| None), None); +} + +// ---- payload validation ---- + +#[test] +fn camel_case_payload_is_accepted() { + let parsed: Result = serde_json::from_value(json!({ + "balanceCents": 1, + "renewalDateTs": RENEWAL, + "currentPeriodPurchasedCents": 0, + "creditGrants": [{ "type": "recurring", "amountCents": 100, "expiresAtTs": null }], + "totalUsageCents": 10, + })); + assert!(parsed.is_ok()); +} + +#[test] +fn missing_or_non_numeric_fields_are_parse_failures() { + let complete = json!({ + "balance_cents": 1, + "renewal_date_ts": RENEWAL, + "current_period_purchased_cents": 0, + "credit_grants": [], + "total_usage_cents": 0, + }); + for field in [ + "balance_cents", + "renewal_date_ts", + "current_period_purchased_cents", + "credit_grants", + "total_usage_cents", + ] { + let mut value = complete.clone(); + value.as_object_mut().expect("object").remove(field); + assert!( + serde_json::from_value::(value).is_err(), + "missing {field} must fail" + ); + } + let mut string_number = complete.clone(); + string_number["total_usage_cents"] = json!("10"); + assert!(serde_json::from_value::(string_number).is_err()); + + let mut bad_grant = complete; + bad_grant["credit_grants"] = json!([{ "type": "recurring" }]); + assert!(serde_json::from_value::(bad_grant).is_err()); +} + +// ---- lanes ---- + +fn base(grants: serde_json::Value, usage: f64, purchased: f64) -> serde_json::Value { + json!({ + "balance_cents": 0, + "renewal_date_ts": RENEWAL, + "current_period_purchased_cents": purchased, + "credit_grants": grants, + "total_usage_cents": usage, + }) +} + +#[test] +fn recurring_credits_fill_the_primary_lane_first() { + let snap = snapshot(base( + json!([ + { "type": "recurring", "amount_cents": 500 }, + { "type": "promotional", "amount_cents": 100, "expires_at_ts": RENEWAL }, + { "type": "purchased", "amount_cents": 200 }, + ]), + 150.0, + 0.0, + )); + assert!((snap.primary.used_percent - 30.0).abs() < 1e-9); + assert_eq!(description(&snap.primary), "150/500 credits"); + assert_eq!( + snap.primary.resets_at, + Utc.timestamp_opt(RENEWAL, 0).single() + ); + let bonus = snap.secondary.expect("secondary lane"); + assert!(bonus.used_percent.abs() < 1e-9); + assert!(description(&bonus).starts_with("0/100 bonus")); + let purchased = snap.tertiary.expect("tertiary lane"); + assert!(purchased.used_percent.abs() < 1e-9); + assert_eq!(description(&purchased), "0/200 credits"); + assert_eq!(snap.login_method.as_deref(), Some("Pro")); +} + +#[test] +fn waterfall_consumes_recurring_then_purchased_then_promotional() { + let snap = snapshot(base( + json!([ + { "type": "recurring", "amount_cents": 100 }, + { "type": "promotional", "amount_cents": 300, "expires_at_ts": RENEWAL }, + { "type": "purchased", "amount_cents": 200 }, + ]), + 450.0, + 0.0, + )); + assert_eq!(description(&snap.primary), "100/100 credits"); + assert_eq!( + description(&snap.tertiary.expect("purchased")), + "200/200 credits" + ); + let bonus = snap.secondary.expect("bonus"); + assert!(description(&bonus).starts_with("150/300 bonus")); + assert!((bonus.used_percent - 50.0).abs() < 1e-9); +} + +#[test] +fn purchased_uses_the_larger_of_grants_and_period_field() { + let from_field = snapshot(base(json!([]), 0.0, 700.0)); + assert_eq!( + description(&from_field.tertiary.expect("tertiary")), + "0/700 credits" + ); + let from_grants = snapshot(base( + json!([{ "type": "purchased", "amount_cents": 900 }]), + 0.0, + 700.0, + )); + assert_eq!( + description(&from_grants.tertiary.expect("tertiary")), + "0/900 credits" + ); +} + +#[test] +fn only_unexpired_promotional_grants_count_and_bonus_type_is_ignored() { + let snap = snapshot(base( + json!([ + { "type": "recurring", "amount_cents": 100 }, + { "type": "promotional", "amount_cents": 50, "expires_at_ts": NOW - 1 }, + { "type": "promotional", "amount_cents": 70, "expires_at_ts": NOW + 86_400 * 3 }, + { "type": "promotional", "amount_cents": 30, "expires_at_ts": null }, + { "type": "bonus", "amount_cents": 999, "expires_at_ts": NOW + 10 }, + ]), + 0.0, + 0.0, + )); + // 70 + 30 = 100; the expired and `bonus`-typed grants are excluded. + let bonus = snap.secondary.expect("bonus"); + assert_eq!(description(&bonus), "0/100 bonus · exp. Jun 18"); +} + +#[test] +fn promotional_expiry_uses_the_earliest_unexpired_grant_in_the_given_zone() { + let value = base( + json!([ + { "type": "recurring", "amount_cents": 100 }, + { "type": "promotional", "amount_cents": 10, "expires_at_ts": 1_750_550_400 }, + { "type": "promotional", "amount_cents": 10, "expires_at_ts": 1_751_000_000 }, + ]), + 0.0, + 0.0, + ); + // 1_750_550_400 is 2025-06-22T00:00:00Z. + let utc = PerplexityProvider::parse_response(&response(value.clone()), now(), &Utc); + assert!(description(&utc.secondary.expect("bonus")).ends_with("exp. Jun 22")); + let west = FixedOffset::west_opt(5 * 3600).expect("offset"); + let local = PerplexityProvider::parse_response(&response(value), now(), &west); + assert!(description(&local.secondary.expect("bonus")).ends_with("exp. Jun 21")); +} + +#[test] +fn promotional_without_expiry_has_no_suffix() { + let snap = snapshot(base( + json!([{ "type": "promotional", "amount_cents": 40 }]), + 0.0, + 0.0, + )); + assert_eq!(description(&snap.secondary.expect("bonus")), "0/40 bonus"); +} + +#[test] +fn primary_is_a_placeholder_when_only_promotional_or_purchased_exist() { + let snap = snapshot(base( + json!([{ "type": "promotional", "amount_cents": 40 }]), + 10.0, + 0.0, + )); + assert!(snap.primary.is_informational); + assert_eq!(description(&snap.primary), NO_RECURRING_CREDITS); + assert!(snap.login_method.is_none()); + + let purchased_only = snapshot(base(json!([]), 0.0, 100.0)); + assert!(purchased_only.primary.is_informational); +} + +#[test] +fn empty_account_reads_zero_of_zero_at_full_usage() { + let snap = snapshot(base(json!([]), 0.0, 0.0)); + assert!((snap.primary.used_percent - 100.0).abs() < 1e-9); + assert_eq!(description(&snap.primary), "0/0 credits"); + let bonus = snap.secondary.expect("secondary always present"); + assert!((bonus.used_percent - 100.0).abs() < 1e-9); + assert_eq!(description(&bonus), "0/0 bonus"); + let purchased = snap.tertiary.expect("tertiary always present"); + assert!((purchased.used_percent - 100.0).abs() < 1e-9); + assert_eq!(description(&purchased), "0/0 credits"); + assert!(snap.login_method.is_none()); +} + +#[test] +fn plan_is_pro_below_5000_recurring_cents_and_max_from_5000() { + let plan = |cents: f64| { + snapshot(base( + json!([{ "type": "recurring", "amount_cents": cents }]), + 0.0, + 0.0, + )) + .login_method + }; + assert_eq!(plan(4999.0).as_deref(), Some("Pro")); + assert_eq!(plan(5000.0).as_deref(), Some("Max")); +} + +#[test] +fn usage_beyond_all_pools_clamps_each_lane_at_full() { + let snap = snapshot(base( + json!([{ "type": "recurring", "amount_cents": 100 }]), + 10_000.0, + 0.0, + )); + assert!((snap.primary.used_percent - 100.0).abs() < 1e-9); + assert_eq!(description(&snap.primary), "100/100 credits"); +} + +#[test] +fn negative_usage_is_shown_as_reported_and_negative_grants_floor_at_zero() { + let snap = snapshot(base( + json!([ + { "type": "recurring", "amount_cents": 100 }, + { "type": "promotional", "amount_cents": -50, "expires_at_ts": RENEWAL }, + ]), + -20.0, + 0.0, + )); + assert_eq!(description(&snap.primary), "-20/100 credits"); + assert!(snap.primary.used_percent.abs() < 1e-9); + let bonus = snap.secondary.expect("bonus"); + assert!(description(&bonus).starts_with("0/0 bonus")); +} + +#[test] +fn descriptions_round_used_and_truncate_total() { + assert_eq!( + PerplexityProvider::credit_description(2.5, 9.9, "credits").as_deref(), + Some("3/9 credits") + ); + assert_eq!( + PerplexityProvider::credit_description(-0.2, 1.0, "credits").as_deref(), + Some("0/1 credits") + ); + assert!(PerplexityProvider::credit_description(f64::INFINITY, 1.0, "credits").is_none()); +} + +#[test] +fn oversized_credit_totals_do_not_render_non_finite_descriptions() { + let snap = snapshot(base( + json!([ + { "type": "recurring", "amount_cents": f64::MAX }, + { "type": "recurring", "amount_cents": f64::MAX }, + ]), + 0.0, + 0.0, + )); + assert!(snap.primary.reset_description.is_none()); +} + +// ---- HTTP ---- + +fn provider_for(server: &mockito::ServerGuard) -> PerplexityProvider { + let mut provider = PerplexityProvider::new(); + provider.credits_url = format!("{}/rest/billing/credits", server.url()); + provider +} + +#[tokio::test] +async fn rejected_cookie_moves_on_and_the_accepted_one_is_used() { + let mut server = mockito::Server::new_async().await; + let rejected = server + .mock("GET", "/rest/billing/credits") + .match_header("cookie", "__Secure-authjs.session-token=bare") + .with_status(401) + .create_async() + .await; + let accepted = server + .mock("GET", "/rest/billing/credits") + .match_header("cookie", "authjs.session-token=bare") + .match_header("origin", ORIGIN) + .match_header("referer", REFERER) + .with_status(200) + .with_body(body(500.0)) + .create_async() + .await; + let provider = provider_for(&server); + + let usage = provider + .fetch_first_accepted(&["bare".to_string()]) + .await + .expect("second cookie name accepted"); + + rejected.assert_async().await; + accepted.assert_async().await; + assert!((usage.primary.used_percent - 50.0).abs() < 1e-9); +} + +#[tokio::test] +async fn later_header_is_tried_after_an_earlier_one_is_rejected() { + let mut server = mockito::Server::new_async().await; + let first = server + .mock("GET", "/rest/billing/credits") + .match_header("cookie", "authjs.session-token=expired") + .with_status(403) + .create_async() + .await; + let second = server + .mock("GET", "/rest/billing/credits") + .match_header("cookie", "authjs.session-token=fresh") + .with_status(200) + .with_body(body(0.0)) + .create_async() + .await; + let provider = provider_for(&server); + + let headers = vec![ + "authjs.session-token=expired; theme=dark".to_string(), + "authjs.session-token=fresh".to_string(), + ]; + provider + .fetch_first_accepted(&headers) + .await + .expect("fresh session accepted"); + + first.assert_async().await; + second.assert_async().await; +} + +#[tokio::test] +async fn identical_cookies_are_requested_once() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/rest/billing/credits") + .match_header("cookie", "authjs.session-token=same") + .with_status(401) + .expect(1) + .create_async() + .await; + let provider = provider_for(&server); + + let headers = vec![ + "authjs.session-token=same".to_string(), + "theme=dark; authjs.session-token=same".to_string(), + ]; + let error = provider + .fetch_first_accepted(&headers) + .await + .expect_err("rejected"); + + mock.assert_async().await; + assert!(matches!(error, ProviderError::AuthRequired)); +} + +#[tokio::test] +async fn no_usable_cookie_reports_missing_cookies_without_a_request() { + let server = mockito::Server::new_async().await; + let provider = provider_for(&server); + + let error = provider + .fetch_first_accepted(&["theme=dark".to_string()]) + .await + .expect_err("nothing to send"); + + assert!(matches!(error, ProviderError::NoCookies)); +} + +#[tokio::test] +async fn server_errors_are_terminal_and_do_not_try_more_cookies() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/rest/billing/credits") + .with_status(500) + .expect(1) + .create_async() + .await; + let provider = provider_for(&server); + + let error = provider + .fetch_first_accepted(&["bare".to_string()]) + .await + .expect_err("HTTP 500"); + + mock.assert_async().await; + assert!(error.to_string().contains("HTTP 500")); +} + +#[tokio::test] +async fn invalid_payloads_are_parse_failures_with_a_friendly_reason() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/rest/billing/credits") + .with_status(200) + .with_body(r#"{"balance_cents": 1}"#) + .create_async() + .await; + let provider = provider_for(&server); + + let error = provider + .fetch_credits("authjs.session-token=x") + .await + .expect_err("missing fields"); + + mock.assert_async().await; + assert!(error.to_string().contains("invalid credit fields")); + + let mut html = mockito::Server::new_async().await; + html.mock("GET", "/rest/billing/credits") + .with_status(200) + .with_body("") + .create_async() + .await; + let error = provider_for(&html) + .fetch_credits("authjs.session-token=x") + .await + .expect_err("not json"); + assert!(error.to_string().contains("invalid JSON")); +} From b8f85c1583fe9aaaaaf1477902e5f42db6539f88 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:43:02 +0700 Subject: [PATCH 110/407] Port upstream 0.66.0: bind Qoder cookies to their regional origin --- rust/src/providers/qoder/mod.rs | 328 ++++++++-- rust/src/providers/qoder/routing.rs | 713 ++++++++++++++++++++++ rust/src/providers/qoder/routing/tests.rs | 451 ++++++++++++++ 3 files changed, 1426 insertions(+), 66 deletions(-) create mode 100644 rust/src/providers/qoder/routing.rs create mode 100644 rust/src/providers/qoder/routing/tests.rs diff --git a/rust/src/providers/qoder/mod.rs b/rust/src/providers/qoder/mod.rs index a7478375d2..ed91b7c472 100644 --- a/rust/src/providers/qoder/mod.rs +++ b/rust/src/providers/qoder/mod.rs @@ -1,21 +1,29 @@ //! Qoder provider implementation. //! -//! Uses browser/manual cookies against the global or China Qoder usage API. +//! Uses browser/manual cookies against the international (`qoder.com`) or +//! China (`qoder.com.cn`) usage API. Every cookie is bound to the origin it +//! belongs to; see `routing`. + +mod routing; use async_trait::async_trait; use chrono::{DateTime, Utc}; use reqwest::Client; +use reqwest::header::HeaderValue; use serde_json::Value; +use std::time::Duration; use crate::core::{ FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; -const GLOBAL_API: &str = "https://qoder.com/api/v2/me/usages/big_model_credits"; -const CHINA_API: &str = "https://qoder.com.cn/api/v2/me/usages/big_model_credits"; +const BX_VERSION: &str = "2.5.35"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"; +use routing::QoderSite; + pub struct QoderProvider { metadata: ProviderMetadata, client: Client, @@ -38,75 +46,155 @@ impl QoderProvider { tertiary_label_key: None, }, client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(15)) + .timeout(REQUEST_TIMEOUT) .build() .unwrap_or_else(|_| Client::new()), } } - async fn fetch_web(&self, cookie_header: &str) -> Result { - let cookie = normalize_cookie_header(cookie_header).ok_or(ProviderError::NoCookies)?; - let regions = [ - (GLOBAL_API, "https://qoder.com/account/usage", "Qoder"), - ( - CHINA_API, - "https://qoder.com.cn/account/usage", - "Qoder China", - ), - ]; - - let mut auth_failed = false; - let mut last_error = None; - for (url, referer, login_method) in regions { - match self.fetch_region(url, referer, &cookie).await { - Ok(value) => { - return Ok(ProviderFetchResult::new( - snapshot_from_payload(&value, login_method)?, - "web", - )); - } - Err(ProviderError::AuthRequired) => auth_failed = true, - Err(err) => last_error = Some(err), - } + async fn fetch_usage_web( + &self, + ctx: &FetchContext, + ) -> Result { + if let Some(raw) = ctx.manual_cookie_header.as_deref() { + // An invalid or unroutable capture never becomes a request. + let credential = routing::manual_credential(raw).ok_or(ProviderError::NoCookies)?; + let candidate = Candidate { + site: credential.site, + cookie_header: credential.cookie_header, + source: format!("manual / {}", credential.site.domain()), + }; + // A manual credential is terminal: a rejection or failure is the + // answer, never a reason to try another site or a browser. + return self.fetch_candidate(&candidate).await; } + self.fetch_browser_candidates().await + } - if auth_failed { - Err(ProviderError::AuthRequired) - } else { - Err(last_error.unwrap_or_else(|| ProviderError::Parse("No Qoder usage payload".into()))) + /// Try every browser session for the international site, then the China + /// site. Each session is sent only to the origin it was read for. + async fn fetch_browser_candidates(&self) -> Result { + let mut tally = AttemptTally::default(); + for site in QoderSite::ALL { + let candidates = + match crate::providers::browser_cookie_headers_for_domain(site.domain()) { + Ok(headers) => browser_candidates(site, headers), + Err(ProviderError::NoCookies) => continue, + Err(error) => { + tally.record_import_error(error); + continue; + } + }; + for candidate in &candidates { + match self.fetch_candidate(candidate).await { + Ok(result) => return Ok(result), + Err(error) => tally.record(error), + } + } } + Err(tally.into_error()) } - async fn fetch_region( + async fn fetch_candidate( &self, - url: &str, - referer: &str, - cookie: &str, - ) -> Result { - let response = self - .client - .get(url) - .header("Cookie", cookie) - .header("Accept", "application/json, text/plain, */*") - .header("Referer", referer) - .header("User-Agent", USER_AGENT) - .send() - .await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { + candidate: &Candidate, + ) -> Result { + let request = build_request(&self.client, candidate)?; + let response = self.client.execute(request).await?; + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { return Err(ProviderError::AuthRequired); } - if !response.status().is_success() { + if !status.is_success() { return Err(ProviderError::Other(format!( - "Qoder usage returned status {}", - response.status() + "Qoder API returned HTTP {}.", + status.as_u16() ))); } - response + let value = response .json::() .await - .map_err(|e| ProviderError::Parse(format!("Failed to parse Qoder usage: {e}"))) + .map_err(|e| ProviderError::Parse(format!("Failed to parse Qoder usage: {e}")))?; + Ok(ProviderFetchResult::new( + snapshot_from_payload(&value, &candidate.source)?, + "web", + )) + } +} + +/// One credential bound to the single site it may be sent to. +#[derive(Debug, Clone, PartialEq, Eq)] +struct Candidate { + site: QoderSite, + cookie_header: String, + /// Shown as the login method: ` / ` or `manual / `. + source: String, +} + +fn browser_candidates(site: QoderSite, headers: Vec<(String, String)>) -> Vec { + headers + .into_iter() + .filter_map(|(browser, header)| { + Some(Candidate { + site, + cookie_header: normalize_cookie_header(&header)?, + source: format!("{browser} / {}", site.domain()), + }) + }) + .collect() +} + +fn build_request( + client: &Client, + candidate: &Candidate, +) -> Result { + let origin = candidate.site.origin(); + let mut cookie = + HeaderValue::from_str(&candidate.cookie_header).map_err(|_| ProviderError::NoCookies)?; + cookie.set_sensitive(true); + Ok(client + .get(candidate.site.usage_url()) + .header("Cookie", cookie) + .header("Accept", "application/json, text/plain, */*") + .header("Accept-Language", "en-US,en;q=0.9") + .header("User-Agent", USER_AGENT) + .header("Origin", origin) + .header("Referer", candidate.site.referer()) + .header("X-Requested-With", "XMLHttpRequest") + .header("Bx-V", BX_VERSION) + .build()?) +} + +/// Outcome bookkeeping across candidates. A rejection (401/403) moves on to +/// the next candidate; any other failure is remembered but does not stop the +/// search either, and outranks a rejection when nothing succeeds. +#[derive(Default)] +struct AttemptTally { + rejected: bool, + failure: Option, + import_error: Option, +} + +impl AttemptTally { + fn record(&mut self, error: ProviderError) { + match error { + ProviderError::AuthRequired => self.rejected = true, + other => self.failure = Some(other), + } + } + + fn record_import_error(&mut self, error: ProviderError) { + self.import_error.get_or_insert(error); + } + + fn into_error(self) -> ProviderError { + if let Some(failure) = self.failure { + failure + } else if self.rejected { + ProviderError::AuthRequired + } else { + self.import_error.unwrap_or(ProviderError::NoCookies) + } } } @@ -118,6 +206,9 @@ impl Default for QoderProvider { fn normalize_cookie_header(raw: &str) -> Option { let mut header = raw.trim(); + if header.chars().any(char::is_control) { + return None; + } if header .get(.."cookie:".len()) .is_some_and(|prefix| prefix.eq_ignore_ascii_case("cookie:")) @@ -424,18 +515,7 @@ impl Provider for QoderProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto | SourceMode::Web => { - let cookie = match ctx.manual_cookie_header.as_deref() { - Some(cookie) => cookie.to_string(), - None => crate::providers::browser_cookie_header(&[ - "qoder.com", - "www.qoder.com", - "qoder.com.cn", - "www.qoder.com.cn", - ])?, - }; - self.fetch_web(&cookie).await - } + SourceMode::Auto | SourceMode::Web => self.fetch_usage_web(ctx).await, SourceMode::OAuth | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) } @@ -521,4 +601,120 @@ mod tests { Some("a=1; b=2") ); } + + #[test] + fn normalize_cookie_header_rejects_control_characters() { + assert_eq!( + normalize_cookie_header( + "a=1 +Host: qoder.com.cn" + ), + None + ); + } + + #[test] + fn browser_candidates_bind_each_header_to_its_site_and_label() { + let candidates = browser_candidates( + QoderSite::China, + vec![ + ("Chrome".to_string(), "Cookie: session=one".to_string()), + ("Edge".to_string(), "empty=".to_string()), + ("Brave".to_string(), "session=two".to_string()), + ], + ); + assert_eq!( + candidates, + vec![ + Candidate { + site: QoderSite::China, + cookie_header: "session=one".to_string(), + source: "Chrome / qoder.com.cn".to_string(), + }, + Candidate { + site: QoderSite::China, + cookie_header: "session=two".to_string(), + source: "Brave / qoder.com.cn".to_string(), + }, + ] + ); + } + + #[test] + fn request_carries_upstream_headers_for_the_candidate_origin_only() { + for (site, origin) in [ + (QoderSite::International, "https://qoder.com"), + (QoderSite::China, "https://qoder.com.cn"), + ] { + let candidate = Candidate { + site, + cookie_header: "session=fixture".to_string(), + source: format!("manual / {}", site.domain()), + }; + let request = build_request(&Client::new(), &candidate).unwrap(); + let header = |name: &str| request.headers().get(name).unwrap().to_str().unwrap(); + assert_eq!( + request.url().as_str(), + format!("{origin}/api/v2/me/usages/big_model_credits") + ); + assert_eq!(header("Cookie"), "session=fixture"); + assert_eq!(header("Origin"), origin); + assert_eq!(header("Referer"), format!("{origin}/account/usage")); + assert_eq!(header("Bx-V"), "2.5.35"); + assert_eq!(header("X-Requested-With"), "XMLHttpRequest"); + assert_eq!(header("Accept-Language"), "en-US,en;q=0.9"); + assert!(request.headers().get("Cookie").unwrap().is_sensitive()); + } + } + + #[test] + fn tally_without_attempts_reports_missing_credential() { + assert!(matches!( + AttemptTally::default().into_error(), + ProviderError::NoCookies + )); + } + + #[test] + fn tally_reports_expired_auth_when_every_candidate_was_rejected() { + let mut tally = AttemptTally::default(); + tally.record(ProviderError::AuthRequired); + tally.record(ProviderError::AuthRequired); + assert!(matches!(tally.into_error(), ProviderError::AuthRequired)); + } + + #[test] + fn tally_keeps_last_non_auth_failure_over_rejections() { + for order in [[500, 0], [0, 503]] { + let mut tally = AttemptTally::default(); + for status in order { + tally.record(if status == 0 { + ProviderError::AuthRequired + } else { + ProviderError::Other(format!("Qoder API returned HTTP {status}.")) + }); + } + let ProviderError::Other(message) = tally.into_error() else { + panic!("expected the non-auth failure"); + }; + assert!(message.contains("HTTP 50") || message.contains("HTTP 503")); + } + + let mut tally = AttemptTally::default(); + tally.record(ProviderError::Other("first".into())); + tally.record(ProviderError::Other("last".into())); + assert!(matches!(tally.into_error(), ProviderError::Other(m) if m == "last")); + } + + #[test] + fn tally_reports_import_error_only_when_nothing_else_happened() { + let mut tally = AttemptTally::default(); + tally.record_import_error(ProviderError::Other("locked".into())); + assert!(matches!(tally.into_error(), ProviderError::Other(m) if m == "locked")); + + let mut tally = AttemptTally::default(); + tally.record_import_error(ProviderError::Other("locked".into())); + tally.record(ProviderError::AuthRequired); + assert!(matches!(tally.into_error(), ProviderError::AuthRequired)); + } } diff --git a/rust/src/providers/qoder/routing.rs b/rust/src/providers/qoder/routing.rs new file mode 100644 index 0000000000..49407ddcbe --- /dev/null +++ b/rust/src/providers/qoder/routing.rs @@ -0,0 +1,713 @@ +//! Qoder regional site routing. +//! +//! Qoder runs two independent sites, `qoder.com` and `qoder.com.cn`. A session +//! cookie is only valid for the site that issued it, so every credential is +//! bound to exactly one site and is never presented to the other one. +//! +//! Browser candidates are bound by the domain they were read for. A manual +//! credential is routed by its content: a plain `Cookie` header is +//! international unless a `Domain=` attribute names a site, a pasted cURL +//! command or raw HTTP request selects the site from its URL or `Host` header, +//! and anything ambiguous or conflicting is rejected before any request is +//! made. Ported from upstream v0.66.0 `QoderCookieRouting.swift`. + +use regex_lite::Regex; +use std::collections::BTreeSet; +use std::sync::OnceLock; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum QoderSite { + International, + China, +} + +impl QoderSite { + pub(super) const ALL: [Self; 2] = [Self::International, Self::China]; + + pub(super) const fn domain(self) -> &'static str { + match self { + Self::International => "qoder.com", + Self::China => "qoder.com.cn", + } + } + + pub(super) const fn origin(self) -> &'static str { + match self { + Self::International => "https://qoder.com", + Self::China => "https://qoder.com.cn", + } + } + + pub(super) fn usage_url(self) -> String { + format!("{}/api/v2/me/usages/big_model_credits", self.origin()) + } + + pub(super) fn referer(self) -> String { + format!("{}/account/usage", self.origin()) + } +} + +/// A manual credential resolved to one site. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct ManualCredential { + pub(super) site: QoderSite, + pub(super) cookie_header: String, +} + +/// Which pasted form selected the site; decides where the cookie comes from. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RouteKind { + Curl, + HttpRequest, + PlainCookie, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Route { + Site(QoderSite, RouteKind), + Invalid, +} + +/// Resolve a manual header to its site and normalized cookie header, or `None` +/// when the routing is invalid or no cookie can be extracted. +pub(super) fn manual_credential(raw: &str) -> Option { + let Route::Site(site, kind) = manual_route(raw) else { + return None; + }; + let cookie_header = match kind { + RouteKind::Curl => super::normalize_cookie_header(&curl_cookie_value(raw)?), + RouteKind::HttpRequest => super::normalize_cookie_header(&http_cookie_value(raw)?), + RouteKind::PlainCookie => plain_cookie_header(raw), + }?; + Some(ManualCredential { + site, + cookie_header, + }) +} + +#[cfg(test)] +pub(super) fn site_for_manual_header(raw: &str) -> Option { + match manual_route(raw) { + Route::Site(site, _) => Some(site), + Route::Invalid => None, + } +} + +fn manual_route(raw: &str) -> Route { + if let Some(route) = curl_request_route(raw) { + return route; + } + if let Some(route) = http_request_route(raw) { + return route; + } + plain_cookie_route(raw) +} + +// --------------------------------------------------------------------------- +// Plain Cookie header +// --------------------------------------------------------------------------- + +/// Swift `split(separator:maxSplits: 1)` with empty subsequences omitted: a +/// pair needs a non-empty name and a non-empty remainder. +fn split_pair(text: &str, separator: char) -> Option<(&str, &str)> { + let (first, rest) = text.trim_start_matches(separator).split_once(separator)?; + (!rest.is_empty()).then_some((first, rest)) +} + +fn plain_cookie_route(raw: &str) -> Route { + let mut routed: Option = None; + for part in raw.split(';').filter(|part| !part.is_empty()) { + let Some((name, value)) = split_pair(part, '=') else { + continue; + }; + if !name.trim().eq_ignore_ascii_case("domain") { + continue; + } + let Some(site) = site_for_host(value.trim().trim_matches(['"', '\''])) else { + return Route::Invalid; + }; + if routed.is_some_and(|previous| previous != site) { + return Route::Invalid; + } + routed = Some(site); + } + Route::Site( + routed.unwrap_or(QoderSite::International), + RouteKind::PlainCookie, + ) +} + +/// A `Domain=` attribute only routes the credential; it is not a cookie, so it +/// is not sent. +fn plain_cookie_header(raw: &str) -> Option { + let normalized = super::normalize_cookie_header(raw)?; + let pairs = normalized + .split("; ") + .filter(|pair| { + !pair + .split_once('=') + .is_some_and(|(name, _)| name.trim().eq_ignore_ascii_case("domain")) + }) + .collect::>(); + (!pairs.is_empty()).then(|| pairs.join("; ")) +} + +// --------------------------------------------------------------------------- +// Raw HTTP request +// --------------------------------------------------------------------------- + +const SUPPORTED_METHODS: [&str; 7] = ["get", "post", "put", "patch", "delete", "head", "options"]; +const UNSUPPORTED_KNOWN_METHODS: [&str; 2] = ["trace", "connect"]; + +fn is_newline(character: char) -> bool { + matches!( + character, + '\n' | '\r' | '\u{0B}' | '\u{0C}' | '\u{85}' | '\u{2028}' | '\u{2029}' + ) +} + +fn lines(raw: &str) -> impl Iterator { + raw.split(is_newline).filter(|line| !line.is_empty()) +} + +fn is_http_request_method_token(token: &str) -> bool { + !token.is_empty() && token.chars().all(|c| c.is_ascii_alphabetic()) +} + +fn http_request_route(raw: &str) -> Option { + let mut request_site: Option = None; + let mut saw_request_line = false; + + for line in lines(raw) { + let parts = line + .trim() + .split([' ', '\t']) + .filter(|part| !part.is_empty()) + .collect::>(); + if parts.len() < 2 { + continue; + } + let method = parts[0].to_ascii_lowercase(); + let is_http_versioned_line = + parts.len() >= 3 && parts[2].to_ascii_lowercase().starts_with("http/"); + if UNSUPPORTED_KNOWN_METHODS.contains(&method.as_str()) + || (is_http_versioned_line && is_http_request_method_token(parts[0])) + { + if !SUPPORTED_METHODS.contains(&method.as_str()) { + return Some(Route::Invalid); + } + } else if !SUPPORTED_METHODS.contains(&method.as_str()) { + continue; + } + + if saw_request_line { + return Some(Route::Invalid); + } + saw_request_line = true; + + let target = parts[1]; + if target.starts_with('/') { + request_site = None; + } else if let Some(site) = site_for_url_text(target) { + request_site = Some(site); + } else { + return Some(Route::Invalid); + } + } + + if !saw_request_line { + return None; + } + let host_sites = host_header_sites(raw); + if host_sites.iter().any(Option::is_none) { + return Some(Route::Invalid); + } + let concrete = host_sites.into_iter().flatten().collect::>(); + if concrete.iter().any(|site| Some(site) != concrete.first()) { + return Some(Route::Invalid); + } + + let site = match (request_site, concrete.first().copied()) { + (Some(request), Some(host)) if request != host => return Some(Route::Invalid), + (Some(request), _) => request, + (None, Some(host)) => host, + (None, None) => return Some(Route::Invalid), + }; + Some(Route::Site(site, RouteKind::HttpRequest)) +} + +fn host_header_sites(raw: &str) -> Vec> { + lines(raw) + .filter_map(|line| split_pair(line, ':')) + .filter(|(name, _)| name.trim().eq_ignore_ascii_case("host")) + .map(|(_, value)| site_for_host(value)) + .collect() +} + +fn http_cookie_value(raw: &str) -> Option { + lines(raw) + .filter_map(|line| split_pair(line, ':')) + .find(|(name, _)| name.trim().eq_ignore_ascii_case("cookie")) + .map(|(_, value)| value.trim().to_string()) +} + +// --------------------------------------------------------------------------- +// cURL command +// --------------------------------------------------------------------------- + +type Targets = Vec<(usize, QoderSite)>; + +fn curl_request_route(raw: &str) -> Option { + let Some(preprocessed) = preprocessed_curl_shell_text(raw) else { + return contains_curl_executable_text(raw).then_some(Route::Invalid); + }; + let tokens = shell_tokens(&preprocessed); + let Some(curl_index) = curl_command_index(&tokens) else { + return tokens + .iter() + .any(|token| is_curl_executable_token(token)) + .then_some(Route::Invalid); + }; + if !tokens.iter().all(|token| is_token_text_safe(token)) { + return Some(Route::Invalid); + } + + let (Some(explicit), Some(url_targets)) = ( + explicit_curl_url_targets(&tokens, curl_index), + url_token_targets(&tokens, curl_index), + ) else { + return Some(Route::Invalid); + }; + + let indices = explicit + .iter() + .chain(&url_targets) + .map(|(index, _)| *index) + .collect::>(); + if indices.len() != 1 { + return Some(Route::Invalid); + } + let target_index = *indices.first()?; + let trusted_index = curl_index + 1; + + let target_site = if let Some((_, site)) = explicit.iter().find(|(i, _)| *i == target_index) { + *site + } else if target_index == trusted_index + && let Some((_, site)) = url_targets.iter().find(|(i, _)| *i == trusted_index) + { + *site + } else { + return Some(Route::Invalid); + }; + + let Some(header_sites) = curl_header_values(&tokens, curl_index).and_then(|headers| { + headers + .iter() + .map(|header| inspect_curl_header_host(header)) + .filter_map(|inspection| match inspection { + HostInspection::Ignored => None, + HostInspection::Site(site) => Some(Some(site)), + HostInspection::Invalid => Some(None), + }) + .collect::>>() + }) else { + return Some(Route::Invalid); + }; + if header_sites.iter().any(|site| *site != target_site) { + return Some(Route::Invalid); + } + Some(Route::Site(target_site, RouteKind::Curl)) +} + +fn curl_cookie_value(raw: &str) -> Option { + let preprocessed = preprocessed_curl_shell_text(raw)?; + let tokens = shell_tokens(&preprocessed); + let curl_index = curl_command_index(&tokens)?; + curl_header_values(&tokens, curl_index)? + .into_iter() + .find_map(|header| { + let (name, value) = header.split_once(':')?; + name.trim() + .eq_ignore_ascii_case("cookie") + .then(|| value.trim().to_string()) + }) +} + +fn curl_command_index(tokens: &[String]) -> Option { + let index = tokens + .iter() + .position(|token| !is_shell_assignment(token))?; + is_curl_executable_token(&tokens[index]).then_some(index) +} + +fn is_curl_executable_token(token: &str) -> bool { + if token.contains('=') || token.contains("://") { + return false; + } + let executable = token + .split('/') + .rfind(|segment| !segment.is_empty()) + .unwrap_or(token); + executable.eq_ignore_ascii_case("curl") +} + +fn contains_curl_executable_text(text: &str) -> bool { + static PATTERN: OnceLock = OnceLock::new(); + if shell_tokens(text) + .iter() + .any(|token| is_curl_executable_token(token)) + { + return true; + } + PATTERN + .get_or_init(|| { + Regex::new(r"(?i)(^|[\s;])(?:[^\s;=]+/)?curl($|[\s;])").expect("curl text regex") + }) + .is_match(text) +} + +fn is_shell_assignment(token: &str) -> bool { + if token.contains(';') { + return false; + } + let Some((name, _)) = token.split_once('=') else { + return false; + }; + let mut characters = name.chars(); + characters + .next() + .is_some_and(|first| first == '_' || first.is_alphabetic()) + && name.chars().all(|c| c == '_' || c.is_alphanumeric()) +} + +fn is_token_text_safe(token: &str) -> bool { + token.chars().all(is_text_safe) +} + +fn is_text_safe(character: char) -> bool { + let value = u32::from(character); + value >= 0x20 && value != 0x7F +} + +fn explicit_curl_url_targets(tokens: &[String], curl_index: usize) -> Option { + let mut targets = Targets::new(); + let mut index = curl_index + 1; + while index < tokens.len() { + let token = &tokens[index]; + let lowercased = token.to_ascii_lowercase(); + if lowercased == "--url" { + let value_index = index + 1; + let site = site_for_url_text(tokens.get(value_index)?)?; + targets.push((value_index, site)); + index = value_index + 1; + continue; + } + if lowercased.starts_with("--url=") { + targets.push((index, site_for_url_text(&token["--url=".len()..])?)); + } + index += 1; + } + Some(targets) +} + +fn url_token_targets(tokens: &[String], curl_index: usize) -> Option { + let mut targets = Targets::new(); + for (index, token) in tokens.iter().enumerate().skip(curl_index + 1) { + if let Some(site) = site_for_url_text(token) { + targets.push((index, site)); + } else if host_for_url_text(token).is_some() { + return None; + } + } + Some(targets) +} + +enum ShortHeader { + Attached(String), + NextToken, + Invalid, +} + +fn short_options_contain(token: &str, option: char) -> bool { + token.starts_with('-') && !token.starts_with("--") && token[1..].contains(option) +} + +fn short_header_value(token: &str) -> Option { + if !token.starts_with('-') || token.starts_with("--") { + return None; + } + let options = &token[1..]; + let header_option = options.find('H')?; + if !options[..header_option] + .chars() + .all(|flag| "fsSL".contains(flag)) + { + return Some(ShortHeader::Invalid); + } + let attached = &options[header_option + 1..]; + Some(if attached.is_empty() { + ShortHeader::NextToken + } else { + ShortHeader::Attached(attached.to_string()) + }) +} + +/// Every `-H` / `--header` value in the command, or `None` when the command +/// can change its target or headers in ways this parser cannot see (config +/// files, variable expansion, header files, `--location-trusted`). +fn curl_header_values(tokens: &[String], curl_index: usize) -> Option> { + let mut values = Vec::new(); + let mut index = curl_index + 1; + while index < tokens.len() { + let token = &tokens[index]; + let lowercased = token.to_ascii_lowercase(); + if lowercased == "--config" + || lowercased.starts_with("--config=") + || lowercased.starts_with("--expand-") + || lowercased == "--location-trusted" + || short_options_contain(token, 'K') + { + return None; + } + if lowercased == "--header" { + values.push(tokens.get(index + 1)?.clone()); + index += 2; + } else if lowercased.starts_with("--header=") { + values.push(token["--header=".len()..].to_string()); + index += 1; + } else if let Some(short) = short_header_value(token) { + match short { + ShortHeader::Attached(value) => { + values.push(value); + index += 1; + } + ShortHeader::NextToken => { + values.push(tokens.get(index + 1)?.clone()); + index += 2; + } + ShortHeader::Invalid => return None, + } + } else { + index += 1; + } + } + Some(values) +} + +enum HostInspection { + Ignored, + Site(QoderSite), + Invalid, +} + +fn inspect_curl_header_host(header_value: &str) -> HostInspection { + let trimmed = header_value.trim(); + if trimmed.is_empty() || trimmed.starts_with('@') { + return HostInspection::Invalid; + } + let Some((name, value)) = trimmed.split_once(':') else { + let lowercased = trimmed.to_lowercase(); + let looks_like_host = lowercased == "host" + || lowercased.starts_with("host ") + || lowercased.starts_with("host\t") + || lowercased == "host;" + || lowercased.starts_with("host;"); + return if looks_like_host { + HostInspection::Invalid + } else { + HostInspection::Ignored + }; + }; + if !name.trim().eq_ignore_ascii_case("host") { + return HostInspection::Ignored; + } + site_for_host(value).map_or(HostInspection::Invalid, HostInspection::Site) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum ShellQuote { + Single, + Double, +} + +/// Normalize line continuations and reject anything a shell would evaluate +/// (control operators, substitutions, expansions, control characters, or an +/// unterminated quote). `None` means the text is not a plain curl command. +fn preprocessed_curl_shell_text(text: &str) -> Option { + let scalars = text.chars().collect::>(); + let mut output = String::with_capacity(text.len()); + let mut index = 0; + let mut quote: Option = None; + + while index < scalars.len() { + let scalar = scalars[index]; + let next = scalars.get(index + 1).copied(); + + if scalar == '\\' { + if quote.is_none() { + if next == Some('\n') { + index += 2; + continue; + } + if next == Some('\r') && scalars.get(index + 2) == Some(&'\n') { + index += 3; + continue; + } + } + if quote != Some(ShellQuote::Single) + && let Some(escaped) = next + && matches!(escaped, '\'' | '"' | '\\') + { + output.push(scalar); + output.push(escaped); + index += 2; + continue; + } + if quote == Some(ShellQuote::Double) { + return None; + } + } + + if !is_text_safe(scalar) { + return None; + } + + match quote { + Some(ShellQuote::Single) => { + if scalar == '\'' { + quote = None; + } + } + Some(ShellQuote::Double) => { + if scalar == '"' { + quote = None; + } else if scalar == '`' || is_unsupported_dollar_expansion(&scalars, index) { + return None; + } + } + None => { + if ";|&<>".contains(scalar) { + return None; + } else if scalar == '\'' { + quote = Some(ShellQuote::Single); + } else if scalar == '"' { + quote = Some(ShellQuote::Double); + } else if scalar == '`' || is_unsupported_dollar_expansion(&scalars, index) { + return None; + } + } + } + + output.push(scalar); + index += 1; + } + + quote.is_none().then_some(output) +} + +fn is_unsupported_dollar_expansion(scalars: &[char], index: usize) -> bool { + if scalars[index] != '$' { + return false; + } + let Some(next) = scalars.get(index + 1).copied() else { + return false; + }; + matches!(next, '\'' | '"' | '(' | '{' | '[') + || next == '_' + || next.is_alphabetic() + || next.is_ascii_digit() + || "*@#?$!-".contains(next) +} + +fn shell_tokens(text: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current = String::new(); + let mut quote: Option = None; + let mut escaped = false; + + for character in text.chars() { + if escaped { + current.push(character); + escaped = false; + continue; + } + if character == '\\' { + escaped = true; + continue; + } + if let Some(active) = quote { + if character == active { + quote = None; + } else { + current.push(character); + } + continue; + } + if character == '\'' || character == '"' { + quote = Some(character); + continue; + } + if character.is_whitespace() { + if !current.is_empty() { + tokens.push(std::mem::take(&mut current)); + } + continue; + } + current.push(character); + } + + if escaped { + current.push('\\'); + } + if !current.is_empty() { + tokens.push(current); + } + tokens +} + +// --------------------------------------------------------------------------- +// Hosts and URLs +// --------------------------------------------------------------------------- + +fn host_for_url_text(text: &str) -> Option { + let trimmed = text.trim().trim_matches(['"', '\'']); + let lowercased = trimmed.to_lowercase(); + if !(lowercased.starts_with("https://") || lowercased.starts_with("http://")) { + return None; + } + reqwest::Url::parse(trimmed) + .ok()? + .host_str() + .map(str::to_lowercase) +} + +fn site_for_url_text(text: &str) -> Option { + site_for_host(&host_for_url_text(text)?) +} + +/// Accepts `qoder.com`, `qoder.com.cn` and their `www.` forms, with an +/// optional leading dot, quotes, and a numeric port. +fn site_for_host(host: &str) -> Option { + let lowercased = host.trim().trim_matches(['"', '\'']).to_lowercase(); + let mut normalized = lowercased.strip_prefix('.').unwrap_or(&lowercased); + if let Some((hostname, port)) = normalized.rsplit_once(':') { + let valid_port = !hostname.contains(':') + && !port.is_empty() + && port.chars().all(|c| c.is_ascii_digit()) + && port + .parse::() + .is_ok_and(|port| (1..=65535).contains(&port)); + if !valid_port { + return None; + } + normalized = hostname; + } + match normalized { + "qoder.com" | "www.qoder.com" => Some(QoderSite::International), + "qoder.com.cn" | "www.qoder.com.cn" => Some(QoderSite::China), + _ => None, + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/providers/qoder/routing/tests.rs b/rust/src/providers/qoder/routing/tests.rs new file mode 100644 index 0000000000..55628655fc --- /dev/null +++ b/rust/src/providers/qoder/routing/tests.rs @@ -0,0 +1,451 @@ +use super::*; + +fn site(raw: &str) -> Option { + site_for_manual_header(raw) +} + +#[test] +fn manual_cookie_header_can_route_to_qoder_china_site() { + assert_eq!(site("sid=abc"), Some(QoderSite::International)); + assert_eq!( + site("sid=qoder.com.cn-looking-value"), + Some(QoderSite::International) + ); + assert_eq!( + site("sid=abc; note=curl https://qoder.com.cn"), + Some(QoderSite::International) + ); + assert_eq!( + site("sid=abc; redirect=https://example.com/curl"), + Some(QoderSite::International) + ); + assert_eq!( + site("sid=abc; Domain=.qoder.com.cn"), + Some(QoderSite::China) + ); + assert_eq!(site("sid=abc; Domain=qoder.com.cn"), Some(QoderSite::China)); + assert_eq!( + site("sid=abc; Domain=www.qoder.com.cn"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("HTTPS_PROXY=http://127.0.0.1:8080 curl https://qoder.com.cn"), + Some(QoderSite::China) + ); + assert_eq!( + site( + "HTTPS_PROXY=http://127.0.0.1:8080 \\\ncurl https://qoder.com.cn -H 'Cookie: sid=abc'" + ), + Some(QoderSite::China) + ); + assert_eq!( + site("\\\ncurl https://qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("\\\r\ncurl https://qoder.com -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); + assert_eq!( + site(concat!( + "curl https://qoder.com -H 'Origin: https://qoder.com' ", + "-H 'Referer: https://qoder.com/account/usage' -H 'Cookie: sid=abc'" + )), + Some(QoderSite::International) + ); + assert_eq!( + site(concat!( + "curl https://qoder.com.cn -H 'Origin: https://qoder.com.cn' ", + "-H 'Referer: https://qoder.com.cn/account/usage' -H 'Cookie: sid=abc'" + )), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://www.qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl --url https://qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl --url https://qoder.com --data 'x=1; Domain=qoder.com.cn'"), + Some(QoderSite::International) + ); + assert_eq!( + site("curl https://qoder.com --data 'GET /account/usage HTTP/1.1\nHost: qoder.com.cn'"), + None + ); + assert_eq!( + site("GET https://qoder.com.cn/account/usage"), + Some(QoderSite::China) + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn"), + Some(QoderSite::China) + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: www.qoder.com.cn"), + Some(QoderSite::China) + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn:443"), + Some(QoderSite::China) + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn:evil"), + None + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn:"), + None + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn:65536"), + None + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com.cn:443:444"), + None + ); + assert_eq!( + site("GET /account/usage HTTP/1.1\nHost: qoder.com"), + Some(QoderSite::International) + ); + assert_eq!( + site("TRACE /account/usage HTTP/1.1\nHost: qoder.com.cn"), + None + ); + assert_eq!(site("CONNECT qoder.com.cn:443 HTTP/1.1"), None); + assert_eq!( + site("BREW /account/usage HTTP/1.1\nHost: qoder.com.cn"), + None + ); + assert_eq!( + site( + "curl -H 'Referer: https://qoder.com.cn/account/usage' https://qoder.com -H 'Cookie: sid=abc'" + ), + None + ); + assert_eq!( + site( + "curl --proxy-header 'X: https://qoder.com.cn' https://qoder.com -H 'Cookie: sid=abc'" + ), + None + ); + assert_eq!(site("curl -X GET https://qoder.com.cn"), None); + assert_eq!(site("sudo curl https://qoder.com.cn"), None); + assert_eq!(site("sid=abc; curl https://qoder.com.cn"), None); + assert_eq!( + site("curl https://qoder.com/account https://qoder.com/profile"), + None + ); + assert_eq!( + site("curl --url https://qoder.com/account https://qoder.com/profile"), + None + ); + assert_eq!( + site("curl https://qoder.com https://qoder.com.cn -H 'Cookie: sid=abc'"), + None + ); + assert_eq!(site("curl https://example.com -H 'Cookie: sid=abc'"), None); + assert_eq!( + site("GET https://qoder.com/account/usage HTTP/1.1\nHost: qoder.com.cn"), + None + ); + assert_eq!( + site("GET https://qoder.com/account/usage HTTP/1.1\nHost: example.com"), + None + ); + assert_eq!(site("GET /account/usage HTTP/1.1\nHost: example.com"), None); +} + +#[test] +fn manual_curl_host_headers_must_match_authoritative_qoder_target() { + assert_eq!( + site("curl https://qoder.com -H 'Host: qoder.com' -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Host: www.qoder.com.cn:443' -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn -sH 'Host: qoder.com.cn' -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn -fsSLHHost:qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn -HHost:qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn --header=Host:qoder.com.cn -H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com.cn \\\n-H 'Host: qoder.com.cn' \\\r\n-H 'Cookie: sid=abc'"), + Some(QoderSite::China) + ); + assert_eq!( + site("curl https://qoder.com -H 'Host: qoder.com.cn' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Host: qoder.com' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Host: qoder.com.cn:evil' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site( + "curl https://qoder.com.cn -H 'Host: qoder.com.cn' -H 'Host: qoder.com' -H 'Cookie: sid=abc'" + ), + None + ); + assert_eq!( + site("curl https://qoder.com -sH 'Host: qoder.com.cn' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -fsSLHHost:qoder.com.cn -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -HHost:qoder.com.cn -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -XH 'Host: qoder.com.cn' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -H @headers.txt -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com --header @- -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Host:' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -H 'Host;' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn --header=Host\\; -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -sHHost\\; -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn -K qoder.curlrc -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn --config qoder.curlrc -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn --config=qoder.curlrc -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site(concat!( + "curl https://qoder.com --variable site=qoder.com.cn --expand-header 'Host: {{site}}' ", + "-H 'Cookie: sid=abc'" + )), + None + ); + assert_eq!( + site(concat!( + "curl https://qoder.com --variable site=qoder.com.cn --expand-url 'https://{{site}}' ", + "-H 'Cookie: sid=abc'" + )), + None + ); + assert_eq!( + site("curl https://qoder.com.cn --expand-config '{{config}}' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn ; echo -H 'Cookie: sid=global'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn | cat -H 'Cookie: sid=global'"), + None + ); + assert_eq!( + site("curl https://qoder.com.cn > headers.txt -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com && echo done -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl 'https://qoder.com/account/usage?a=1&b=2;next=ok' -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); + assert_eq!( + site("curl https://qoder.com -H 'X-Note: a;b|c&d=' -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); +} + +#[test] +fn manual_curl_rejects_shell_synthesis_and_injected_controls() { + assert_eq!( + site("curl https://qoder.com --location-trusted -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A $'agent\r\nHost: qoder.com.cn' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site(concat!( + "curl https://qoder.com --referer $'https://qoder.com\r\nHost: qoder.com.cn' ", + "-H 'Cookie: sid=abc'" + )), + None + ); + assert_eq!( + site( + "curl https://qoder.com -A \\'$'agent\\r\\nHost: qoder.com.cn'\\' -H 'Cookie: sid=abc'" + ), + None + ); + assert_eq!( + site("curl https://qoder.com -A \\'$'agent\r\nHost: qoder.com.cn'\\' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site( + "curl https://qoder.com -H 'User-Agent: agent\\\nHost: qoder.com.cn' -H 'Cookie: sid=abc'" + ), + None + ); + assert_eq!( + site("curl https://qoder.com -A $'agent\\r\\nHost: qoder.com.cn' -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A $(printf agent) -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A `printf agent` -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A $AGENT -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("\"curl\" https://qoder.com.cn -A $AGENT -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("'/usr/bin/curl' https://qoder.com.cn -A $AGENT -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("\\curl https://qoder.com.cn -A $AGENT -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("QODER_AGENT=$AGENT \\\ncurl https://qoder.com.cn -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -H \"User-Agent: $AGENT\" -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A $\"agent\" -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -H @<(printf 'Host: qoder.com.cn') -H 'Cookie: sid=abc'"), + None + ); + assert_eq!( + site("curl https://qoder.com -A \\'literal\\' -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); + assert_eq!( + site("curl https://qoder.com -A \\\"literal\\\" -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); + assert_eq!( + site("curl https://qoder.com -A literal\\\\slash -H 'Cookie: sid=abc'"), + Some(QoderSite::International) + ); +} + +#[test] +fn manual_credential_binds_cookie_to_its_routed_site() { + let china = manual_credential("curl https://qoder.com.cn -H 'Cookie: session=china'").unwrap(); + assert_eq!(china.site, QoderSite::China); + assert_eq!(china.cookie_header, "session=china"); + + let http = + manual_credential("GET /account/usage HTTP/1.1\nHost: qoder.com.cn\nCookie: session=http") + .unwrap(); + assert_eq!(http.site, QoderSite::China); + assert_eq!(http.cookie_header, "session=http"); + + let plain = manual_credential("Cookie: session=global").unwrap(); + assert_eq!(plain.site, QoderSite::International); + assert_eq!(plain.cookie_header, "session=global"); + + let lookalike = manual_credential("session=qoder.com.cn-looking-value").unwrap(); + assert_eq!(lookalike.site, QoderSite::International); +} + +#[test] +fn manual_domain_attribute_routes_but_is_not_sent_as_a_cookie() { + let credential = manual_credential("sid=abc; Domain=.qoder.com.cn").unwrap(); + assert_eq!(credential.site, QoderSite::China); + assert_eq!(credential.cookie_header, "sid=abc"); +} + +#[test] +fn manual_credential_rejects_invalid_or_conflicting_routing() { + for raw in [ + "curl https://qoder.com -H 'Host: qoder.com.cn' -H 'Cookie: sid=fixture'", + "sid=abc; Domain=qoder.com; Domain=qoder.com.cn", + "sid=abc; Domain=example.com", + "curl https://example.com -H 'Cookie: sid=abc'", + "curl https://qoder.com.cn", + "", + ] { + assert_eq!(manual_credential(raw), None, "{raw}"); + } +} + +#[test] +fn site_urls_stay_on_their_own_origin() { + for site in QoderSite::ALL { + assert!(site.usage_url().starts_with(site.origin())); + assert_eq!(site.referer(), format!("{}/account/usage", site.origin())); + } + assert_eq!(QoderSite::International.domain(), "qoder.com"); + assert_eq!(QoderSite::China.domain(), "qoder.com.cn"); +} From 44dada5dc4e77d0c2ff4117abc0dab480faab484 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:45:22 +0700 Subject: [PATCH 111/407] Port upstream 0.66.0: OpenAI Admin pagination, token totals, balance-fallback gate, transient retry --- rust/src/providers/openaiapi/mod.rs | 673 +++++++++++++--------- rust/src/providers/openaiapi/tests.rs | 795 ++++++++++++++++++++++++++ 2 files changed, 1183 insertions(+), 285 deletions(-) create mode 100644 rust/src/providers/openaiapi/tests.rs diff --git a/rust/src/providers/openaiapi/mod.rs b/rust/src/providers/openaiapi/mod.rs index 91e5abf697..bc92d1ebd8 100644 --- a/rust/src/providers/openaiapi/mod.rs +++ b/rust/src/providers/openaiapi/mod.rs @@ -1,13 +1,28 @@ //! OpenAI API usage provider. //! //! Tracks organization usage from the Admin API, with the older platform credit -//! balance endpoint as a fallback for project/user keys. +//! balance endpoint as a fallback for keys that are not project-scoped Admin keys. +//! +//! Upstream v0.66.0 `Resources/Plugins/openai.js` parity: +//! - `/v1/organization/costs` (`group_by=line_item`) and +//! `/v1/organization/usage/completions` (`group_by=model`) use `bucket_width=1d`, +//! `limit <= 31` and UTC-day-aligned 31-day ranges, following `has_more` / `next_page` +//! (at most 100 pages per range; a missing or repeated cursor is an error). +//! - Token totals are `input + input_audio + output + output_audio`; cached input is a +//! subset of input and is never added on top. +//! - Each Admin GET gets one transient retry (see [`RetryPolicy`]). +//! +//! The history window is fixed at 30 days. Upstream's `OPENAI_HISTORY_DAYS` (1-365) is +//! deferred until a Windows setting exists for it; [`usage_ranges`] already takes the day +//! count, so honoring it later only needs to pass the setting through. use async_trait::async_trait; -use chrono::{DateTime, Duration, TimeZone, Utc}; +use chrono::{DateTime, TimeZone, Utc}; use reqwest::Client; +use reqwest::header::{HeaderValue, RETRY_AFTER}; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::time::Duration; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, @@ -19,6 +34,16 @@ const OPENAI_ORG_COSTS_URL: &str = "https://api.openai.com/v1/organization/costs const OPENAI_ORG_COMPLETIONS_URL: &str = "https://api.openai.com/v1/organization/usage/completions"; const OPENAI_API_CREDENTIAL_TARGET: &str = "codexbar-openaiapi"; +/// Days of history requested from the Admin API (period label `Last 30 days`). +const HISTORY_DAYS: u32 = 30; +/// Endpoint bucket limit: each request covers at most this many daily buckets. +const MAX_BUCKETS_PER_REQUEST: u32 = 31; +const SECONDS_PER_DAY: i64 = 86_400; +/// Upstream `pages()` stops with a parse failure after this many pages per range. +const MAX_PAGES_PER_RANGE: usize = 100; +/// Upstream fetches with `timeoutSeconds: 20`. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(20); + #[derive(Debug, Deserialize)] struct CreditGrantsResponse { total_granted: f64, @@ -37,9 +62,12 @@ struct CreditGrant { expires_at: Option, } +/// One page of an Admin API listing. `has_more` is required; a page without it is malformed. #[derive(Debug, Deserialize)] -struct CostsResponse { - data: Vec, +struct Page { + data: Vec, + has_more: bool, + next_page: Option, } #[derive(Debug, Deserialize)] @@ -64,11 +92,6 @@ struct CostAmount { value: serde_json::Value, } -#[derive(Debug, Deserialize)] -struct CompletionsUsageResponse { - data: Vec, -} - #[derive(Debug, Deserialize)] struct CompletionsUsageBucket { start_time: i64, @@ -84,6 +107,10 @@ struct CompletionsUsageBucket { struct CompletionsUsageResult { model: Option, input_tokens: Option, + #[allow( + dead_code, + reason = "cached input is a subset of input_tokens, so it is tracked but never added to totals" + )] input_cached_tokens: Option, output_tokens: Option, input_audio_tokens: Option, @@ -91,9 +118,95 @@ struct CompletionsUsageResult { num_model_requests: Option, } +impl CompletionsUsageResult { + /// Upstream `tokens = input + input_audio + output + output_audio`. + fn total_tokens(&self) -> i64 { + [ + self.input_tokens, + self.input_audio_tokens, + self.output_tokens, + self.output_audio_tokens, + ] + .into_iter() + .map(|tokens| tokens.unwrap_or(0)) + .sum() + } +} + +/// One `start_time..end_time` request window of at most [`MAX_BUCKETS_PER_REQUEST`] days. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct UsageRange { + start: i64, + end: i64, + limit: u32, +} + +/// Transient retry for the Admin GETs, mirroring upstream `transientIdempotent` +/// (one retry, base delay 1 s, `Retry-After` honored up to 10 s). +#[derive(Debug, Clone, Copy)] +struct RetryPolicy { + base_delay: Duration, + max_delay: Duration, +} + +impl RetryPolicy { + const DEFAULT: Self = Self { + base_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(10), + }; + + /// Upstream `retryableStatusCodes`. + fn retries_status(status: reqwest::StatusCode) -> bool { + matches!(status.as_u16(), 408 | 429 | 500 | 502 | 503 | 504) + } + + /// Numeric `Retry-After` seconds capped at `max_delay`; anything else uses `base_delay`. + fn delay(&self, retry_after: Option<&HeaderValue>) -> Duration { + retry_after + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|seconds| seconds.is_finite() && *seconds >= 0.0) + .map_or(self.base_delay, |seconds| { + Duration::from_secs_f64(seconds.min(self.max_delay.as_secs_f64())) + }) + } +} + +#[derive(Debug, Clone)] +struct Endpoints { + credit_grants: String, + costs: String, + completions: String, +} + +impl Endpoints { + fn production() -> Self { + Self { + credit_grants: OPENAI_CREDIT_GRANTS_URL.to_string(), + costs: OPENAI_ORG_COSTS_URL.to_string(), + completions: OPENAI_ORG_COMPLETIONS_URL.to_string(), + } + } +} + +/// Resolved key plus whether it came from an Admin key source (Preferences, keychain, +/// `OPENAI_ADMIN_KEY`) rather than a plain `OPENAI_API_KEY`-style variable. +struct ApiCredential { + key: String, + is_admin: bool, +} + +/// Upstream `OpenAIAPIUsageCredential.allowsLegacyBalanceFallback`: the credit-grants +/// endpoint is not project-filtered, so a project-scoped Admin key never falls back to it. +fn allows_legacy_balance_fallback(project_id: Option<&str>, is_admin: bool) -> bool { + project_id.is_none() || !is_admin +} + pub struct OpenAIApiProvider { metadata: ProviderMetadata, client: Client, + endpoints: Endpoints, + retry: RetryPolicy, } impl OpenAIApiProvider { @@ -113,21 +226,23 @@ impl OpenAIApiProvider { tertiary_label_key: None, }, client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(15)) + .timeout(REQUEST_TIMEOUT) .build() .unwrap_or_else(|_| Client::new()), + endpoints: Endpoints::production(), + retry: RetryPolicy::DEFAULT, } } - fn api_key(api_key: Option<&str>) -> Result { + fn credential(api_key: Option<&str>) -> Result { resolve_api_key( api_key, OPENAI_API_CREDENTIAL_TARGET, &[ - "OPENAI_ADMIN_KEY", - "OPENAI_ADMIN_API_KEY", - "OPENAI_API_KEY", - "OPENAI_PLATFORM_API_KEY", + ("OPENAI_ADMIN_KEY", true), + ("OPENAI_ADMIN_API_KEY", true), + ("OPENAI_API_KEY", false), + ("OPENAI_PLATFORM_API_KEY", false), ], ) } @@ -135,7 +250,7 @@ impl OpenAIApiProvider { async fn fetch_api(&self, api_key: &str) -> Result { let response = self .client - .get(OPENAI_CREDIT_GRANTS_URL) + .get(&self.endpoints.credit_grants) .bearer_auth(api_key) .header("Accept", "application/json") .send() @@ -164,49 +279,117 @@ impl OpenAIApiProvider { &self, api_key: &str, project_id: Option<&str>, + now: DateTime, ) -> Result { - let now = Utc::now(); - let (start, end) = admin_time_range(now)?; let project_id = clean_project_id(project_id); - - let costs: CostsResponse = self - .fetch_admin_json( - OPENAI_ORG_COSTS_URL, - &admin_query( - ("start_time", start.to_string()), - ("end_time", end.to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "line_item".to_string()), - project_id.as_deref(), - ), + let ranges = usage_ranges(now, HISTORY_DAYS); + + let costs: Vec = self + .fetch_pages( + &self.endpoints.costs, + AdminQuery { + group_by: "line_item", + project_id: project_id.as_deref(), + label: "costs", + }, + &ranges, api_key, - "costs", ) .await?; - - let completions: CompletionsUsageResponse = self - .fetch_admin_json( - OPENAI_ORG_COMPLETIONS_URL, - &admin_query( - ("start_time", start.to_string()), - ("end_time", end.to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "model".to_string()), - project_id.as_deref(), - ), + let completions: Vec = self + .fetch_pages( + &self.endpoints.completions, + AdminQuery { + group_by: "model", + project_id: project_id.as_deref(), + label: "completions", + }, + &ranges, api_key, - "completions", ) .await?; - Ok(result_from_admin_usage( - &costs, - &completions, - now, - project_id.as_deref(), - )) + result_from_admin_usage(&costs, &completions, now, project_id.as_deref()) + } + + /// Admin usage first; the legacy balance endpoint only when + /// [`allows_legacy_balance_fallback`] permits it. + async fn fetch_admin_or_balance( + &self, + credential: &ApiCredential, + project_id: Option<&str>, + now: DateTime, + ) -> Result { + let admin_error = match self + .fetch_admin_usage(&credential.key, project_id, now) + .await + { + Ok(result) => return Ok(result), + Err(admin_error) => admin_error, + }; + if !allows_legacy_balance_fallback(project_id, credential.is_admin) { + return Err(admin_error); + } + match self.fetch_api(&credential.key).await { + Ok(result) => Ok(ProviderFetchResult { + source_label: "billing-api".to_string(), + ..result + }), + Err(balance_error) => { + if matches!(admin_error, ProviderError::AuthRequired) { + Err(balance_error) + } else { + Err(admin_error) + } + } + } + } + + /// Fetch every page of every range, following `has_more` / `next_page`. + async fn fetch_pages( + &self, + url: &str, + query: AdminQuery<'_>, + ranges: &[UsageRange], + api_key: &str, + ) -> Result, ProviderError> { + let label = query.label; + let mut buckets = Vec::new(); + for range in ranges { + let mut page: Option = None; + let mut seen = HashSet::new(); + for count in 0..MAX_PAGES_PER_RANGE { + let params = admin_query(range, &query, page.as_deref()); + let decoded: Page = self.fetch_admin_json(url, ¶ms, api_key, label).await?; + buckets.extend(decoded.data); + if !decoded.has_more { + break; + } + let cursor = decoded + .next_page + .as_deref() + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .ok_or_else(|| { + ProviderError::Parse(format!( + "OpenAI API {label} pagination cursor missing" + )) + })? + .to_string(); + if !seen.insert(cursor.clone()) { + return Err(ProviderError::Parse(format!( + "OpenAI API {label} pagination cursor repeated" + ))); + } + if count + 1 == MAX_PAGES_PER_RANGE { + return Err(ProviderError::Parse(format!( + "OpenAI API {label} pagination exceeded {MAX_PAGES_PER_RANGE} pages" + ))); + } + page = Some(cursor); + } + } + Ok(buckets) } async fn fetch_admin_json( @@ -216,14 +399,7 @@ impl OpenAIApiProvider { api_key: &str, label: &str, ) -> Result { - let response = self - .client - .get(url) - .query(query) - .bearer_auth(api_key) - .header("Accept", "application/json") - .send() - .await?; + let response = self.get_with_retry(url, query, api_key).await?; if response.status() == reqwest::StatusCode::UNAUTHORIZED || response.status() == reqwest::StatusCode::FORBIDDEN { @@ -244,6 +420,40 @@ impl OpenAIApiProvider { .await .map_err(|e| ProviderError::Parse(format!("Failed to parse OpenAI API {label}: {e}"))) } + + /// GET with at most one retry on a transient failure: 408/429/5xx, a timeout, or a + /// refused connection. Auth failures, TLS errors and other transport errors are + /// returned at once, and dropping the future (cancellation) never triggers a retry. + async fn get_with_retry( + &self, + url: &str, + query: &[(&str, String)], + api_key: &str, + ) -> Result { + let send = || { + self.client + .get(url) + .query(query) + .bearer_auth(api_key) + .header("Accept", "application/json") + .send() + }; + match send().await { + Ok(response) if RetryPolicy::retries_status(response.status()) => { + let delay = self.retry.delay(response.headers().get(RETRY_AFTER)); + tokio::time::sleep(delay).await; + } + Ok(response) => return Ok(response), + Err(error) => { + let error = ProviderError::Network(error); + if !error.is_transport_failure() { + return Err(error); + } + tokio::time::sleep(self.retry.delay(None)).await; + } + } + Ok(send().await?) + } } fn result_from_grants(grants: &CreditGrantsResponse) -> ProviderFetchResult { @@ -288,85 +498,52 @@ fn result_from_grants(grants: &CreditGrantsResponse) -> ProviderFetchResult { } fn result_from_admin_usage( - costs: &CostsResponse, - completions: &CompletionsUsageResponse, + costs: &[CostBucket], + completions: &[CompletionsUsageBucket], now: DateTime, project_id: Option<&str>, -) -> ProviderFetchResult { - let cost_total: f64 = costs - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| { - result - .amount - .as_ref() - .and_then(|a| number_value(&a.value)) - .unwrap_or(0.0) - }) - .sum(); - let request_total: i64 = completions - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| result.num_model_requests.unwrap_or(0)) - .sum(); - let token_total: i64 = completions - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| { - result.input_tokens.unwrap_or(0) - + result.input_cached_tokens.unwrap_or(0) - + result.output_tokens.unwrap_or(0) - + result.input_audio_tokens.unwrap_or(0) - + result.output_audio_tokens.unwrap_or(0) - }) - .sum(); - - let first_bucket = costs - .data - .first() - .map(|b| b.start_time) - .or_else(|| completions.data.first().map(|b| b.start_time)); - let start = first_bucket.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); +) -> Result { + let mut cost_total = 0.0; + let mut line_item_costs: HashMap = HashMap::new(); + for result in costs.iter().flat_map(|bucket| &bucket.results) { + let amount = cost_amount(result)?; + cost_total += amount; + let line_item = result + .line_item + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or("API"); + *line_item_costs.entry(line_item.to_string()).or_default() += amount; + } + let mut request_total: i64 = 0; + let mut token_total: i64 = 0; let mut model_tokens: HashMap = HashMap::new(); - for result in completions.data.iter().flat_map(|bucket| &bucket.results) { + for result in completions.iter().flat_map(|bucket| &bucket.results) { + let tokens = result.total_tokens(); + request_total += result.num_model_requests.unwrap_or(0); + token_total += tokens; let model = result .model .as_deref() .map(str::trim) .filter(|s| !s.is_empty()) .unwrap_or("Responses and Chat Completions"); - let tokens = result.input_tokens.unwrap_or(0) - + result.input_cached_tokens.unwrap_or(0) - + result.output_tokens.unwrap_or(0) - + result.input_audio_tokens.unwrap_or(0) - + result.output_audio_tokens.unwrap_or(0); *model_tokens.entry(model.to_string()).or_default() += tokens; } - let mut line_item_costs: HashMap = HashMap::new(); - for result in costs.data.iter().flat_map(|bucket| &bucket.results) { - let line_item = result - .line_item - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - .unwrap_or("API"); - *line_item_costs.entry(line_item.to_string()).or_default() += result - .amount - .as_ref() - .and_then(|a| number_value(&a.value)) - .unwrap_or(0.0); - } + let first_bucket = costs + .first() + .map(|b| b.start_time) + .or_else(|| completions.first().map(|b| b.start_time)); + let start = first_bucket.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); let mut usage = UsageSnapshot::new(RateWindow::with_details( 0.0, None, start, - Some(format!("${cost_total:.2} over last 30 days")), + Some(format!("${cost_total:.2} over last {HISTORY_DAYS} days")), )) .with_extra_rate_window( "requests", @@ -409,36 +586,82 @@ fn result_from_admin_usage( ); } - ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( - cost_total, - "USD", - "Last 30 days", - )) + Ok( + ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( + cost_total, + "USD", + format!("Last {HISTORY_DAYS} days"), + )), + ) +} + +/// A missing, null or blank amount counts as zero; a present but non-numeric or +/// non-finite one (`NaN`, `Infinity`, `1e309`) is a parse failure, as upstream. +fn cost_amount(result: &CostResult) -> Result { + let Some(amount) = &result.amount else { + return Ok(0.0); + }; + match &amount.value { + serde_json::Value::Null => Ok(0.0), + serde_json::Value::String(text) if text.trim().is_empty() => Ok(0.0), + value => number_value(value).ok_or_else(|| { + ProviderError::Parse("OpenAI API costs amount must be numeric".to_string()) + }), + } +} + +/// Query parameters shared by every page of one Admin API listing. +struct AdminQuery<'a> { + group_by: &'static str, + project_id: Option<&'a str>, + label: &'static str, } fn admin_query( - start: (&'static str, String), - end: (&'static str, String), - bucket_width: (&'static str, String), - limit: (&'static str, String), - group_by: (&'static str, String), - project_id: Option<&str>, + range: &UsageRange, + query: &AdminQuery<'_>, + page: Option<&str>, ) -> Vec<(&'static str, String)> { - let mut query = vec![start, end, bucket_width, limit, group_by]; - if let Some(project_id) = clean_project_id(project_id) { - query.push(("project_ids", project_id)); + let mut params = vec![ + ("start_time", range.start.to_string()), + ("end_time", range.end.to_string()), + ("bucket_width", "1d".to_string()), + ("limit", range.limit.to_string()), + ("group_by", query.group_by.to_string()), + ]; + if let Some(project_id) = clean_project_id(query.project_id) { + params.push(("project_ids", project_id)); } - query + if let Some(page) = page { + params.push(("page", page.to_string())); + } + params } -fn admin_time_range(now: DateTime) -> Result<(i64, i64), ProviderError> { - let start = (now.date_naive() - Duration::days(29)) - .and_hms_opt(0, 0, 0) - .ok_or_else(|| ProviderError::Parse("Invalid OpenAI usage start date".to_string()))? - .and_utc() - .timestamp(); - Ok((start, now.timestamp())) +/// Upstream `ranges()`: UTC-day-aligned windows of at most [`MAX_BUCKETS_PER_REQUEST`] days +/// covering the last `history_days` days including today. Upstream ends the final window +/// at tomorrow's midnight; the API rejects a future `end_time` +/// (`end_time must not be in the future`), so the end is clamped to `now`. +fn usage_ranges(now: DateTime, history_days: u32) -> Vec { + let now = now.timestamp(); + let today = now - now.rem_euclid(SECONDS_PER_DAY); + let mut start = today - (i64::from(history_days) - 1) * SECONDS_PER_DAY; + let mut remaining = history_days; + let mut ranges = Vec::new(); + while remaining > 0 { + let limit = remaining.min(MAX_BUCKETS_PER_REQUEST); + let end = start + i64::from(limit) * SECONDS_PER_DAY; + ranges.push(UsageRange { + start, + end: end.min(now), + limit, + }); + start = end; + remaining -= limit; + } + ranges } + fn clean_project_id(project_id: Option<&str>) -> Option { project_id .map(str::trim) @@ -493,26 +716,10 @@ impl Provider for OpenAIApiProvider { 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())?; - match self - .fetch_admin_usage(&api_key, ctx.workspace_id.as_deref()) + let credential = Self::credential(ctx.api_key.as_deref())?; + let project_id = clean_project_id(ctx.workspace_id.as_deref()); + self.fetch_admin_or_balance(&credential, project_id.as_deref(), Utc::now()) .await - { - Ok(result) => Ok(result), - Err(admin_error) => match self.fetch_api(&api_key).await { - Ok(result) => Ok(ProviderFetchResult { - source_label: "billing-api".to_string(), - ..result - }), - Err(balance_error) => { - if matches!(admin_error, ProviderError::AuthRequired) { - Err(balance_error) - } else { - Err(admin_error) - } - } - }, - } } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -525,32 +732,44 @@ impl Provider for OpenAIApiProvider { } } +/// Resolve the key from Preferences, the credential store, then the environment. +/// Each `env_names` entry carries whether that variable holds an Admin key. fn resolve_api_key( explicit: Option<&str>, credential_target: &str, - env_names: &[&str], -) -> Result { + env_names: &[(&str, bool)], +) -> Result { if let Some(key) = explicit && !key.trim().is_empty() { - return Ok(key.trim().to_string()); + return Ok(ApiCredential { + key: key.trim().to_string(), + is_admin: true, + }); } if let Ok(entry) = keyring::Entry::new(credential_target, "api_key") && let Ok(key) = entry.get_password() && !key.trim().is_empty() { - return Ok(key); + return Ok(ApiCredential { + key, + is_admin: true, + }); } - for env in env_names { + for (env, is_admin) in env_names { if let Ok(key) = std::env::var(env) && !key.trim().is_empty() { - return Ok(key); + return Ok(ApiCredential { + key, + is_admin: *is_admin, + }); } } + let names: Vec<&str> = env_names.iter().map(|(name, _)| *name).collect(); Err(ProviderError::NotInstalled(format!( "API key not found. Set {} in Preferences or environment.", - env_names.join(" / ") + names.join(" / ") ))) } @@ -561,120 +780,4 @@ fn resolve_api_key( fn _assert_datetime_send(_: DateTime) {} #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn openai_api_credit_snapshot_formats_available_balance() { - let result = result_from_grants(&CreditGrantsResponse { - total_granted: 100.0, - total_used: 25.0, - total_available: 75.0, - grants: None, - }); - assert_eq!(result.usage.primary.used_percent, 25.0); - assert_eq!(result.cost.unwrap().remaining(), Some(75.0)); - } - - #[test] - fn openai_admin_usage_accepts_numeric_string_cost_amounts() { - let costs = CostsResponse { - data: vec![CostBucket { - start_time: 1_797_638_400, - end_time: 1_797_724_800, - results: vec![CostResult { - amount: Some(CostAmount { - value: serde_json::json!("12.50"), - }), - line_item: Some("API".to_string()), - }], - }], - }; - let completions = CompletionsUsageResponse { - data: vec![CompletionsUsageBucket { - start_time: 1_797_638_400, - end_time: 1_797_724_800, - results: vec![CompletionsUsageResult { - model: Some("gpt-5.1".to_string()), - input_tokens: Some(100), - input_cached_tokens: Some(25), - output_tokens: Some(50), - input_audio_tokens: None, - output_audio_tokens: None, - num_model_requests: Some(7), - }], - }], - }; - let now = Utc.timestamp_opt(1_797_724_800, 0).single().unwrap(); - let result = result_from_admin_usage(&costs, &completions, now, Some("proj_demo")); - assert_eq!(result.cost.unwrap().used, 12.5); - assert_eq!( - result.usage.account_organization.as_deref(), - Some("Project: proj_demo") - ); - assert_eq!( - result.usage.login_method.as_deref(), - Some("Admin API: proj_demo") - ); - assert!(result.usage.secondary.is_none()); - let requests = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == "requests") - .unwrap(); - assert!(requests.window.is_informational); - assert_eq!( - requests.window.reset_description.as_deref(), - Some("7 requests") - ); - assert!( - result.usage.extra_rate_windows.iter().any(|window| window - .window - .reset_description - .as_deref() - == Some("175 tokens")) - ); - } - - #[test] - fn openai_admin_usage_rejects_nonfinite_cost_amounts() { - assert_eq!(number_value(&serde_json::json!("NaN")), None); - assert_eq!(number_value(&serde_json::json!("Infinity")), None); - } - - #[test] - fn openai_admin_query_scopes_project_ids_when_configured() { - let query = admin_query( - ("start_time", "1".to_string()), - ("end_time", "2".to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "model".to_string()), - Some(" proj_123 "), - ); - assert!(query.contains(&("project_ids", "proj_123".to_string()))); - } - - #[test] - fn openai_admin_time_range_does_not_end_in_the_future() { - let now = Utc.timestamp_opt(1_783_981_862, 0).single().unwrap(); - let (start, end) = admin_time_range(now).unwrap(); - - assert_eq!(end, now.timestamp()); - assert!(start < end); - } - - #[test] - fn openai_response_error_detail_prefers_api_message() { - assert_eq!( - response_error_detail(r#"{"error":{"message":"end_time must not be in the future"}}"#), - "end_time must not be in the future" - ); - assert_eq!( - response_error_detail("plain text failure"), - "plain text failure" - ); - assert_eq!(response_error_detail(" "), ""); - } -} +mod tests; diff --git a/rust/src/providers/openaiapi/tests.rs b/rust/src/providers/openaiapi/tests.rs new file mode 100644 index 0000000000..37493fef4b --- /dev/null +++ b/rust/src/providers/openaiapi/tests.rs @@ -0,0 +1,795 @@ +use super::*; +use mockito::{Matcher, Server, ServerGuard}; + +const COSTS_PATH: &str = "/v1/organization/costs"; +const COMPLETIONS_PATH: &str = "/v1/organization/usage/completions"; +const GRANTS_PATH: &str = "/v1/dashboard/billing/credit_grants"; +const EMPTY_PAGE: &str = r#"{"object":"page","data":[],"has_more":false,"next_page":null}"#; +/// 2023-11-17T00:00:00Z, the fixture clock used by upstream `OpenAIAPIUsageFetcherTests`. +const NOW: i64 = 1_700_179_200; + +fn fixed_now(offset_seconds: i64) -> DateTime { + Utc.timestamp_opt(NOW + offset_seconds, 0).single().unwrap() +} + +/// A provider pointed at the mock server, with retry delays removed so tests stay fast. +fn provider(server: &ServerGuard) -> OpenAIApiProvider { + let mut provider = OpenAIApiProvider::new(); + provider.client = Client::builder() + .no_proxy() + .build() + .expect("the test client should build"); + provider.endpoints = Endpoints { + credit_grants: format!("{}{GRANTS_PATH}", server.url()), + costs: format!("{}{COSTS_PATH}", server.url()), + completions: format!("{}{COMPLETIONS_PATH}", server.url()), + }; + provider.retry = RetryPolicy { + base_delay: Duration::ZERO, + max_delay: Duration::ZERO, + }; + provider +} + +fn admin_credential() -> ApiCredential { + ApiCredential { + key: "sk-test".to_string(), + is_admin: true, + } +} + +async fn mock_page( + server: &mut ServerGuard, + path: &str, + query: Matcher, + body: &str, +) -> mockito::Mock { + server + .mock("GET", path) + .match_query(query) + .with_status(200) + .with_body(body) + .create_async() + .await +} + +fn extra_window<'a>(result: &'a ProviderFetchResult, id: &str) -> &'a str { + result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == id) + .and_then(|window| window.window.reset_description.as_deref()) + .unwrap_or_else(|| panic!("missing extra window {id}")) +} + +fn cost_bucket(amount: serde_json::Value) -> CostBucket { + CostBucket { + start_time: 1_797_638_400, + end_time: 1_797_724_800, + results: vec![CostResult { + amount: Some(CostAmount { value: amount }), + line_item: Some("API".to_string()), + }], + } +} + +#[test] +fn openai_api_credit_snapshot_formats_available_balance() { + let result = result_from_grants(&CreditGrantsResponse { + total_granted: 100.0, + total_used: 25.0, + total_available: 75.0, + grants: None, + }); + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!(result.cost.unwrap().remaining(), Some(75.0)); +} + +#[test] +fn openai_admin_usage_accepts_numeric_string_cost_amounts() { + let costs = [cost_bucket(serde_json::json!("12.50"))]; + let completions = [CompletionsUsageBucket { + start_time: 1_797_638_400, + end_time: 1_797_724_800, + results: vec![CompletionsUsageResult { + model: Some("gpt-5.1".to_string()), + input_tokens: Some(100), + input_cached_tokens: Some(25), + output_tokens: Some(50), + input_audio_tokens: None, + output_audio_tokens: None, + num_model_requests: Some(7), + }], + }]; + let now = Utc.timestamp_opt(1_797_724_800, 0).single().unwrap(); + let result = result_from_admin_usage(&costs, &completions, now, Some("proj_demo")).unwrap(); + assert_eq!(result.cost.as_ref().unwrap().used, 12.5); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Project: proj_demo") + ); + assert_eq!( + result.usage.login_method.as_deref(), + Some("Admin API: proj_demo") + ); + assert!(result.usage.secondary.is_none()); + let requests = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == "requests") + .unwrap(); + assert!(requests.window.is_informational); + assert_eq!( + requests.window.reset_description.as_deref(), + Some("7 requests") + ); + // Cached input is a subset of input: 100 + 50, not 100 + 25 + 50. + assert_eq!(extra_window(&result, "tokens"), "150 tokens"); + assert_eq!(extra_window(&result, "model-0"), "150 tokens"); +} + +#[test] +fn openai_admin_tokens_sum_text_and_audio_without_cached() { + // Upstream `tokens = input + input_audio + output + output_audio`. + let completions = [CompletionsUsageBucket { + start_time: 1_700_000_000, + end_time: 1_700_086_400, + results: vec![ + CompletionsUsageResult { + model: Some("gpt-5.2".to_string()), + input_tokens: Some(1000), + input_cached_tokens: Some(250), + output_tokens: Some(500), + input_audio_tokens: Some(40), + output_audio_tokens: Some(10), + num_model_requests: Some(7), + }, + CompletionsUsageResult { + model: Some("gpt-5.2-codex".to_string()), + input_tokens: Some(300), + input_cached_tokens: None, + output_tokens: Some(200), + input_audio_tokens: None, + output_audio_tokens: None, + num_model_requests: Some(3), + }, + ], + }]; + let result = result_from_admin_usage(&[], &completions, fixed_now(0), None).unwrap(); + assert_eq!(extra_window(&result, "tokens"), "2050 tokens"); + assert_eq!(extra_window(&result, "requests"), "10 requests"); + assert_eq!(extra_window(&result, "model-0"), "1550 tokens"); + assert_eq!(extra_window(&result, "model-1"), "500 tokens"); +} + +#[test] +fn openai_admin_usage_sums_costs_and_ranks_line_items() { + // Cases from upstream `parses admin costs and completions usage into daily summaries`. + let costs = [ + CostBucket { + start_time: 1_700_000_000, + end_time: 1_700_086_400, + results: vec![ + CostResult { + amount: Some(CostAmount { + value: serde_json::json!(12.50), + }), + line_item: Some("Text tokens".to_string()), + }, + CostResult { + amount: Some(CostAmount { + value: serde_json::json!("2.25"), + }), + line_item: Some("Web search tool calls".to_string()), + }, + ], + }, + CostBucket { + start_time: 1_700_086_400, + end_time: 1_700_172_800, + results: vec![CostResult { + amount: Some(CostAmount { + value: serde_json::json!(4.00), + }), + line_item: Some("Text tokens".to_string()), + }], + }, + ]; + let result = result_from_admin_usage(&costs, &[], fixed_now(0), None).unwrap(); + assert_eq!(result.cost.as_ref().unwrap().used, 18.75); + assert_eq!(result.cost.as_ref().unwrap().period, "Last 30 days"); + assert_eq!(extra_window(&result, "line-item-0"), "$16.50"); + assert_eq!(extra_window(&result, "line-item-1"), "$2.25"); +} + +#[test] +fn openai_admin_usage_rejects_nonfinite_cost_amounts() { + assert_eq!(number_value(&serde_json::json!("NaN")), None); + assert_eq!(number_value(&serde_json::json!("Infinity")), None); + for value in ["NaN", "Infinity", "-Infinity", "1e309", "-1e309"] { + let costs = [cost_bucket(serde_json::json!(value))]; + let error = result_from_admin_usage(&costs, &[], fixed_now(0), None) + .err() + .unwrap_or_else(|| panic!("{value} should be rejected")); + assert!(matches!(error, ProviderError::Parse(_)), "{value}: {error}"); + } +} + +#[test] +fn openai_admin_usage_treats_blank_amounts_as_zero() { + let costs = [ + cost_bucket(serde_json::Value::Null), + cost_bucket(serde_json::json!(" ")), + CostBucket { + start_time: 1, + end_time: 2, + results: vec![CostResult { + amount: None, + line_item: None, + }], + }, + ]; + let result = result_from_admin_usage(&costs, &[], fixed_now(0), None).unwrap(); + assert_eq!(result.cost.unwrap().used, 0.0); +} + +#[test] +fn openai_admin_query_scopes_project_ids_and_page_when_configured() { + let range = UsageRange { + start: 1, + end: 2, + limit: 31, + }; + let query = AdminQuery { + group_by: "model", + project_id: Some(" proj_123 "), + label: "completions", + }; + let params = admin_query(&range, &query, Some("cursor_2")); + assert!(params.contains(&("project_ids", "proj_123".to_string()))); + assert!(params.contains(&("page", "cursor_2".to_string()))); + assert!(params.contains(&("bucket_width", "1d".to_string()))); + assert!(params.contains(&("group_by", "model".to_string()))); + + let unscoped = admin_query( + &range, + &AdminQuery { + project_id: None, + ..query + }, + None, + ); + assert!(unscoped.iter().all(|(name, _)| *name != "project_ids")); + assert!(unscoped.iter().all(|(name, _)| *name != "page")); +} + +#[test] +fn openai_admin_time_range_does_not_end_in_the_future() { + let now = Utc.timestamp_opt(1_783_981_862, 0).single().unwrap(); + let ranges = usage_ranges(now, HISTORY_DAYS); + + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].end, now.timestamp()); + assert_eq!( + ranges[0].start, + 1_783_981_862 - 1_783_981_862 % 86_400 - 29 * 86_400 + ); + assert_eq!(ranges[0].limit, 30); + assert!(ranges[0].start < ranges[0].end); +} + +#[test] +fn openai_usage_ranges_page_long_history_within_endpoint_bucket_limit() { + // Upstream `admin usage fetch pages long history within endpoint bucket limit`: 90 days + // is three ranges of 31, 31 and 28 buckets, contiguous and aligned to UTC day starts. + let ranges = usage_ranges(fixed_now(0), 90); + let limits: Vec = ranges.iter().map(|range| range.limit).collect(); + assert_eq!(limits, [31, 31, 28]); + assert!(limits.iter().all(|limit| *limit <= 31)); + assert_eq!(ranges[0].start, NOW - 89 * 86_400); + assert!(ranges.iter().all(|range| range.start % 86_400 == 0)); + assert_eq!(ranges[1].start, ranges[0].start + 31 * 86_400); + assert_eq!(ranges[2].start, ranges[1].start + 31 * 86_400); + assert_eq!(ranges[0].end, ranges[1].start); + assert_eq!(ranges[2].end, NOW); + + assert_eq!(usage_ranges(fixed_now(3_600), 1).len(), 1); + assert_eq!(usage_ranges(fixed_now(3_600), 1)[0].start, NOW); + assert_eq!(usage_ranges(fixed_now(3_600), 31).len(), 1); + assert_eq!(usage_ranges(fixed_now(3_600), 32).len(), 2); +} + +#[test] +fn openai_response_error_detail_prefers_api_message() { + assert_eq!( + response_error_detail(r#"{"error":{"message":"end_time must not be in the future"}}"#), + "end_time must not be in the future" + ); + assert_eq!( + response_error_detail("plain text failure"), + "plain text failure" + ); + assert_eq!(response_error_detail(" "), ""); +} + +#[test] +fn openai_retry_policy_honors_retry_after_up_to_ten_seconds() { + let policy = RetryPolicy::DEFAULT; + let header = |value: &'static str| HeaderValue::from_static(value); + + assert_eq!(policy.delay(None), Duration::from_secs(1)); + assert_eq!(policy.delay(Some(&header("3"))), Duration::from_secs(3)); + assert_eq!(policy.delay(Some(&header(" 2 "))), Duration::from_secs(2)); + assert_eq!(policy.delay(Some(&header("0"))), Duration::ZERO); + assert_eq!(policy.delay(Some(&header("600"))), Duration::from_secs(10)); + for invalid in ["soon", "-5", "inf", "NaN", ""] { + assert_eq!( + policy.delay(Some(&header(invalid))), + Duration::from_secs(1), + "{invalid:?}" + ); + } +} + +#[test] +fn openai_retry_policy_retries_only_transient_statuses() { + for status in [408, 429, 500, 502, 503, 504] { + assert!(RetryPolicy::retries_status( + reqwest::StatusCode::from_u16(status).unwrap() + )); + } + for status in [200, 400, 401, 403, 404, 422] { + assert!(!RetryPolicy::retries_status( + reqwest::StatusCode::from_u16(status).unwrap() + )); + } +} + +#[test] +fn openai_balance_fallback_is_blocked_only_for_project_scoped_admin_keys() { + assert!(allows_legacy_balance_fallback(None, true)); + assert!(allows_legacy_balance_fallback(None, false)); + assert!(allows_legacy_balance_fallback(Some("proj_abc"), false)); + assert!(!allows_legacy_balance_fallback(Some("proj_abc"), true)); +} + +#[test] +fn openai_credential_kind_follows_the_key_source() { + let credential = resolve_api_key( + Some(" sk-admin-configured "), + "codexbar-openaiapi-test-unused", + &[("CODEXBAR_TEST_OPENAI_UNSET", false)], + ) + .unwrap(); + assert_eq!(credential.key, "sk-admin-configured"); + assert!(credential.is_admin); +} + +#[tokio::test] +async fn openai_admin_usage_filters_costs_and_completions_by_project() { + let mut server = Server::new_async().await; + let query = |group_by: &str| { + Matcher::AllOf(vec![ + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + Matcher::UrlEncoded("group_by".into(), group_by.into()), + Matcher::UrlEncoded("bucket_width".into(), "1d".into()), + Matcher::UrlEncoded("limit".into(), "30".into()), + ]) + }; + let costs = mock_page(&mut server, COSTS_PATH, query("line_item"), EMPTY_PAGE) + .await + .expect(1); + let completions = mock_page(&mut server, COMPLETIONS_PATH, query("model"), EMPTY_PAGE) + .await + .expect(1); + + let result = provider(&server) + .fetch_admin_usage("sk-test", Some(" proj_abc "), fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + completions.assert_async().await; + assert_eq!(result.source_label, "admin-api"); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Project: proj_abc") + ); +} + +#[tokio::test] +async fn openai_admin_usage_pages_each_range_of_a_long_history() { + let mut server = Server::new_async().await; + let costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(3); + let completions = mock_page(&mut server, COMPLETIONS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(3); + + let provider = provider(&server); + let ranges = usage_ranges(fixed_now(0), 90); + let buckets: Vec = provider + .fetch_pages( + &provider.endpoints.costs, + AdminQuery { + group_by: "line_item", + project_id: None, + label: "costs", + }, + &ranges, + "sk-test", + ) + .await + .unwrap(); + assert!(buckets.is_empty()); + let _: Vec = provider + .fetch_pages( + &provider.endpoints.completions, + AdminQuery { + group_by: "model", + project_id: None, + label: "completions", + }, + &ranges, + "sk-test", + ) + .await + .unwrap(); + + costs.assert_async().await; + completions.assert_async().await; +} + +const COSTS_PAGE_1: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.costs.result","amount":{"value":1.25,"currency":"usd"},"line_item":"Text tokens"}]}],"has_more":true,"next_page":"costs_page_2"}"#; +const COSTS_PAGE_2: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.costs.result","amount":{"value":2.75,"currency":"usd"},"line_item":"Web search tool calls"}]}],"has_more":false,"next_page":null}"#; +const COMPLETIONS_PAGE_1: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.usage.completions.result","input_tokens":10,"output_tokens":5,"num_model_requests":1,"model":"gpt-5.2"}]}],"has_more":true,"next_page":"completions_page_2"}"#; +const COMPLETIONS_PAGE_2: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.usage.completions.result","input_tokens":20,"output_tokens":10,"num_model_requests":2,"model":"gpt-5.2"}]}],"has_more":false,"next_page":null}"#; + +#[tokio::test] +async fn openai_admin_usage_follows_costs_and_completions_pagination_cursors() { + let mut server = Server::new_async().await; + let page = |name: &str| Matcher::UrlEncoded("page".into(), name.into()); + let mocks = vec![ + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::AllOf(vec![Matcher::Regex( + "project_ids=proj_abc$".into(), + )])) + .with_body(COSTS_PAGE_1) + .expect(1) + .create_async() + .await, + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::AllOf(vec![ + page("costs_page_2"), + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + ])) + .with_body(COSTS_PAGE_2) + .expect(1) + .create_async() + .await, + server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::AllOf(vec![Matcher::Regex( + "project_ids=proj_abc$".into(), + )])) + .with_body(COMPLETIONS_PAGE_1) + .expect(1) + .create_async() + .await, + server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::AllOf(vec![ + page("completions_page_2"), + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + ])) + .with_body(COMPLETIONS_PAGE_2) + .expect(1) + .create_async() + .await, + ]; + + let result = provider(&server) + .fetch_admin_usage("sk-test", Some("proj_abc"), fixed_now(3_600)) + .await + .unwrap(); + + for mock in &mocks { + mock.assert_async().await; + } + assert_eq!(result.cost.as_ref().unwrap().used, 4.0); + assert_eq!(extra_window(&result, "requests"), "3 requests"); + assert_eq!(extra_window(&result, "tokens"), "45 tokens"); +} + +async fn fetch_with_page_body(body: &str) -> ProviderError { + let mut server = Server::new_async().await; + let _costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, body).await; + let _completions = mock_page(&mut server, COMPLETIONS_PATH, Matcher::Any, EMPTY_PAGE).await; + provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("the page should be rejected") +} + +#[tokio::test] +async fn openai_admin_usage_rejects_repeated_pagination_cursor() { + let error = fetch_with_page_body( + r#"{"object":"page","data":[],"has_more":true,"next_page":"same_page"}"#, + ) + .await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("cursor repeated")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_rejects_missing_pagination_cursor() { + for body in [ + r#"{"object":"page","data":[],"has_more":true,"next_page":null}"#, + r#"{"object":"page","data":[],"has_more":true,"next_page":" "}"#, + r#"{"object":"page","data":[],"has_more":true}"#, + ] { + let error = fetch_with_page_body(body).await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("cursor missing")), + "{body}: {error}" + ); + } +} + +#[tokio::test] +async fn openai_admin_usage_rejects_malformed_pages() { + for body in [ + r#"{"object":"page","has_more":false,"next_page":null}"#, + r#"{"object":"page","data":[],"next_page":null}"#, + r#"{"object":"page","data":[],"has_more":false,"next_page":7}"#, + ] { + let error = fetch_with_page_body(body).await; + assert!(matches!(error, ProviderError::Parse(_)), "{body}: {error}"); + } +} + +#[tokio::test] +async fn openai_admin_usage_rejects_pagination_beyond_one_hundred_pages() { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_body_from_request(|request| { + let path = request.path_and_query(); + let current = path + .split("page=p") + .nth(1) + .and_then(|rest| rest.split('&').next()) + .and_then(|number| number.parse::().ok()) + .unwrap_or(0); + format!( + r#"{{"data":[],"has_more":true,"next_page":"p{}"}}"#, + current + 1 + ) + .into_bytes() + }) + .expect(MAX_PAGES_PER_RANGE) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("the endless cursor chain should be rejected"); + + costs.assert_async().await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("exceeded 100 pages")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_retries_transient_completions_failure_once() { + let mut server = Server::new_async().await; + let costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(1); + let failure = server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::Any) + .with_status(503) + .expect(1) + .create_async() + .await; + let success = mock_page( + &mut server, + COMPLETIONS_PATH, + Matcher::Any, + r#"{"object":"page","data":[{"start_time":1700000000,"end_time":1700086400,"results":[{"input_tokens":10,"output_tokens":5,"num_model_requests":1,"model":"gpt-5.2"}]}],"has_more":false,"next_page":null}"#, + ) + .await + .expect(1); + + let result = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + failure.assert_async().await; + success.assert_async().await; + assert_eq!(extra_window(&result, "tokens"), "15 tokens"); + assert_eq!(extra_window(&result, "requests"), "1 requests"); +} + +#[tokio::test] +async fn openai_admin_usage_retries_once_then_reports_the_status() { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(429) + .with_header("Retry-After", "600") + .with_body(r#"{"error":{"message":"slow down"}}"#) + .expect(2) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("a persistent 429 should fail"); + + costs.assert_async().await; + assert!( + matches!(&error, ProviderError::Other(message) + if message.contains("429") && message.contains("slow down")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_never_retries_auth_failures() { + for status in [401, 403] { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(status) + .expect(1) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("auth failures should surface"); + + costs.assert_async().await; + assert!(matches!(error, ProviderError::AuthRequired), "{status}"); + } +} + +const GRANTS_BODY: &str = r#"{"total_granted":100.0,"total_used":25.0,"total_available":75.0}"#; + +async fn mock_admin_status(server: &mut ServerGuard, status: usize) -> mockito::Mock { + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(status) + .create_async() + .await +} + +#[tokio::test] +async fn openai_unscoped_key_falls_back_to_balance_on_admin_auth_failure() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + + let result = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .unwrap(); + + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); + assert_eq!( + result.usage.login_method.as_deref(), + Some("API balance: $75.00") + ); +} + +#[tokio::test] +async fn openai_unscoped_key_falls_back_to_balance_during_admin_outage() { + let mut server = Server::new_async().await; + let costs = mock_admin_status(&mut server, 500).await.expect(2); + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + + let result = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); +} + +#[tokio::test] +async fn openai_project_scoped_admin_key_never_uses_the_unfiltered_balance() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(0) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_or_balance(&admin_credential(), Some("proj_abc"), fixed_now(3_600)) + .await + .expect_err("a project-scoped admin failure should not fall back"); + + grants.assert_async().await; + assert!(matches!(error, ProviderError::AuthRequired)); +} + +#[tokio::test] +async fn openai_project_scoped_plain_key_keeps_the_balance_fallback() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + let credential = ApiCredential { + key: "sk-plain".to_string(), + is_admin: false, + }; + + let result = provider(&server) + .fetch_admin_or_balance(&credential, Some("proj_abc"), fixed_now(3_600)) + .await + .unwrap(); + + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); +} + +#[tokio::test] +async fn openai_balance_failure_keeps_the_admin_outage_error() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 500).await; + let _grants = server + .mock("GET", GRANTS_PATH) + .with_status(404) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .expect_err("both endpoints failed"); + + assert!( + matches!(&error, ProviderError::Other(message) if message.contains("costs returned status 500")), + "{error}" + ); +} From 4393cd437ba94770c5ef2838cc26fb79fae965e3 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:48:32 +0700 Subject: [PATCH 112/407] Port upstream 0.66.0: add hand-written Chromium LevelDB and Snappy reader --- rust/src/browser/leveldb/local_storage.rs | 70 +++ rust/src/browser/leveldb/log.rs | 114 ++++ rust/src/browser/leveldb/mod.rs | 134 +++++ rust/src/browser/leveldb/snappy.rs | 116 ++++ rust/src/browser/leveldb/table.rs | 208 +++++++ rust/src/browser/leveldb/tests.rs | 625 ++++++++++++++++++++++ rust/src/browser/leveldb/varint.rs | 41 ++ rust/src/browser/mod.rs | 1 + 8 files changed, 1309 insertions(+) create mode 100644 rust/src/browser/leveldb/local_storage.rs create mode 100644 rust/src/browser/leveldb/log.rs create mode 100644 rust/src/browser/leveldb/mod.rs create mode 100644 rust/src/browser/leveldb/snappy.rs create mode 100644 rust/src/browser/leveldb/table.rs create mode 100644 rust/src/browser/leveldb/tests.rs create mode 100644 rust/src/browser/leveldb/varint.rs diff --git a/rust/src/browser/leveldb/local_storage.rs b/rust/src/browser/leveldb/local_storage.rs new file mode 100644 index 0000000000..a2e3d5236e --- /dev/null +++ b/rust/src/browser/leveldb/local_storage.rs @@ -0,0 +1,70 @@ +//! Chromium `Local Storage` layer on top of the LevelDB reader. +//! +//! Chromium stores each origin's `localStorage` items under the key +//! `_\0` where `` is the serialized origin (`https://www.kimi.ai`) and +//! `` starts with a format byte: `0x00` for UTF-16LE text, `0x01` for Latin-1. Values use +//! the same format byte. Other keys in the database (`VERSION`, `META:`, +//! `METAACCESS:`) are bookkeeping and are ignored. + +use super::{Entry, LevelDbError, read_entries}; +use std::path::{Path, PathBuf}; + +const KEY_PREFIX: u8 = b'_'; +const ORIGIN_TERMINATOR: u8 = 0; +const FORMAT_UTF16LE: u8 = 0; +const FORMAT_LATIN1: u8 = 1; + +/// One decoded `localStorage` item. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LocalStorageEntry { + pub key: String, + pub value: String, +} + +/// 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") +} + +/// Read the `localStorage` items stored for `origin` (for example `https://www.kimi.ai`) in the +/// LevelDB directory `dir`, sorted by key. A trailing slash on `origin` is ignored. +pub fn read_local_storage_entries( + dir: &Path, + origin: &str, +) -> Result, LevelDbError> { + Ok(decode_origin_entries(&read_entries(dir)?, origin)) +} + +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); + + entries + .iter() + .filter_map(|entry| { + let key = decode_text(entry.key.strip_prefix(prefix.as_slice())?)?; + let value = decode_text(&entry.value)?; + Some(LocalStorageEntry { key, value }) + }) + .collect() +} + +/// 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()?; + match format { + FORMAT_LATIN1 => Some(data.iter().map(|&byte| char::from(byte)).collect()), + FORMAT_UTF16LE if data.len() % 2 == 0 => { + let units: Vec = data + .as_chunks::<2>() + .0 + .iter() + .map(|pair| u16::from_le_bytes(*pair)) + .collect(); + String::from_utf16(&units).ok() + } + _ => None, + } +} diff --git a/rust/src/browser/leveldb/log.rs b/rust/src/browser/leveldb/log.rs new file mode 100644 index 0000000000..98cf322948 --- /dev/null +++ b/rust/src/browser/leveldb/log.rs @@ -0,0 +1,114 @@ +//! LevelDB write-ahead log (`*.log`) reader. +//! +//! A log is a sequence of 32 KiB blocks holding physical records (7 byte header: masked CRC32C, +//! length, type) that are reassembled into logical records. Each logical record is a write batch: +//! `sequence: u64le`, `count: u32le`, then `count` operations (`1 key value` puts and `0 key` +//! deletes, with varint-length-prefixed byte strings). Operation `i` has sequence `sequence + i`. +//! +//! The live log of a running browser routinely ends in a half-written record, so a truncated or +//! 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}; + +const BLOCK_SIZE: usize = 32 * 1024; +const HEADER_SIZE: usize = 7; +const BATCH_HEADER_SIZE: usize = 12; + +const TYPE_ZERO: u8 = 0; +const TYPE_FULL: u8 = 1; +const TYPE_FIRST: u8 = 2; +const TYPE_MIDDLE: u8 = 3; +const TYPE_LAST: u8 = 4; + +const OP_DELETE: u8 = 0; +const OP_PUT: u8 = 1; + +/// Feed every put/delete found in `data` to `emit`. +pub(super) fn read_log(data: &[u8], emit: &mut impl FnMut(Record)) { + let mut assembled: Vec = Vec::new(); + let mut in_fragmented = false; + + let mut offset = 0usize; + while offset < data.len() { + let block_remaining = BLOCK_SIZE - (offset % BLOCK_SIZE); + if block_remaining < HEADER_SIZE { + offset += block_remaining; + continue; + } + let Some(header) = data.get(offset..offset + HEADER_SIZE) else { + return; + }; + let length = usize::from(u16::from_le_bytes([header[4], header[5]])); + let kind = header[6]; + if kind == TYPE_ZERO && length == 0 { + // Zero padding: the rest of this block is unused. + offset += block_remaining; + continue; + } + let payload_start = offset + HEADER_SIZE; + if HEADER_SIZE + length > block_remaining { + return; + } + let Some(payload) = data.get(payload_start..payload_start + length) else { + return; + }; + offset = payload_start + length; + + match kind { + TYPE_FULL => { + assembled.clear(); + in_fragmented = false; + read_batch(payload, emit); + } + TYPE_FIRST => { + assembled.clear(); + assembled.extend_from_slice(payload); + in_fragmented = true; + } + TYPE_MIDDLE if in_fragmented => assembled.extend_from_slice(payload), + TYPE_LAST if in_fragmented => { + assembled.extend_from_slice(payload); + in_fragmented = false; + read_batch(&assembled, emit); + assembled.clear(); + } + _ => return, + } + } +} + +fn read_batch(batch: &[u8], emit: &mut impl FnMut(Record)) { + let (Some(sequence), Some(count)) = (read_u64_le(batch, 0), read_u32_le(batch, 8)) else { + return; + }; + 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; + }; + let Some((key, after_key)) = read_length_prefixed(after_op) else { + return; + }; + let value = match op { + OP_PUT => { + let Some((value, after_value)) = read_length_prefixed(after_key) else { + return; + }; + rest = after_value; + Some(value.to_vec()) + } + OP_DELETE => { + rest = after_key; + None + } + _ => return, + }; + emit(Record { + key: key.to_vec(), + sequence: sequence.wrapping_add(index), + value, + }); + } +} diff --git a/rust/src/browser/leveldb/mod.rs b/rust/src/browser/leveldb/mod.rs new file mode 100644 index 0000000000..19af9081dd --- /dev/null +++ b/rust/src/browser/leveldb/mod.rs @@ -0,0 +1,134 @@ +//! Read-only LevelDB reader for Chromium browser storage (Local Storage, and other LevelDB +//! directories a browser keeps under a profile). +//! +//! This is a small hand-written reader rather than a dependency: it understands the write-ahead +//! log and sorted-table formats, including Snappy-compressed table blocks, and resolves each user +//! key to its newest value by sequence number. It does not open the database, take its `LOCK`, or +//! write anything, so it can be pointed at a profile of a running browser. +//! +//! Known limits, all acceptable for a best-effort credential import: +//! - The `MANIFEST` is not consulted, so a table that compaction already obsoleted but the +//! browser has not yet removed can still be read. Newer sequence numbers win, so this only +//! matters for a deleted key whose tombstone was compacted away. +//! - 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. + +pub mod local_storage; +mod log; +pub mod snappy; +mod table; +mod varint; + +#[cfg(test)] +mod tests; + +use std::collections::BTreeMap; +use std::ffi::OsStr; +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. +pub const MAX_BLOCK_BYTES: usize = 16 * 1024 * 1024; + +/// One live key/value pair of the database. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Entry { + pub key: Vec, + pub value: Vec, +} + +/// 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), +} + +/// A put (`value: Some`) or delete (`value: None`) with the sequence number it was written at. +pub(crate) struct Record { + pub(crate) key: Vec, + pub(crate) sequence: u64, + pub(crate) value: Option>, +} + +/// 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. +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)); + } + }; + + for dir_entry in std::fs::read_dir(dir)? { + 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) { + Ok(data) => data, + Err(error) => { + tracing::debug!(file = ?path.file_name(), %error, "skipping unreadable LevelDB file"); + continue; + } + }; + 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"); + } + }, + } + } + + Ok(newest + .into_iter() + .filter_map(|(key, (_, value))| value.map(|value| Entry { key, value })) + .collect()) +} + +#[derive(Clone, Copy)] +enum FileKind { + Log, + Table, +} + +impl FileKind { + fn of(path: &Path) -> Option { + let extension = path.extension().and_then(OsStr::to_str)?; + if extension.eq_ignore_ascii_case("log") { + Some(Self::Log) + } else if extension.eq_ignore_ascii_case("ldb") || extension.eq_ignore_ascii_case("sst") { + Some(Self::Table) + } else { + None + } + } +} + +fn read_bounded_file(path: &Path) -> std::io::Result> { + let size = std::fs::metadata(path)?.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) +} diff --git a/rust/src/browser/leveldb/snappy.rs b/rust/src/browser/leveldb/snappy.rs new file mode 100644 index 0000000000..3436a621df --- /dev/null +++ b/rust/src/browser/leveldb/snappy.rs @@ -0,0 +1,116 @@ +//! Raw Snappy block decompression (the format LevelDB uses for compressed table blocks). +//! +//! Format reference: . +//! The stream is a varint uncompressed length followed by literal and copy elements. +//! Output size is bounded by the caller so a hostile length preamble cannot exhaust memory. + +use super::varint::read_varint32; + +/// Why a Snappy block could not be decoded. +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum SnappyError { + #[error("snappy length preamble is missing or malformed")] + BadPreamble, + #[error("snappy output of {0} bytes exceeds the {1} byte limit")] + TooLarge(usize, usize), + #[error("snappy stream is truncated")] + Truncated, + #[error("snappy copy references data before the start of the output")] + BadOffset, + #[error("snappy stream length does not match its preamble")] + LengthMismatch, +} + +/// Decompress one raw Snappy block, refusing to produce more than `max_len` bytes. +pub fn decompress(input: &[u8], max_len: usize) -> Result, SnappyError> { + let (expected, mut pos) = read_varint32(input).ok_or(SnappyError::BadPreamble)?; + let expected = expected as usize; + if expected > max_len { + return Err(SnappyError::TooLarge(expected, max_len)); + } + + let mut out = Vec::with_capacity(expected); + while pos < input.len() { + let tag = input[pos]; + pos += 1; + match tag & 0b11 { + 0b00 => { + let mut len = usize::from(tag >> 2); + if len >= 60 { + let extra = len - 59; + let bytes = input.get(pos..pos + extra).ok_or(SnappyError::Truncated)?; + pos += extra; + len = bytes + .iter() + .rev() + .fold(0usize, |acc, b| (acc << 8) | usize::from(*b)); + } + len = len.checked_add(1).ok_or(SnappyError::Truncated)?; + let literal = pos + .checked_add(len) + .and_then(|end| input.get(pos..end)) + .ok_or(SnappyError::Truncated)?; + if out.len() + literal.len() > expected { + return Err(SnappyError::LengthMismatch); + } + out.extend_from_slice(literal); + pos += len; + } + kind => { + let (len, offset) = match kind { + 0b01 => { + let low = *input.get(pos).ok_or(SnappyError::Truncated)?; + pos += 1; + ( + 4 + usize::from((tag >> 2) & 0b111), + (usize::from(tag >> 5) << 8) | usize::from(low), + ) + } + 0b10 => { + let bytes = input.get(pos..pos + 2).ok_or(SnappyError::Truncated)?; + pos += 2; + ( + 1 + usize::from(tag >> 2), + usize::from(u16::from_le_bytes([bytes[0], bytes[1]])), + ) + } + _ => { + let bytes = input.get(pos..pos + 4).ok_or(SnappyError::Truncated)?; + pos += 4; + ( + 1 + usize::from(tag >> 2), + u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize, + ) + } + }; + copy_within_output(&mut out, offset, len, expected)?; + } + } + } + + if out.len() == expected { + Ok(out) + } else { + Err(SnappyError::LengthMismatch) + } +} + +/// Append `len` bytes copied from `offset` bytes back; the ranges may overlap (run-length style). +fn copy_within_output( + out: &mut Vec, + offset: usize, + len: usize, + expected: usize, +) -> Result<(), SnappyError> { + if offset == 0 || offset > out.len() { + return Err(SnappyError::BadOffset); + } + if out.len() + len > expected { + return Err(SnappyError::LengthMismatch); + } + let start = out.len() - offset; + for i in 0..len { + out.push(out[start + i]); + } + Ok(()) +} diff --git a/rust/src/browser/leveldb/table.rs b/rust/src/browser/leveldb/table.rs new file mode 100644 index 0000000000..0bf805c6bf --- /dev/null +++ b/rust/src/browser/leveldb/table.rs @@ -0,0 +1,208 @@ +//! LevelDB sorted-table (`*.ldb` / `*.sst`) reader. +//! +//! Layout: data blocks, a metaindex block, an index block, and a 48 byte footer that ends with a +//! magic number. The footer holds block handles (varint offset, varint size) for the metaindex and +//! index blocks; index entries map a separator key to the handle of a data block. Every block is +//! followed by a 5 byte trailer (compression type, CRC32C). Entries inside a block share key +//! prefixes with the previous entry, and a restart array closes the block. +//! +//! Only the index and data blocks are read; the filter/metaindex blocks are irrelevant to a full +//! scan. Data-block keys are internal keys: the user key followed by `sequence << 8 | kind`. +//! Checksums are not verified (see `log.rs`); structural checks keep malformed input from +//! panicking or over-allocating. + +use super::snappy; +use super::varint::{read_u32_le, read_u64_le, read_varint32, read_varint64}; +use super::{MAX_BLOCK_BYTES, Record}; + +const FOOTER_SIZE: usize = 48; +const TABLE_MAGIC: u64 = 0xdb47_7524_8b80_fb57; +const BLOCK_TRAILER_SIZE: usize = 5; + +const COMPRESSION_NONE: u8 = 0; +const COMPRESSION_SNAPPY: u8 = 1; + +const KIND_DELETE: u64 = 0; +const KIND_VALUE: u64 = 1; + +#[derive(Debug, thiserror::Error)] +pub(super) enum TableError { + #[error("file is too short or has a bad table footer")] + BadFooter, + #[error("block handle points outside the file")] + BadHandle, + #[error("unsupported block compression type {0}")] + UnsupportedCompression(u8), + #[error("block is malformed")] + BadBlock, + #[error("snappy block failed to decode: {0}")] + Snappy(#[from] snappy::SnappyError), +} + +#[derive(Clone, Copy)] +struct BlockHandle { + offset: usize, + size: usize, +} + +impl BlockHandle { + /// Parse a handle from the front of `input`; returns it with the bytes consumed. + fn parse(input: &[u8]) -> Option<(Self, usize)> { + let (offset, first) = read_varint64(input)?; + let (size, second) = read_varint64(input.get(first..)?)?; + Some(( + Self { + offset: usize::try_from(offset).ok()?, + size: usize::try_from(size).ok()?, + }, + first + second, + )) + } +} + +/// Feed every put/delete found in the table `data` to `emit`. +/// +/// 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. +pub(super) fn read_table(data: &[u8], emit: &mut impl FnMut(Record)) -> Result { + let footer_start = data + .len() + .checked_sub(FOOTER_SIZE) + .ok_or(TableError::BadFooter)?; + let footer = &data[footer_start..]; + if read_u64_le(footer, FOOTER_SIZE - 8) != Some(TABLE_MAGIC) { + return Err(TableError::BadFooter); + } + let (_metaindex, used) = BlockHandle::parse(footer).ok_or(TableError::BadFooter)?; + let (index, _) = BlockHandle::parse(&footer[used..]).ok_or(TableError::BadFooter)?; + + let index_block = read_block(data, index)?; + let mut skipped_blocks = 0usize; + for (_separator, handle_bytes) in BlockEntries::new(&index_block)? { + let entry = BlockHandle::parse(handle_bytes) + .map(|(handle, _)| handle) + .ok_or(TableError::BadBlock); + let block = entry.and_then(|handle| read_block(data, handle)); + let Ok(block) = block else { + skipped_blocks += 1; + continue; + }; + if emit_block_records(&block, emit).is_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)? { + 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_DELETE => None, + _ => continue, + }; + emit(Record { + key: user_key.to_vec(), + sequence: packed >> 8, + value, + }); + } + Ok(()) +} + +/// Return the decompressed contents of the block at `handle` (without its trailer). +fn read_block(data: &[u8], handle: BlockHandle) -> Result, TableError> { + 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 { + return Err(TableError::BadHandle); + } + let contents_end = end - BLOCK_TRAILER_SIZE; + let raw = data + .get(handle.offset..contents_end) + .ok_or(TableError::BadHandle)?; + let compression = data[contents_end]; + match compression { + COMPRESSION_NONE => Ok(raw.to_vec()), + COMPRESSION_SNAPPY => Ok(snappy::decompress(raw, MAX_BLOCK_BYTES)?), + other => Err(TableError::UnsupportedCompression(other)), + } +} + +/// Iterator over the `(key, value)` entries of a decoded block, undoing prefix compression. +struct BlockEntries<'a> { + entries: &'a [u8], + key: Vec, + failed: bool, +} + +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; + let restarts_len = restart_count + .checked_mul(4) + .and_then(|len| len.checked_add(4)) + .ok_or(TableError::BadBlock)?; + let entries_end = block + .len() + .checked_sub(restarts_len) + .ok_or(TableError::BadBlock)?; + Ok(Self { + entries: &block[..entries_end], + key: Vec::new(), + failed: false, + }) + } +} + +impl<'a> Iterator for BlockEntries<'a> { + type Item = (Vec, &'a [u8]); + + fn next(&mut self) -> Option { + if self.failed || self.entries.is_empty() { + return None; + } + let parsed = parse_entry(self.entries, &self.key); + match parsed { + Some((key, value, rest)) => { + self.key.clone_from(&key); + self.entries = rest; + Some((key, value)) + } + None => { + self.failed = true; + None + } + } + } +} + +/// Decode one entry: `shared`, `non_shared`, `value_len` varints, key delta, value. +fn parse_entry<'a>( + entries: &'a [u8], + previous_key: &[u8], +) -> Option<(Vec, &'a [u8], &'a [u8])> { + let (shared, a) = read_varint32(entries)?; + let (non_shared, b) = read_varint32(entries.get(a..)?)?; + let (value_len, c) = read_varint32(entries.get(a + b..)?)?; + let body = entries.get(a + b + c..)?; + let (shared, non_shared, value_len) = + (shared as usize, non_shared as usize, value_len as usize); + let payload_len = non_shared.checked_add(value_len)?; + if shared > previous_key.len() || body.len() < payload_len { + return None; + } + let mut key = Vec::with_capacity(shared + non_shared); + 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/browser/leveldb/tests.rs b/rust/src/browser/leveldb/tests.rs new file mode 100644 index 0000000000..7110e9d215 --- /dev/null +++ b/rust/src/browser/leveldb/tests.rs @@ -0,0 +1,625 @@ +//! Tests for the LevelDB reader. Fixtures are built in-memory from the documented on-disk +//! formats (log block/record framing, table block/footer layout, Snappy elements). + +#![allow( + clippy::cast_possible_truncation, + reason = "fixture builders encode small in-memory sizes into fixed-width fields" +)] + +use super::local_storage::{ + LocalStorageEntry, decode_origin_entries, local_storage_dir, read_local_storage_entries, +}; +use super::snappy::{self, SnappyError}; +use super::table::{TableError, read_table}; +use super::{Entry, Record, log, read_entries}; + +// --------------------------------------------------------------------------------------------- +// Fixture builders +// --------------------------------------------------------------------------------------------- + +fn varint(mut value: u64) -> Vec { + let mut out = Vec::new(); + while value >= 0x80 { + out.push((value as u8 & 0x7f) | 0x80); + value >>= 7; + } + out.push(value as u8); + out +} + +/// A valid Snappy stream made only of literal elements (no back-references). +fn snappy_literal(data: &[u8]) -> Vec { + let mut out = varint(data.len() as u64); + for chunk in data.chunks(60) { + out.push(((chunk.len() - 1) as u8) << 2); + out.extend_from_slice(chunk); + } + out +} + +enum Op<'a> { + Put(&'a [u8], &'a [u8]), + Delete(&'a [u8]), +} + +fn write_batch(sequence: u64, ops: &[Op]) -> Vec { + let mut out = sequence.to_le_bytes().to_vec(); + out.extend_from_slice(&(ops.len() as u32).to_le_bytes()); + for op in ops { + match op { + Op::Put(key, value) => { + out.push(1); + out.extend(varint(key.len() as u64)); + out.extend_from_slice(key); + out.extend(varint(value.len() as u64)); + out.extend_from_slice(value); + } + Op::Delete(key) => { + out.push(0); + out.extend(varint(key.len() as u64)); + out.extend_from_slice(key); + } + } + } + out +} + +/// Frame batches as a LevelDB log: 32 KiB blocks, fragmenting records that cross a block edge. +fn write_log(batches: &[Vec]) -> Vec { + const BLOCK: usize = 32 * 1024; + let mut out = Vec::new(); + for batch in batches { + let mut remaining = batch.as_slice(); + let mut first = true; + loop { + let left_in_block = BLOCK - out.len() % BLOCK; + if left_in_block < 7 { + out.extend(std::iter::repeat_n(0u8, left_in_block)); + continue; + } + let chunk_len = remaining.len().min(left_in_block - 7); + let last = chunk_len == remaining.len(); + let kind: u8 = match (first, last) { + (true, true) => 1, + (true, false) => 2, + (false, false) => 3, + (false, true) => 4, + }; + out.extend_from_slice(&[0, 0, 0, 0]); // checksum (not verified by the reader) + out.extend_from_slice(&(chunk_len as u16).to_le_bytes()); + out.push(kind); + out.extend_from_slice(&remaining[..chunk_len]); + remaining = &remaining[chunk_len..]; + first = false; + if last { + break; + } + } + } + out +} + +fn internal_key(user_key: &[u8], sequence: u64, kind: u8) -> Vec { + let mut key = user_key.to_vec(); + key.extend_from_slice(&((sequence << 8) | u64::from(kind)).to_le_bytes()); + key +} + +/// Encode a block with prefix compression and a restart every `restart_interval` entries. +fn write_block(entries: &[(Vec, Vec)], restart_interval: usize) -> Vec { + let mut out = Vec::new(); + let mut restarts = vec![0u32]; + let mut previous: &[u8] = &[]; + for (index, (key, value)) in entries.iter().enumerate() { + let shared = if index % restart_interval == 0 { + if index > 0 { + restarts.push(out.len() as u32); + } + 0 + } else { + previous.iter().zip(key).take_while(|(a, b)| a == b).count() + }; + out.extend(varint(shared as u64)); + out.extend(varint((key.len() - shared) as u64)); + out.extend(varint(value.len() as u64)); + out.extend_from_slice(&key[shared..]); + out.extend_from_slice(value); + previous = key; + } + for restart in &restarts { + out.extend_from_slice(&restart.to_le_bytes()); + } + out.extend_from_slice(&(restarts.len() as u32).to_le_bytes()); + out +} + +#[derive(Clone, Copy, PartialEq)] +enum Compression { + None, + Snappy, +} + +/// A table record: user key, sequence, kind (1 = value, 0 = delete), value. +type TableRecord<'a> = (&'a [u8], u64, u8, &'a [u8]); + +/// Build a table with one data block per slice of `blocks`. +fn write_table(blocks: &[(&[TableRecord], Compression)]) -> Vec { + let mut file = Vec::new(); + let mut index_entries = Vec::new(); + for (records, compression) in blocks { + let entries: Vec<_> = records + .iter() + .map(|(key, sequence, kind, value)| { + (internal_key(key, *sequence, *kind), value.to_vec()) + }) + .collect(); + let block = write_block(&entries, 2); + let (contents, tag) = match compression { + Compression::None => (block, 0u8), + Compression::Snappy => (snappy_literal(&block), 1u8), + }; + let mut handle = varint(file.len() as u64); + handle.extend(varint(contents.len() as u64)); + file.extend_from_slice(&contents); + file.push(tag); + file.extend_from_slice(&[0, 0, 0, 0]); + index_entries.push((entries.last().unwrap().0.clone(), handle)); + } + + let index = write_block(&index_entries, 1); + let mut index_handle = varint(file.len() as u64); + index_handle.extend(varint(index.len() as u64)); + file.extend_from_slice(&index); + file.extend_from_slice(&[0, 0, 0, 0, 0]); + + let mut footer = varint(0); // metaindex handle: offset 0, size 0 (never read) + footer.extend(varint(0)); + footer.extend_from_slice(&index_handle); + footer.resize(40, 0); + footer.extend_from_slice(&0xdb47_7524_8b80_fb57u64.to_le_bytes()); + file.extend_from_slice(&footer); + file +} + +type Collected = Vec<(Vec, u64, Option>)>; + +fn collect_log(data: &[u8]) -> Collected { + let mut records = Vec::new(); + log::read_log(data, &mut |r: Record| { + records.push((r.key, r.sequence, r.value)); + }); + records +} + +fn collect_table(data: &[u8]) -> Result { + let mut records = Vec::new(); + read_table(data, &mut |r: Record| { + records.push((r.key, r.sequence, r.value)); + })?; + Ok(records) +} + +// --------------------------------------------------------------------------------------------- +// Snappy +// --------------------------------------------------------------------------------------------- + +#[test] +fn snappy_decodes_literal() { + let mut stream = vec![5, 4 << 2]; + stream.extend_from_slice(b"hello"); + assert_eq!(snappy::decompress(&stream, 1024).unwrap(), b"hello"); +} + +#[test] +fn snappy_decodes_long_literal_with_extra_length_byte() { + let data: Vec = (0..100u8).collect(); + let mut stream = varint(100); + stream.extend_from_slice(&[60 << 2, 99]); + stream.extend_from_slice(&data); + assert_eq!(snappy::decompress(&stream, 1024).unwrap(), data); +} + +#[test] +fn snappy_decodes_overlapping_copy_with_one_byte_offset() { + // "abc" literal, then copy of length 9 at offset 3 (copy-1 tag: len-4 in bits 2..5). + let stream = [12, 2 << 2, b'a', b'b', b'c', 0b01 | (5 << 2), 3]; + assert_eq!(snappy::decompress(&stream, 1024).unwrap(), b"abcabcabcabc"); +} + +#[test] +fn snappy_decodes_copy_with_two_byte_offset() { + let mut stream = vec![15, 9 << 2]; + stream.extend_from_slice(b"0123456789"); + stream.extend_from_slice(&[0b10 | (4 << 2), 10, 0]); + assert_eq!( + snappy::decompress(&stream, 1024).unwrap(), + b"012345678901234" + ); +} + +#[test] +fn snappy_decodes_copy_with_four_byte_offset() { + let mut stream = vec![14, 9 << 2]; + stream.extend_from_slice(b"0123456789"); + stream.extend_from_slice(&[0b11 | (3 << 2), 10, 0, 0, 0]); + assert_eq!( + snappy::decompress(&stream, 1024).unwrap(), + b"01234567890123" + ); +} + +#[test] +fn snappy_rejects_malformed_streams() { + assert_eq!(snappy::decompress(&[], 16), Err(SnappyError::BadPreamble)); + // Declared length larger than the limit is refused before allocating. + assert_eq!( + snappy::decompress(&varint(u64::from(u32::MAX)), 1024), + Err(SnappyError::TooLarge(u32::MAX as usize, 1024)) + ); + // Literal runs past the end of input. + assert_eq!( + snappy::decompress(&[5, 4 << 2, b'h'], 16), + Err(SnappyError::Truncated) + ); + // Copy before any output exists. + assert_eq!( + snappy::decompress(&[4, 0b01, 1], 16), + Err(SnappyError::BadOffset) + ); + // Zero offset is invalid. + assert_eq!( + snappy::decompress(&[5, 0, b'a', 0b01, 0], 16), + Err(SnappyError::BadOffset) + ); + // Output shorter than the preamble promised. + assert_eq!( + snappy::decompress(&[5, 0, b'a'], 16), + Err(SnappyError::LengthMismatch) + ); + // Output longer than the preamble promised. + assert_eq!( + snappy::decompress(&[1, 1 << 2, b'a', b'b'], 16), + Err(SnappyError::LengthMismatch) + ); +} + +// --------------------------------------------------------------------------------------------- +// Write-ahead log +// --------------------------------------------------------------------------------------------- + +#[test] +fn log_yields_puts_and_deletes_with_incrementing_sequences() { + let data = write_log(&[write_batch( + 10, + &[Op::Put(b"a", b"1"), Op::Delete(b"b"), Op::Put(b"c", b"3")], + )]); + assert_eq!( + collect_log(&data), + vec![ + (b"a".to_vec(), 10, Some(b"1".to_vec())), + (b"b".to_vec(), 11, None), + (b"c".to_vec(), 12, Some(b"3".to_vec())), + ] + ); +} + +#[test] +fn log_reassembles_records_that_span_blocks() { + let big = vec![0xabu8; 80 * 1024]; + let data = write_log(&[ + write_batch(1, &[Op::Put(b"small", b"x")]), + write_batch(2, &[Op::Put(b"big", &big)]), + write_batch(3, &[Op::Put(b"after", b"y")]), + ]); + assert!(data.len() > 64 * 1024); + let records = collect_log(&data); + assert_eq!(records.len(), 3); + assert_eq!(records[1].2.as_deref(), Some(big.as_slice())); + assert_eq!(records[2].0, b"after"); +} + +#[test] +fn log_skips_block_trailer_padding() { + // Leave fewer than 7 bytes at the end of the first block so the writer pads it. + let filler = vec![1u8; 32 * 1024 - 7 - 12 - 1 - 2 - 1 - 1 - 4]; + let data = write_log(&[ + write_batch(1, &[Op::Put(b"k", &filler)]), + write_batch(2, &[Op::Put(b"z", b"after-padding")]), + ]); + let records = collect_log(&data); + assert_eq!(records.len(), 2); + assert_eq!(records[1].2.as_deref(), Some(b"after-padding".as_slice())); +} + +#[test] +fn log_keeps_records_before_a_truncated_tail() { + let mut data = write_log(&[ + write_batch(1, &[Op::Put(b"kept", b"1")]), + write_batch(2, &[Op::Put(b"lost", b"2")]), + ]); + data.truncate(data.len() - 3); + let records = collect_log(&data); + assert_eq!(records.len(), 1); + assert_eq!(records[0].0, b"kept"); +} + +#[test] +fn log_ignores_garbage_without_panicking() { + assert!(collect_log(&[0xff; 100]).is_empty()); + assert!(collect_log(&[]).is_empty()); + // Middle fragment with no preceding first fragment. + let orphan = [0, 0, 0, 0, 1, 0, 3, 9]; + assert!(collect_log(&orphan).is_empty()); +} + +// --------------------------------------------------------------------------------------------- +// Tables +// --------------------------------------------------------------------------------------------- + +#[test] +fn table_reads_prefix_compressed_multi_block_data() { + let first: &[TableRecord] = &[ + (b"app:alpha", 5, 1, b"one"), + (b"app:alpha2", 6, 1, b"two"), + (b"app:beta", 7, 1, b"three"), + ]; + let second: &[TableRecord] = &[(b"zeta", 8, 1, b"four"), (b"zeta-gone", 9, 0, b"")]; + let data = write_table(&[(first, Compression::None), (second, Compression::None)]); + + let records = collect_table(&data).unwrap(); + assert_eq!(records.len(), 5); + assert_eq!( + records[1], + (b"app:alpha2".to_vec(), 6, Some(b"two".to_vec())) + ); + assert_eq!(records[4], (b"zeta-gone".to_vec(), 9, None)); +} + +#[test] +fn table_reads_snappy_compressed_blocks() { + let long_value = vec![b'v'; 500]; + let records: &[TableRecord] = &[(b"k1", 1, 1, &long_value), (b"k2", 2, 1, b"short")]; + let data = write_table(&[(records, Compression::Snappy)]); + + let decoded = collect_table(&data).unwrap(); + assert_eq!(decoded.len(), 2); + assert_eq!(decoded[0].2.as_deref(), Some(long_value.as_slice())); +} + +#[test] +fn table_with_bad_magic_or_length_is_rejected() { + let records: &[TableRecord] = &[(b"k", 1, 1, b"v")]; + let mut data = write_table(&[(records, Compression::None)]); + assert!(matches!( + collect_table(&data[..20]), + Err(TableError::BadFooter) + )); + let last = data.len() - 1; + data[last] ^= 0xff; + assert!(matches!(collect_table(&data), Err(TableError::BadFooter))); +} + +#[test] +fn table_skips_an_undecodable_block_but_keeps_the_rest() { + let good: &[TableRecord] = &[(b"good", 1, 1, b"ok")]; + let bad: &[TableRecord] = &[(b"bad", 2, 1, b"nope")]; + let mut data = write_table(&[(bad, Compression::Snappy), (good, Compression::None)]); + // Corrupt the first block's compression tag to an unsupported type (2). + let bad_block_len = snappy_literal(&write_block( + &[(internal_key(b"bad", 2, 1), b"nope".to_vec())], + 2, + )) + .len(); + data[bad_block_len] = 2; + + let mut records = Vec::new(); + let skipped = read_table(&data, &mut |r: Record| records.push(r.key)).unwrap(); + assert_eq!(skipped, 1); + assert_eq!(records, vec![b"good".to_vec()]); +} + +#[test] +fn table_with_out_of_range_handle_does_not_panic() { + let records: &[TableRecord] = &[(b"k", 1, 1, b"v")]; + let mut data = write_table(&[(records, Compression::None)]); + // Point the index handle far past the end of the file. + let footer = data.len() - 48; + data[footer..footer + 40].fill(0); + data[footer + 2..footer + 6].copy_from_slice(&[0xff, 0xff, 0xff, 0x0f]); + assert!(collect_table(&data).is_err()); +} + +// --------------------------------------------------------------------------------------------- +// Directory reader +// --------------------------------------------------------------------------------------------- + +#[test] +fn read_entries_resolves_newest_sequence_across_log_and_table() { + let dir = tempfile::tempdir().unwrap(); + let table_records: &[TableRecord] = &[ + (b"keep", 1, 1, b"table-keep"), + (b"stale", 2, 1, b"table-old"), + (b"removed", 3, 1, b"table-removed"), + ]; + std::fs::write( + dir.path().join("000005.ldb"), + write_table(&[(table_records, Compression::None)]), + ) + .unwrap(); + std::fs::write( + dir.path().join("000006.log"), + write_log(&[write_batch( + 10, + &[ + Op::Put(b"stale", b"log-new"), + Op::Delete(b"removed"), + Op::Put(b"fresh", b"log-fresh"), + ], + )]), + ) + .unwrap(); + // Bookkeeping files and unrelated data are ignored. + std::fs::write(dir.path().join("LOCK"), b"").unwrap(); + std::fs::write(dir.path().join("CURRENT"), b"MANIFEST-000001\n").unwrap(); + std::fs::write(dir.path().join("notes.txt"), b"not a database").unwrap(); + + let entries = read_entries(dir.path()).unwrap(); + let pairs: Vec<(&[u8], &[u8])> = entries + .iter() + .map(|e| (e.key.as_slice(), e.value.as_slice())) + .collect(); + assert_eq!( + pairs, + vec![ + (b"fresh".as_slice(), b"log-fresh".as_slice()), + (b"keep".as_slice(), b"table-keep".as_slice()), + (b"stale".as_slice(), b"log-new".as_slice()), + ] + ); +} + +#[test] +fn read_entries_lets_a_newer_table_delete_beat_an_older_log_put() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("000001.log"), + write_log(&[write_batch(1, &[Op::Put(b"gone", b"old")])]), + ) + .unwrap(); + let deletes: &[TableRecord] = &[(b"gone", 2, 0, b"")]; + std::fs::write( + dir.path().join("000002.ldb"), + write_table(&[(deletes, Compression::None)]), + ) + .unwrap(); + assert!(read_entries(dir.path()).unwrap().is_empty()); +} + +#[test] +fn read_entries_skips_corrupt_files_and_keeps_good_ones() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("000001.ldb"), b"definitely not a table").unwrap(); + std::fs::write( + dir.path().join("000002.log"), + write_log(&[write_batch(1, &[Op::Put(b"ok", b"1")])]), + ) + .unwrap(); + assert_eq!( + read_entries(dir.path()).unwrap(), + vec![Entry { + key: b"ok".to_vec(), + value: b"1".to_vec() + }] + ); +} + +#[test] +fn read_entries_reports_a_missing_directory() { + let dir = tempfile::tempdir().unwrap(); + assert!(read_entries(&dir.path().join("absent")).is_err()); +} + +// --------------------------------------------------------------------------------------------- +// Chromium Local Storage layer +// --------------------------------------------------------------------------------------------- + +fn latin1(text: &str) -> Vec { + let mut out = vec![1u8]; + out.extend(text.chars().map(|c| c as u8)); + out +} + +fn utf16(text: &str) -> Vec { + let mut out = vec![0u8]; + out.extend(text.encode_utf16().flat_map(u16::to_le_bytes)); + out +} + +fn storage_key(origin: &str, encoded_key: &[u8]) -> Vec { + let mut key = format!("_{origin}").into_bytes(); + key.push(0); + key.extend_from_slice(encoded_key); + key +} + +#[test] +fn local_storage_decodes_latin1_and_utf16_items_for_one_origin() { + let entries = vec![ + Entry { + key: b"META:https://www.kimi.ai".to_vec(), + value: b"\x08\x01".to_vec(), + }, + Entry { + key: storage_key("https://www.kimi.ai", &latin1("access_token")), + value: latin1("eyJ.body.sig"), + }, + Entry { + key: storage_key("https://www.kimi.ai", &utf16("name")), + value: utf16("Zoë \u{1F600}"), + }, + Entry { + key: storage_key("https://www.kimi.com", &latin1("access_token")), + value: latin1("other-origin"), + }, + Entry { + key: storage_key("https://www.kimi.ai", &latin1("bad_format")), + value: vec![7, b'x'], + }, + ]; + + let decoded = decode_origin_entries(&entries, "https://www.kimi.ai/"); + assert_eq!( + decoded, + vec![ + LocalStorageEntry { + key: "access_token".into(), + value: "eyJ.body.sig".into() + }, + LocalStorageEntry { + key: "name".into(), + value: "Zoë \u{1F600}".into() + }, + ] + ); +} + +#[test] +fn local_storage_origin_match_is_exact_not_a_prefix() { + let entries = vec![Entry { + key: storage_key("https://www.kimi.ai.evil.example", &latin1("access_token")), + value: latin1("nope"), + }]; + assert!(decode_origin_entries(&entries, "https://www.kimi.ai").is_empty()); +} + +#[test] +fn local_storage_rejects_odd_length_utf16_values() { + let entries = vec![Entry { + key: storage_key("https://a.example", &latin1("k")), + value: vec![0, b'x'], + }]; + assert!(decode_origin_entries(&entries, "https://a.example").is_empty()); +} + +#[test] +fn read_local_storage_entries_reads_a_profile_directory() { + let profile = tempfile::tempdir().unwrap(); + let dir = local_storage_dir(profile.path()); + std::fs::create_dir_all(&dir).unwrap(); + + let key = storage_key("https://www.kimi.ai", &latin1("access_token")); + let value = latin1("a.b.c"); + std::fs::write( + dir.join("000003.log"), + write_log(&[write_batch(4, &[Op::Put(&key, &value)])]), + ) + .unwrap(); + + assert_eq!( + read_local_storage_entries(&dir, "https://www.kimi.ai").unwrap(), + vec![LocalStorageEntry { + key: "access_token".into(), + value: "a.b.c".into() + }] + ); +} diff --git a/rust/src/browser/leveldb/varint.rs b/rust/src/browser/leveldb/varint.rs new file mode 100644 index 0000000000..c51511a711 --- /dev/null +++ b/rust/src/browser/leveldb/varint.rs @@ -0,0 +1,41 @@ +//! LevelDB varint and fixed-width helpers. + +/// Decode a little-endian base-128 varint of at most 32 bits; returns the value and bytes used. +pub(super) fn read_varint32(input: &[u8]) -> Option<(u32, usize)> { + let (value, used) = read_varint64(input)?; + Some((u32::try_from(value).ok()?, used)) +} + +/// Decode a little-endian base-128 varint of at most 64 bits; returns the value and bytes used. +pub(super) fn read_varint64(input: &[u8]) -> Option<(u64, usize)> { + let mut value = 0u64; + for (index, byte) in input.iter().take(10).enumerate() { + let bits = u64::from(byte & 0x7f); + if index == 9 && bits > 1 { + return None; + } + value |= bits << (7 * index); + if byte & 0x80 == 0 { + return Some((value, index + 1)); + } + } + None +} + +/// Split a varint-length-prefixed slice off the front of `input`, returning `(slice, rest)`. +pub(super) fn read_length_prefixed(input: &[u8]) -> Option<(&[u8], &[u8])> { + let (len, used) = read_varint32(input)?; + let rest = &input[used..]; + let len = len as usize; + (rest.len() >= len).then(|| rest.split_at(len)) +} + +pub(super) fn read_u32_le(input: &[u8], at: usize) -> Option { + let bytes = input.get(at..at.checked_add(4)?)?; + Some(u32::from_le_bytes(bytes.try_into().ok()?)) +} + +pub(super) fn read_u64_le(input: &[u8], at: usize) -> Option { + let bytes = input.get(at..at.checked_add(8)?)?; + Some(u64::from_le_bytes(bytes.try_into().ok()?)) +} diff --git a/rust/src/browser/mod.rs b/rust/src/browser/mod.rs index f428e1b41a..035fb6e6d2 100755 --- a/rust/src/browser/mod.rs +++ b/rust/src/browser/mod.rs @@ -3,6 +3,7 @@ pub mod cookie_cache; pub mod cookies; pub mod detection; +pub mod leveldb; pub mod watchdog; pub mod wsl_paths; From cc07049a85b7e8c4b517a757e40c75285bf087c6 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:53:57 +0700 Subject: [PATCH 113/407] Port upstream 0.67.0: count Mistral plan-covered tokens from chat and Vibe Code Token totals now count consumed units (value, falling back to value_paid) across completion, chat, and vibe_code.completion; spend keeps counting billed units (value_paid, falling back to value). Other categories stay cost-only. Checked-overflow behavior applies across all three categories. --- rust/src/providers/mistral/mod.rs | 152 +++++++++++++++++++++++++++++- 1 file changed, 147 insertions(+), 5 deletions(-) diff --git a/rust/src/providers/mistral/mod.rs b/rust/src/providers/mistral/mod.rs index 7ee469ee35..bcaa15e48c 100644 --- a/rust/src/providers/mistral/mod.rs +++ b/rust/src/providers/mistral/mod.rs @@ -31,6 +31,11 @@ const SUBSCRIPTION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs #[derive(Debug, Deserialize)] struct BillingResponse { completion: Option, + /// Le Chat consumption; counted like `completion`. + chat: Option, + /// Vibe Code consumption; its `completion` counts like API completions. + #[serde(rename = "vibe_code")] + vibe_code: Option, ocr: Option, connectors: Option, audio: Option, @@ -53,6 +58,11 @@ struct ModelUsageCategory { models: Option>, } +#[derive(Debug, Deserialize)] +struct VibeCodeCategory { + completion: Option, +} + #[derive(Debug, Deserialize)] struct LibrariesUsageCategory { pages: Option, @@ -283,7 +293,18 @@ impl MistralProvider { let mut total_tokens = TokenCounts::default(); let mut model_count = 0; - if let Some(models) = billing.completion.and_then(|c| c.models) { + // API, Le Chat, and Vibe Code completions share consumed-token and + // billed-cost accounting. + let token_categories = [ + billing.completion, + billing.chat, + billing.vibe_code.and_then(|vibe| vibe.completion), + ]; + for models in token_categories + .into_iter() + .flatten() + .filter_map(|category| category.models) + { model_count += models.len(); for data in models.values() { match Self::aggregate_model(data, &prices, AggregationMode::CostAndTokens)? { @@ -452,13 +473,15 @@ impl MistralProvider { (TokenKind::Cached, data.cached.as_deref()), ] { for entry in entries.unwrap_or_default() { - let units = entry.value_paid.or(entry.value).unwrap_or(0); + // Tokens count consumed units (plan-covered usage included); + // spend counts only billed units. if matches!(mode, AggregationMode::CostAndTokens) { - tokens.add_lane(units, kind)?; + tokens.add_lane(entry.value.or(entry.value_paid).unwrap_or(0), kind)?; } if let (Some(metric), Some(group)) = (&entry.billing_metric, &entry.billing_group) { - let entry_cost = - (units as f64) * prices.get(&format!("{metric}::{group}")).unwrap_or(&0.0); + let billed_units = entry.value_paid.or(entry.value).unwrap_or(0); + let entry_cost = (billed_units as f64) + * prices.get(&format!("{metric}::{group}")).unwrap_or(&0.0); Self::accumulate_finite_cost(entry_cost, &mut cost); } } @@ -711,4 +734,123 @@ mod tests { Err(ProviderError::Parse(message)) if message.contains("token count") )); } + + /// Model `fixture` priced at 0.25 per billed unit, wrapped for one category. + fn fixture_billing(category: &str, model: serde_json::Value) -> BillingResponse { + let models = serde_json::json!({ "models": { "fixture": model } }); + let mut payload = serde_json::json!({ + "prices": [{ "billing_metric": "fixture", "billing_group": "unit", "price": "0.25" }], + }); + payload[category] = if category == "vibe_code" { + serde_json::json!({ "completion": models }) + } else { + models + }; + serde_json::from_value(payload).unwrap() + } + + fn fixture_entry(value: Option, paid: Option) -> serde_json::Value { + let mut entry = serde_json::json!({ "billing_metric": "fixture", "billing_group": "unit" }); + if let Some(value) = value { + entry["value"] = value.into(); + } + if let Some(paid) = paid { + entry["value_paid"] = paid.into(); + } + entry + } + + #[test] + fn plan_covered_consumption_counts_tokens_separately_from_billed_spend() { + for category in ["completion", "chat", "vibe_code"] { + let billing = fixture_billing( + category, + serde_json::json!({ + "input": [fixture_entry(Some(1000), Some(0))], + "output": [fixture_entry(Some(500), Some(200))], + "cached": [fixture_entry(Some(300), Some(0))], + }), + ); + let summary = MistralProvider::summarize_billing(billing).unwrap(); + assert_eq!(summary.total_input_tokens, 1000, "{category}"); + assert_eq!(summary.total_output_tokens, 500, "{category}"); + assert_eq!(summary.total_cached_tokens, 300, "{category}"); + assert_eq!(summary.model_count, 1, "{category}"); + assert_eq!(summary.total_cost, 50.0, "{category}"); + } + } + + #[test] + fn missing_consumed_or_billed_units_use_the_available_count() { + for entry in [fixture_entry(Some(20), None), fixture_entry(None, Some(20))] { + let billing = fixture_billing("completion", serde_json::json!({ "input": [entry] })); + let summary = MistralProvider::summarize_billing(billing).unwrap(); + assert_eq!(summary.total_input_tokens, 20); + assert_eq!(summary.total_cost, 5.0); + } + } + + #[test] + fn categories_aggregate_together_and_count_each_model_entry_once() { + let billing: BillingResponse = serde_json::from_value(serde_json::json!({ + "prices": [{ "billing_metric": "fixture", "billing_group": "unit", "price": "0.25" }], + "completion": { "models": { "a": { "input": [fixture_entry(Some(10), Some(10))] } } }, + "chat": { "models": { "b": { "input": [fixture_entry(Some(20), Some(0))] } } }, + "vibe_code": { "completion": { "models": { + "c": { "output": [fixture_entry(Some(30), Some(4))] } + } } }, + })) + .unwrap(); + let summary = MistralProvider::summarize_billing(billing).unwrap(); + assert_eq!(summary.total_input_tokens, 30); + assert_eq!(summary.total_output_tokens, 30); + assert_eq!(summary.model_count, 3); + assert_eq!(summary.total_cost, 3.5); + } + + #[test] + fn cost_only_categories_do_not_sum_token_lanes() { + let entries = serde_json::json!({ + "input": [fixture_entry(Some(i64::MAX), None), fixture_entry(Some(1), None)] + }); + let billing: BillingResponse = serde_json::from_value(serde_json::json!({ + "prices": [{ "billing_metric": "fixture", "billing_group": "unit", "price": "0.25" }], + "ocr": { "models": { "fixture": entries } }, + })) + .unwrap(); + let summary = MistralProvider::summarize_billing(billing).unwrap(); + assert_eq!(summary.total_input_tokens, 0); + assert_eq!(summary.model_count, 0); + assert!(summary.total_cost > 0.0); + } + + #[test] + fn paid_zero_does_not_hide_an_unrepresentable_consumed_value() { + let entry = fixture_entry(Some(i64::MAX), Some(0)); + let billing = fixture_billing( + "completion", + serde_json::json!({ "input": [entry.clone(), entry] }), + ); + assert!(matches!( + MistralProvider::summarize_billing(billing), + Err(ProviderError::Parse(message)) if message.contains("token count") + )); + } + + #[test] + fn consumed_totals_stay_checked_across_api_chat_and_vibe_categories() { + let big = serde_json::json!({ "input": [fixture_entry(Some(i64::MAX), None)] }); + let one = serde_json::json!({ "input": [fixture_entry(Some(1), None)] }); + for mut payload in [ + serde_json::json!({ "chat": { "models": { "fixture": one } } }), + serde_json::json!({ "vibe_code": { "completion": { "models": { "fixture": one } } } }), + ] { + payload["completion"] = serde_json::json!({ "models": { "fixture": big } }); + let billing: BillingResponse = serde_json::from_value(payload).unwrap(); + assert!(matches!( + MistralProvider::summarize_billing(billing), + Err(ProviderError::Parse(message)) if message.contains("token count") + )); + } + } } From a9c501b99827ac044a144a958bb37969eb6f49a4 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:54:31 +0700 Subject: [PATCH 114/407] Port upstream 0.67.0: add CostReportingPeriod core and range-aware cost scanners --- rust/src/cost_reporting_period.rs | 456 ++++++++++++++++++ rust/src/cost_scanner.rs | 69 ++- rust/src/cost_scanner/codex.rs | 23 + rust/src/cost_scanner/codex/scan.rs | 14 +- rust/src/cost_scanner/tests.rs | 4 + rust/src/cost_scanner/tests/period.rs | 158 ++++++ rust/src/cost_scanner/window.rs | 61 +++ rust/src/lib.rs | 1 + rust/src/pi_session_cost.rs | 20 +- .../providers/antigravity/local_sessions.rs | 4 +- .../src/providers/antigravity/local_sqlite.rs | 4 +- rust/src/providers/bedrock/mod.rs | 79 ++- rust/src/providers/muse/local_usage/mod.rs | 2 +- rust/src/providers/opencodego/local.rs | 4 +- rust/src/settings.rs | 7 + rust/src/settings/raw.rs | 4 + rust/src/settings/tests.rs | 36 ++ 17 files changed, 900 insertions(+), 46 deletions(-) create mode 100644 rust/src/cost_reporting_period.rs create mode 100644 rust/src/cost_scanner/tests/period.rs create mode 100644 rust/src/cost_scanner/window.rs diff --git a/rust/src/cost_reporting_period.rs b/rust/src/cost_reporting_period.rs new file mode 100644 index 0000000000..c1721afaa3 --- /dev/null +++ b/rust/src/cost_reporting_period.rs @@ -0,0 +1,456 @@ +//! Calendar selection shared by local cost surfaces (upstream 0.67.0). +//! +//! A [`CostReportingPeriod`] is a *semantic* choice (rolling days, month to +//! date, or all available history). Every operation resolves it again into +//! inclusive local-day bounds at its own boundary, so a saved selection never +//! goes stale across midnight or a month rollover. All arithmetic runs on +//! calendar dates, never on 24-hour multiples, so leap years and 23/25-hour +//! daylight-saving days need no special cases. + +use std::fmt; + +use chrono::{DateTime, Datelike, Duration, Local, NaiveDate, TimeZone, Utc}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +/// Largest rolling window a user can select. +pub const MAX_ROLLING_DAYS: u32 = 365; +/// Rolling window used when nothing is saved. +pub const DEFAULT_ROLLING_DAYS: u32 = 30; +/// Sanity ceiling for a resolved day count handed to a scanner (100 years). +/// +/// Rolling selections stop at [`MAX_ROLLING_DAYS`], but *all available* history +/// resolves to many more days. The ceiling only keeps date arithmetic far from +/// `NaiveDate` overflow when a caller passes an arbitrary `u32`. +pub const MAX_WINDOW_DAYS: u32 = 36_500; + +/// First day considered by an all-available window when no earliest source day +/// is known. Nothing this app reads predates it. +const ALL_AVAILABLE_FLOOR: NaiveDate = match NaiveDate::from_ymd_opt(1970, 1, 1) { + Some(date) => date, + None => panic!("valid epoch date"), +}; + +/// Clamp a resolved day count into `1..=MAX_WINDOW_DAYS`. +pub fn clamp_window_days(days: u32) -> u32 { + days.clamp(1, MAX_WINDOW_DAYS) +} + +/// Zone whose midnights bound a reporting day. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CostTimeZone { + /// The machine's local zone (what the local scanners bucket by today). + Local, + /// A pinned IANA zone. + Named(chrono_tz::Tz), +} + +impl CostTimeZone { + /// Zone identifier used in cache identities. + pub fn identifier(&self) -> String { + match self { + Self::Local => crate::core::local_timezone_name(), + Self::Named(tz) => tz.name().to_string(), + } + } + + /// The calendar date `now` falls on in this zone. + pub fn date(&self, now: DateTime) -> NaiveDate { + match self { + Self::Local => now.with_timezone(&Local).date_naive(), + Self::Named(tz) => now.with_timezone(tz).date_naive(), + } + } + + /// The instant a calendar date begins in this zone. + /// + /// A midnight skipped by a DST gap begins at the first valid instant of + /// that date instead. + pub fn start_of_day_utc(&self, date: NaiveDate) -> DateTime { + match self { + Self::Local => Self::resolve(&Local, date), + Self::Named(tz) => Self::resolve(tz, date), + } + } + + fn resolve(tz: &Tz, date: NaiveDate) -> DateTime { + let midnight = date.and_time(chrono::NaiveTime::MIN); + // A gap at midnight (a few zones spring forward at 00:00) has no local + // midnight; the first valid instant is at most a few hours later. + (0..=6) + .find_map(|hours| { + tz.from_local_datetime(&(midnight + Duration::hours(hours))) + .earliest() + }) + .map_or_else( + || Utc.from_utc_datetime(&midnight), + |local| local.with_timezone(&Utc), + ) + } +} + +/// Inclusive local-day bounds resolved from a [`CostReportingPeriod`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CostDayBounds { + pub start: NaiveDate, + pub end: NaiveDate, +} + +impl CostDayBounds { + /// Number of calendar days covered, counting both ends. + pub fn days(&self) -> u32 { + u32::try_from((self.end - self.start).num_days() + 1).unwrap_or(u32::MAX) + } + + /// Whether a `YYYY-MM-DD` day key falls inside the bounds. + pub fn contains_day_key(&self, key: &str) -> bool { + NaiveDate::parse_from_str(key, "%Y-%m-%d") + .is_ok_and(|day| day >= self.start && day <= self.end) + } +} + +/// Which slice of local history a cost surface reports. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum CostReportingPeriod { + /// The last N calendar days including today. + /// + /// Parsing and [`Self::rolling`] keep N within `1..=365`; the scanner + /// entry point [`crate::cost_scanner::CostScanner::new`] wraps whatever + /// day count its existing callers pass. + Rolling(u32), + /// From midnight on the first of the current month through today. + MonthToDate, + /// Every day with available source history. Not clamped to a year. + AllAvailable, +} + +impl Default for CostReportingPeriod { + fn default() -> Self { + Self::Rolling(DEFAULT_ROLLING_DAYS) + } +} + +impl CostReportingPeriod { + /// A rolling selection clamped to `1..=365`. + pub fn rolling(days: u32) -> Self { + Self::Rolling(days.clamp(1, MAX_ROLLING_DAYS)) + } + + /// Parse the persisted form: `rolling:N`, `month-to-date`, or `all`. + /// + /// `rolling:0` and non-numeric counts are rejected; counts above 365 clamp. + pub fn parse(raw: &str) -> Option { + match raw.trim() { + "month-to-date" => Some(Self::MonthToDate), + "all" => Some(Self::AllAvailable), + other => { + let days = other.strip_prefix("rolling:")?.parse::().ok()?; + (days > 0).then(|| Self::rolling(days)) + } + } + } + + /// The persisted form, the inverse of [`Self::parse`]. + pub fn raw(&self) -> String { + match self { + Self::Rolling(days) => format!("rolling:{days}"), + Self::MonthToDate => "month-to-date".to_string(), + Self::AllAvailable => "all".to_string(), + } + } + + /// Resolve a saved selection, falling back to a legacy day count and then + /// to the 30-day default. An unreadable saved value migrates like a + /// missing one. + pub fn migrated(raw: Option<&str>, legacy_days: Option) -> Self { + raw.and_then(Self::parse) + .unwrap_or_else(|| Self::rolling(legacy_days.unwrap_or(DEFAULT_ROLLING_DAYS))) + } + + /// Inclusive day bounds for `now` in `tz`. + /// + /// `earliest` only affects [`Self::AllAvailable`]: it is the first day + /// with source data. Without it the window opens at 1970-01-01, which + /// callers that walk day by day must avoid by supplying the real first + /// day. The start never passes today. + pub fn bounds( + &self, + now: DateTime, + tz: CostTimeZone, + earliest: Option, + ) -> CostDayBounds { + let end = tz.date(now); + let start = match self { + Self::Rolling(days) => end - Duration::days(i64::from((*days).max(1) - 1)), + Self::MonthToDate => end.with_day(1).unwrap_or(end), + Self::AllAvailable => earliest.unwrap_or(ALL_AVAILABLE_FLOOR), + }; + CostDayBounds { + start: start.min(end), + end, + } + } + + /// Number of days in the resolved window. + pub fn days(&self, now: DateTime, tz: CostTimeZone, earliest: Option) -> u32 { + self.bounds(now, tz, earliest).days() + } + + /// Keep the entries whose day key lies inside the resolved window. + pub fn entries( + &self, + entries: Vec, + day_key: impl Fn(&T) -> &str, + now: DateTime, + tz: CostTimeZone, + ) -> Vec { + let bounds = self.bounds(now, tz, None); + entries + .into_iter() + .filter(|entry| bounds.contains_day_key(day_key(entry))) + .collect() + } + + /// Cache identity `raw|zone|start|end`. + /// + /// It changes with the selection, the pinned zone, and the resolved dates, + /// so a month rollover or a zone change never reuses a previous window. + pub fn identity(&self, now: DateTime, tz: CostTimeZone) -> String { + let bounds = self.bounds(now, tz, None); + format!( + "{}|{}|{}|{}", + self.raw(), + tz.identifier(), + bounds.start.format("%Y-%m-%d"), + bounds.end.format("%Y-%m-%d"), + ) + } +} + +impl fmt::Display for CostReportingPeriod { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.raw()) + } +} + +impl Serialize for CostReportingPeriod { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(&self.raw()) + } +} + +/// 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)?; + Ok(Self::migrated(raw.as_deref(), None)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono_tz::{America::Los_Angeles, Asia::Tokyo}; + + const LA: CostTimeZone = CostTimeZone::Named(Los_Angeles); + const UTC: CostTimeZone = CostTimeZone::Named(chrono_tz::UTC); + + fn at(value: &str) -> DateTime { + DateTime::parse_from_rfc3339(value) + .expect("fixture instant") + .with_timezone(&Utc) + } + + fn day(value: &str) -> NaiveDate { + NaiveDate::parse_from_str(value, "%Y-%m-%d").expect("fixture day") + } + + #[test] + fn month_starts_and_day_counts_follow_the_zone_calendar() { + // (now, days into the month, month start): copied in spirit from + // upstream CostReportingPeriodTests (Los Angeles bucket calendar). + let cases = [ + ("2026-01-31T20:00:00Z", 31, "2026-01-01"), + ("2026-02-01T08:00:00Z", 1, "2026-02-01"), + ("2026-02-28T20:00:00Z", 28, "2026-02-01"), + ("2024-02-29T20:00:00Z", 29, "2024-02-01"), + // 23-hour spring-forward day (2026-03-08) earlier in the month. + ("2026-03-09T07:00:00Z", 9, "2026-03-01"), + // 25-hour fall-back day (2026-11-01) earlier in the month. + ("2026-11-02T08:00:00Z", 2, "2026-11-01"), + ]; + for (now, days, start) in cases { + let now = at(now); + let period = CostReportingPeriod::MonthToDate; + assert_eq!(period.days(now, LA, None), days, "{now}"); + assert_eq!(period.bounds(now, LA, None).start, day(start), "{now}"); + assert_eq!(period.bounds(now, LA, None).end, LA.date(now), "{now}"); + } + } + + #[test] + fn identity_distinguishes_periods_rollover_and_zones() { + let january = at("2026-02-01T07:59:59Z"); + let february = january + Duration::seconds(1); + let month = CostReportingPeriod::MonthToDate; + assert_ne!(month.identity(january, LA), month.identity(february, LA)); + assert_ne!( + month.identity(february, LA), + CostReportingPeriod::Rolling(1).identity(february, LA) + ); + assert_eq!(month.days(january, LA, None), 31); + assert_eq!(month.days(january, UTC, None), 1); + assert_ne!(month.identity(january, LA), month.identity(january, UTC)); + assert_eq!( + month.identity(january, LA), + "month-to-date|America/Los_Angeles|2026-01-01|2026-01-31" + ); + } + + #[test] + fn rolling_bounds_include_today_and_never_use_24_hour_math() { + let now = at("2026-03-09T07:00:00Z"); + // Seven days back from Mar 9 crosses the 2026-03-08 spring-forward day. + let bounds = CostReportingPeriod::Rolling(7).bounds(now, LA, None); + assert_eq!( + (bounds.start, bounds.end), + (day("2026-03-03"), day("2026-03-09")) + ); + assert_eq!(bounds.days(), 7); + let today = CostReportingPeriod::Rolling(1).bounds(now, LA, None); + assert_eq!( + (today.start, today.end), + (day("2026-03-09"), day("2026-03-09")) + ); + // A degenerate zero-day rolling window still covers today. + assert_eq!(CostReportingPeriod::Rolling(0).days(now, LA, None), 1); + } + + #[test] + fn month_to_date_on_the_first_covers_only_today() { + let now = at("2026-02-01T08:00:00Z"); + let bounds = CostReportingPeriod::MonthToDate.bounds(now, LA, None); + assert_eq!(bounds.start, bounds.end); + assert_eq!(bounds.start, day("2026-02-01")); + } + + #[test] + fn month_start_follows_the_pinned_zone_across_the_date_line() { + let now = at("2026-02-28T16:00:00Z"); + // Tokyo is already on March 1; Los Angeles is still on February 28. + assert_eq!( + CostReportingPeriod::MonthToDate + .bounds(now, CostTimeZone::Named(Tokyo), None) + .start, + day("2026-03-01") + ); + assert_eq!( + CostReportingPeriod::MonthToDate.bounds(now, LA, None).start, + day("2026-02-01") + ); + } + + #[test] + fn all_available_uses_the_earliest_day_and_is_not_clamped_to_a_year() { + let now = at("2026-07-15T12:00:00Z"); + let earliest = day("2023-05-04"); + let bounds = CostReportingPeriod::AllAvailable.bounds(now, UTC, Some(earliest)); + assert_eq!((bounds.start, bounds.end), (earliest, day("2026-07-15"))); + assert!(bounds.days() > MAX_ROLLING_DAYS); + // Without a known earliest day the window opens at the epoch floor. + assert_eq!( + CostReportingPeriod::AllAvailable + .bounds(now, UTC, None) + .start, + day("1970-01-01") + ); + // A future earliest day cannot push the start past today. + let future = CostReportingPeriod::AllAvailable.bounds(now, UTC, Some(day("2030-01-01"))); + assert_eq!(future.start, future.end); + } + + #[test] + fn raw_forms_round_trip_and_reject_invalid_input() { + for (raw, period) in [ + ("rolling:30", CostReportingPeriod::Rolling(30)), + ("rolling:1", CostReportingPeriod::Rolling(1)), + ("month-to-date", CostReportingPeriod::MonthToDate), + ("all", CostReportingPeriod::AllAvailable), + ] { + assert_eq!(CostReportingPeriod::parse(raw), Some(period)); + assert_eq!(period.raw(), raw); + } + assert_eq!( + CostReportingPeriod::parse("rolling:999"), + Some(CostReportingPeriod::Rolling(365)) + ); + for invalid in [ + "", + "rolling:0", + "rolling:", + "rolling:-3", + "rolling:x", + "week", + "ALL", + ] { + assert_eq!(CostReportingPeriod::parse(invalid), None, "{invalid:?}"); + } + } + + #[test] + fn legacy_windows_migrate_unchanged() { + let migrated = CostReportingPeriod::migrated; + assert_eq!(migrated(None, Some(90)), CostReportingPeriod::Rolling(90)); + assert_eq!(migrated(None, None), CostReportingPeriod::Rolling(30)); + assert_eq!( + migrated(Some("invalid"), Some(999)), + CostReportingPeriod::Rolling(365) + ); + assert_eq!( + migrated(Some("month-to-date"), Some(90)), + CostReportingPeriod::MonthToDate + ); + } + + #[test] + fn entries_filter_by_the_resolved_window() { + let now = at("2026-02-02T20:00:00Z"); + let rows = vec![ + ("2026-01-31", 100), + ("2026-02-01", 3), + ("2026-02-02", 4), + ("2026-02-03", 9), + ("not-a-day", 1), + ]; + let kept = CostReportingPeriod::MonthToDate.entries(rows, |row| row.0, now, LA); + assert_eq!(kept.iter().map(|row| row.1).sum::(), 7); + } + + #[test] + fn start_of_day_is_the_local_midnight_instant() { + assert_eq!( + LA.start_of_day_utc(day("2026-02-01")), + at("2026-02-01T08:00:00Z") + ); + // Daylight time: the first of November 2026 is still PDT (UTC-7). + assert_eq!( + LA.start_of_day_utc(day("2026-11-01")), + at("2026-11-01T07:00:00Z") + ); + assert_eq!( + UTC.start_of_day_utc(day("2026-02-01")), + at("2026-02-01T00:00:00Z") + ); + } + + #[test] + fn serde_uses_the_raw_string_and_tolerates_unreadable_values() { + assert_eq!( + serde_json::to_string(&CostReportingPeriod::MonthToDate).unwrap(), + "\"month-to-date\"" + ); + let read = |json: &str| serde_json::from_str::(json).unwrap(); + assert_eq!(read("\"all\""), CostReportingPeriod::AllAvailable); + assert_eq!(read("\"nonsense\""), CostReportingPeriod::Rolling(30)); + assert_eq!(read("null"), CostReportingPeriod::Rolling(30)); + } +} diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 619f9ba7b0..ce1c98ebe2 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -21,14 +21,15 @@ use std::time::{SystemTime, UNIX_EPOCH}; #[cfg(test)] use crate::codex_costs::scan_codex_file_cost; use crate::codex_costs::{ - add_codex_days_map_to_summary, add_codex_records_to_summary, codex_period_start, - codex_scan_dates, merge_codex_records_into_days, + add_codex_days_map_to_summary, add_codex_records_to_summary, codex_scan_dates, + merge_codex_records_into_days, }; use crate::codex_sessions::{codex_sessions_dir_candidates, default_wsl_roots}; use crate::core::{ CachedCostReport, CodexScanPauseReason, CostScanOptions, CostUsageCache, CostUsageDayRange, CostUsageFileUsage, JsonlScanner, ProviderId, }; +use crate::cost_reporting_period::{CostReportingPeriod, MAX_ROLLING_DAYS}; use crate::providers::claude::quota_history::{ ClaudeHistoryAttribution, ClaudeQuotaDedupKey, ClaudeQuotaHistoryRecord, }; @@ -39,6 +40,7 @@ mod claude_usage; mod codex; mod read_receipt; mod stats; +mod window; use claude_pricing::ClaudeScanPricingResolver; #[cfg(test)] use claude_pricing::{ClaudePricing, FALLBACK_CLAUDE_MODEL}; @@ -487,7 +489,8 @@ impl ClaudeFileScanResult { #[derive(Debug, Clone)] pub struct CostScanner { - days: u32, + /// The window every scan resolves at its own start (never cached as dates). + period: CostReportingPeriod, options: CostScanOptions, cache_root: Option, /// When set, bypass normal sessions-dir discovery (tests / inject roots). @@ -497,14 +500,25 @@ pub struct CostScanner { impl CostScanner { /// Create a new scanner for the last N days (default 60s cache debounce). pub fn new(days: u32) -> Self { + Self::for_period(CostReportingPeriod::Rolling(days)) + } + + /// Create a scanner for a reporting period. `new(days)` is + /// `for_period(CostReportingPeriod::Rolling(days))`. + pub fn for_period(period: CostReportingPeriod) -> Self { Self { - days, + period, options: CostScanOptions::default(), cache_root: None, sessions_dirs_override: None, } } + /// The reporting period this scanner resolves on every scan. + pub fn period(&self) -> CostReportingPeriod { + self.period + } + /// Override scan options (e.g. [`CostScanOptions::app_driven`] for force refresh). pub fn with_options(mut self, options: CostScanOptions) -> Self { self.options = options; @@ -533,16 +547,17 @@ impl CostScanner { } pub fn scan_pi_with_cancel(&self, cancel: Option<&AtomicBool>) -> CostSummary { - let today = Utc::now().date_naive(); + let now = Utc::now(); + let window = self.transcript_window(now, now.date_naive()); let mut summary = CostSummary { - period_start: Some(today - Duration::days(self.days as i64)), - period_end: Some(today), + period_start: Some(window.start), + period_end: Some(window.end), ..CostSummary::default() }; let mut seen_entries = HashSet::new(); let evidence = crate::pi_session_cost::scan_pi_into( &mut summary, - self.days, + window.cutoff, cancel, &mut seen_entries, ); @@ -577,12 +592,12 @@ impl CostScanner { ) -> CostSummary { let projects_dir = self.get_claude_projects_dir(); let mut summary = CostSummary::default(); - let today = Utc::now().date_naive(); - let start_date = today - Duration::days(self.days as i64); - let cutoff = Utc::now() - Duration::days(self.days as i64); + let now = Utc::now(); + let window = self.transcript_window(now, now.date_naive()); + let cutoff = window.cutoff; - summary.period_start = Some(start_date); - summary.period_end = Some(today); + summary.period_start = Some(window.start); + summary.period_end = Some(window.end); // Walk through projects directory, de-duplicating usage records // that appear across multiple files. @@ -615,7 +630,7 @@ impl CostScanner { crate::pi_session_cost::scan_pi_compatible_into( &mut summary, crate::pi_session_cost::PiMappedProvider::Claude, - self.days, + cutoff, cancel, &mut seen_pi, ); @@ -643,17 +658,26 @@ impl CostScanner { cancel: Option<&AtomicBool>, ) -> ClaudeChartSnapshot { let projects_dir = self.get_claude_projects_dir(); - let today = Local::now().date_naive(); - let cutoff = Utc::now() - Duration::days(self.days as i64); + let now = Utc::now(); + let window = self.transcript_window(now, Local::now().date_naive()); + let cutoff = window.cutoff; let mut summary = CostSummary { - period_start: Some(today - Duration::days(self.days as i64)), - period_end: Some(today), + period_start: Some(window.start), + period_end: Some(window.end), ..CostSummary::default() }; let mut daily_cost = HashMap::new(); let mut daily_tokens = HashMap::new(); - for days_ago in 0..self.days { - let date = today - Duration::days(days_ago as i64); + // The chart never shows more than a year of daily slots, so an + // all-available window does not allocate one per day since 1970. + let slot_days = match self.period { + CostReportingPeriod::Rolling(days) => days, + CostReportingPeriod::MonthToDate | CostReportingPeriod::AllAvailable => { + window.days.min(MAX_ROLLING_DAYS) + } + }; + for days_ago in 0..slot_days { + let date = window.end - Duration::days(days_ago as i64); let key = date.format("%Y-%m-%d").to_string(); daily_cost.insert(key.clone(), None); daily_tokens.insert(key, 0); @@ -695,7 +719,7 @@ impl CostScanner { crate::pi_session_cost::scan_pi_compatible_into( &mut summary, crate::pi_session_cost::PiMappedProvider::Claude, - self.days, + cutoff, cancel, &mut HashSet::new(), ); @@ -763,7 +787,8 @@ impl CostScanner { return CostSummary::default(); } let now = Utc::now(); - let Some(local) = opencodego_local::model_cost_summary_scan(now, self.days) else { + let days = self.calendar_window(now, None).days; + let Some(local) = opencodego_local::model_cost_summary_scan(now, days) else { return CostSummary::default(); }; CostSummary { diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index f054c6eb8f..4491c87bf3 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -102,6 +102,29 @@ fn codex_parent_baseline( baseline } +/// First day that can hold a Codex date partition: January 1 of the earliest +/// `YYYY` directory under any sessions root (upstream `firstPartitionDate`). +/// +/// An all-available window starts here so the `YYYY/MM/DD` walk never probes +/// empty history. Years before 1970 are stray directories, not sessions. +pub(super) fn first_codex_partition_date(sessions_dirs: &[PathBuf]) -> Option { + sessions_dirs + .iter() + .filter_map(|root| fs::read_dir(root).ok()) + .flat_map(|entries| entries.flatten()) + .filter(|entry| entry.file_type().is_ok_and(|kind| kind.is_dir())) + .filter_map(|entry| { + let name = entry.file_name(); + let name = name.to_str()?; + (name.len() == 4 && name.bytes().all(|byte| byte.is_ascii_digit())) + .then(|| name.parse::().ok()) + .flatten() + }) + .filter(|year| *year >= 1970) + .min() + .and_then(|year| NaiveDate::from_ymd_opt(year, 1, 1)) +} + fn is_codex_path_in_scan_window( path: &Path, sessions_dirs: &[PathBuf], diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index e983adfd6a..a1a87533f0 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -89,8 +89,15 @@ pub(super) fn scan_codex_detailed_with_cache( ) -> (CostSummary, CostScanStats, CostUsageCache) { let mut summary = CostSummary::default(); let mut stats = CostScanStats::default(); + let sessions_dirs = scanner.get_codex_sessions_dirs(); + let now = Utc::now(); let today = Local::now().date_naive(); - let start_date = codex_period_start(today, scanner.days); + // All-available history opens at the first existing partition; with no + // partitions it is just today, so an empty tree costs no directory probes. + let earliest = (scanner.period == CostReportingPeriod::AllAvailable) + .then(|| first_codex_partition_date(&sessions_dirs).unwrap_or(today)); + let window = scanner.calendar_window(now, earliest); + let (start_date, today) = (window.start, window.end); let range = CostUsageDayRange::new(start_date, today); let now_ms = unix_now_ms(); @@ -99,7 +106,6 @@ pub(super) fn scan_codex_detailed_with_cache( let cache_root = scanner.cache_root.as_deref(); let mut cache = JsonlScanner::load_cache(ProviderId::Codex, cache_root); - let sessions_dirs = scanner.get_codex_sessions_dirs(); let pending_scan = CodexPendingScanContext::new( &cache, &range, @@ -452,11 +458,13 @@ fn append_pi_compatible_costs( return; } + let now = Utc::now(); + let cutoff = scanner.transcript_window(now, now.date_naive()).cutoff; let mut seen_pi = HashSet::new(); crate::pi_session_cost::scan_pi_compatible_into( summary, crate::pi_session_cost::PiMappedProvider::Codex, - scanner.days, + cutoff, cancel, &mut seen_pi, ); diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 3b16c92a1c..19acf2a9d0 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::codex_costs::codex_period_start; use crate::core::{CodexSessionLineage, CostUsagePricing}; use std::io::Write; @@ -3039,3 +3040,6 @@ fn incomplete_or_buffered_empty_codex_fragment_is_not_marked_complete() { #[cfg(test)] #[path = "tests/paginated.rs"] mod paginated; + +#[path = "tests/period.rs"] +mod period; diff --git a/rust/src/cost_scanner/tests/period.rs b/rust/src/cost_scanner/tests/period.rs new file mode 100644 index 0000000000..68d84c72a0 --- /dev/null +++ b/rust/src/cost_scanner/tests/period.rs @@ -0,0 +1,158 @@ +//! Reporting-period behavior of the scanner: window resolution and the Codex +//! partition walk for month-to-date and all-available history. + +use super::*; +use crate::cost_reporting_period::CostReportingPeriod; +use chrono::{Datelike, TimeZone}; + +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() +} + +/// Writes one Codex session under the `YYYY/MM/DD` partition for `day`. +fn write_session_on(sessions_root: &Path, day: NaiveDate, input_tokens: u64) { + 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(); + // Local noon keeps the record on `day` in the local zone. + let local_noon = Local + .from_local_datetime(&day.and_hms_opt(12, 0, 0).unwrap()) + .earliest() + .unwrap() + .with_timezone(&Utc); + let ts = local_noon.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); + let body = format!( + r#"{{"timestamp":"{ts}","type":"event_msg","payload":{{"type":"token_count","info":{{"model":"gpt-5","total_token_usage":{{"input_tokens":{input_tokens},"cached_input_tokens":0,"output_tokens":5}}}}}}}} +"# + ); + std::fs::write(day_dir.join(format!("s-{input_tokens}.jsonl")), body).unwrap(); +} + +fn scan_input_tokens(period: CostReportingPeriod, sessions: &Path, cache: &Path) -> u64 { + let scanner = CostScanner::for_period(period) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(cache) + .with_sessions_dirs(vec![sessions.to_path_buf()]); + let (summary, _) = scanner.scan_codex_detailed(None); + summary.input_tokens +} + +#[test] +fn for_period_rolling_matches_new() { + assert_eq!( + CostScanner::new(14).period(), + CostScanner::for_period(CostReportingPeriod::Rolling(14)).period() + ); +} + +#[test] +fn zero_day_scanner_still_resolves_a_one_day_window() { + let now = Utc::now(); + let window = CostScanner::new(0).calendar_window(now, None); + assert_eq!(window.days, 1); + assert_eq!(window.start, window.end); +} + +#[test] +fn transcript_window_rolling_keeps_the_legacy_shape() { + let now = utc(2026, 5, 15, 12); + let today = date(2026, 5, 15); + let window = CostScanner::new(30).transcript_window(now, today); + assert_eq!(window.start, date(2026, 4, 15)); + assert_eq!(window.end, today); + assert_eq!(window.cutoff, now - Duration::days(30)); + assert_eq!(window.days, 30); +} + +#[test] +fn transcript_window_month_to_date_starts_at_the_month() { + let now = Utc::now(); + let scanner = CostScanner::for_period(CostReportingPeriod::MonthToDate); + let window = scanner.transcript_window(now, now.date_naive()); + let local_today = Local::now().date_naive(); + assert_eq!(window.start.day(), 1); + assert_eq!(window.start.month(), local_today.month()); + assert!(window.cutoff <= now); + assert!(window.days >= 1 && window.days <= 31); +} + +#[test] +fn calendar_window_all_uses_the_earliest_day_without_clamping() { + let now = Utc::now(); + let earliest = date(2000, 1, 1); + let scanner = CostScanner::for_period(CostReportingPeriod::AllAvailable); + let window = scanner.calendar_window(now, Some(earliest)); + assert_eq!(window.start, earliest); + assert!(window.days > 365 * 20); +} + +#[test] +fn first_partition_skips_years_before_the_first_existing_one() { + let root = tempfile::tempdir().unwrap(); + for name in ["2024", "2023", "0500", "abcd", "20240"] { + std::fs::create_dir_all(root.path().join(name)).unwrap(); + } + std::fs::write(root.path().join("1999"), b"file, not a partition").unwrap(); + assert_eq!( + codex::first_codex_partition_date(&[root.path().to_path_buf()]), + Some(date(2023, 1, 1)) + ); +} + +#[test] +fn first_partition_is_none_for_empty_or_missing_roots() { + let root = tempfile::tempdir().unwrap(); + assert_eq!( + codex::first_codex_partition_date(&[root.path().to_path_buf()]), + None + ); + assert_eq!( + codex::first_codex_partition_date(&[root.path().join("missing")]), + None + ); +} + +#[test] +fn codex_all_available_reads_history_beyond_a_year() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let today = Local::now().date_naive(); + write_session_on(&sessions, today, 100); + write_session_on(&sessions, today - Duration::days(500), 1_000); + + let rolling = scan_input_tokens( + CostReportingPeriod::Rolling(30), + &sessions, + &root.path().join("cache-rolling"), + ); + let all = scan_input_tokens( + CostReportingPeriod::AllAvailable, + &sessions, + &root.path().join("cache-all"), + ); + assert_eq!(rolling, 100, "rolling 30 excludes the 500 day old session"); + assert_eq!(all, 1_100, "all available includes the old partition"); +} + +#[test] +fn codex_month_to_date_excludes_the_previous_month() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let today = Local::now().date_naive(); + let month_start = today.with_day(1).unwrap(); + write_session_on(&sessions, today, 100); + write_session_on(&sessions, month_start - Duration::days(3), 1_000); + + let mtd = scan_input_tokens( + CostReportingPeriod::MonthToDate, + &sessions, + &root.path().join("cache-mtd"), + ); + assert_eq!(mtd, 100); +} diff --git a/rust/src/cost_scanner/window.rs b/rust/src/cost_scanner/window.rs new file mode 100644 index 0000000000..e108b7e063 --- /dev/null +++ b/rust/src/cost_scanner/window.rs @@ -0,0 +1,61 @@ +//! Resolves a scanner's [`CostReportingPeriod`] into concrete scan windows. + +use chrono::{DateTime, Duration, NaiveDate, Utc}; + +use super::CostScanner; +use crate::cost_reporting_period::{CostReportingPeriod, CostTimeZone, clamp_window_days}; + +/// Inclusive local-day window plus the instant transcript scanners cut off at. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct ScanWindow { + /// First day of the window (also reported as `period_start`). + pub start: NaiveDate, + /// Last day of the window (today). + pub end: NaiveDate, + /// Records before this instant are outside the window. + pub cutoff: DateTime, + /// Resolved day count for scanners that take a day count. + pub days: u32, +} + +impl CostScanner { + /// The window in the local zone, resolved from the calendar. + /// + /// `earliest` is the first day with source data; it only matters for + /// [`CostReportingPeriod::AllAvailable`]. Rolling(N) covers N days + /// including today, and the cutoff is the start of the first day. + pub(super) fn calendar_window( + &self, + now: DateTime, + earliest: Option, + ) -> ScanWindow { + let tz = CostTimeZone::Local; + let bounds = self.period.bounds(now, tz, earliest); + ScanWindow { + start: bounds.start, + end: bounds.end, + cutoff: tz.start_of_day_utc(bounds.start), + days: clamp_window_days(bounds.days()), + } + } + + /// The window used by the Claude and Pi transcript scanners. + /// + /// A rolling window keeps its historical shape (a `now - N * 24h` cutoff + /// and `today - N` as the reported start, with `today` supplied by the + /// caller); month to date and all available history use exact calendar + /// midnights so no record from the previous month leaks in. + pub(super) fn transcript_window(&self, now: DateTime, today: NaiveDate) -> ScanWindow { + match self.period { + CostReportingPeriod::Rolling(days) => ScanWindow { + start: today - Duration::days(i64::from(days)), + end: today, + cutoff: now - Duration::days(i64::from(days)), + days, + }, + CostReportingPeriod::MonthToDate | CostReportingPeriod::AllAvailable => { + self.calendar_window(now, None) + } + } + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 1dfa17cff5..17e12e3554 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -11,6 +11,7 @@ pub mod codex_accounts; pub mod codex_cli; pub mod codex_workspaces; pub mod core; +pub mod cost_reporting_period; pub mod cost_scanner; pub mod host; pub mod locale; diff --git a/rust/src/pi_session_cost.rs b/rust/src/pi_session_cost.rs index a9f8f1b82c..4ff684f8d5 100644 --- a/rust/src/pi_session_cost.rs +++ b/rust/src/pi_session_cost.rs @@ -113,13 +113,13 @@ fn resolve_environment_path(value: &str, cwd: &Path) -> Option { pub fn scan_pi_compatible_into( summary: &mut CostSummary, target: PiMappedProvider, - days: u32, + cutoff: DateTime, cancel: Option<&AtomicBool>, seen_entries: &mut HashSet, ) -> PiScanEvidence { scan_roots_into( summary, - days, + cutoff, cancel, seen_entries, pi_compatible_session_roots(dirs::home_dir()), @@ -135,13 +135,13 @@ pub fn scan_pi_compatible_into( /// both roots so standalone Pi history is not double-counted. pub fn scan_pi_into( summary: &mut CostSummary, - days: u32, + cutoff: DateTime, cancel: Option<&AtomicBool>, seen_entries: &mut HashSet, ) -> PiScanEvidence { scan_roots_into( summary, - days, + cutoff, cancel, seen_entries, pi_compatible_session_roots(dirs::home_dir()), @@ -209,13 +209,12 @@ fn scan_pi_daily_from_roots( fn scan_roots_into( summary: &mut CostSummary, - days: u32, + cutoff: DateTime, 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; let mut evidence = PiScanEvidence::default(); for root in roots { @@ -895,7 +894,14 @@ mod tests { let mut summary = CostSummary::default(); let mut seen = HashSet::new(); - let evidence = scan_roots_into(&mut summary, 365, None, &mut seen, vec![sessions], None); + let evidence = scan_roots_into( + &mut summary, + Utc::now() - Duration::days(365), + None, + &mut seen, + vec![sessions], + None, + ); assert!(!evidence.complete); assert_eq!(summary.input_tokens, 11); } diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 458bb97f69..5e785b8c0c 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -141,7 +141,9 @@ fn summarize_paths( truncated: bool, ) -> LocalSessionSummary { let first_day = now.with_timezone(&Local).date_naive() - - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + - Duration::days(i64::from( + crate::cost_reporting_period::clamp_window_days(days).saturating_sub(1), + )); let mut total_tokens = 0_u64; let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index b050660411..fb229feeca 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -139,7 +139,9 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let first_day = now.with_timezone(&Local).date_naive() - - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + - Duration::days(i64::from( + crate::cost_reporting_period::clamp_window_days(days).saturating_sub(1), + )); let mut complete = discovery_complete && budget.check(); let mut events = Vec::new(); let mut authoritative_database = false; diff --git a/rust/src/providers/bedrock/mod.rs b/rust/src/providers/bedrock/mod.rs index ad5fdef00b..4068eab03b 100644 --- a/rust/src/providers/bedrock/mod.rs +++ b/rust/src/providers/bedrock/mod.rs @@ -3,7 +3,7 @@ //! Fetches current-month Bedrock spend from AWS Cost Explorer using SigV4. use async_trait::async_trait; -use chrono::{Datelike, Duration, TimeZone, Utc}; +use chrono::{DateTime, Datelike, Duration, Months, NaiveDate, TimeZone, Utc}; use reqwest::Client; use serde_json::{Value, json}; @@ -11,6 +11,7 @@ use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, hex, hmac_sha256, sha256_hex, }; +use crate::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; const COST_EXPLORER_URL: &str = "https://ce.us-east-1.amazonaws.com"; const COST_EXPLORER_TARGET: &str = "AWSInsightsIndexService.GetCostAndUsage"; @@ -656,19 +657,36 @@ fn parse_aws_profile_credentials(stdout: &[u8]) -> Result (String, String) { - let now = Utc::now(); - let start = Utc - .with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0) - .single() - .unwrap_or(now); - let tomorrow = (now + Duration::days(1)).date_naive(); +/// 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. +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); ( - start.format("%Y-%m-%d").to_string(), + 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 { @@ -805,6 +823,49 @@ 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!({ diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 49a2f7ec34..9ae2ce6d30 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -155,7 +155,7 @@ fn cache_root() -> PathBuf { } fn day_window(days: u32) -> (String, String) { - let days = days.clamp(1, 365); + let days = crate::cost_reporting_period::clamp_window_days(days); let today = Local::now().date_naive(); ( (today - chrono::Duration::days(i64::from(days - 1))).to_string(), diff --git a/rust/src/providers/opencodego/local.rs b/rust/src/providers/opencodego/local.rs index 5777706864..0af1b3b2ad 100644 --- a/rust/src/providers/opencodego/local.rs +++ b/rust/src/providers/opencodego/local.rs @@ -378,7 +378,7 @@ pub fn daily_model_costs( now: DateTime, history_days: u32, ) -> Vec { - let clamped = history_days.clamp(1, 365); + let clamped = crate::cost_reporting_period::clamp_window_days(history_days); let today = local_today_from_utc(now); let since = today - Duration::days(clamped as i64 - 1); let since_ms = Local @@ -431,7 +431,7 @@ pub fn model_cost_summary_from_rows( now: DateTime, days: u32, ) -> ModelCostSummary { - let clamped = days.clamp(1, 365); + let clamped = crate::cost_reporting_period::clamp_window_days(days); let today = local_today_from_utc(now); let since = today - Duration::days(clamped as i64 - 1); let since_ms = Local diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..f1e9b19d23 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -16,6 +16,7 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use crate::core::ProviderId; +use crate::cost_reporting_period::CostReportingPeriod; /// Stable namespace used by the desktop bridge for quota metric rows. pub const USAGE_ITEM_METRIC_PREFIX: &str = "metric:"; @@ -369,6 +370,11 @@ pub struct Settings { /// Hide native Codex spend rows when an OpenCodex import is present. #[serde(default)] pub hide_native_codex_cost_when_open_codex_present: bool, + + /// History window for local cost surfaces: `rolling:N`, `month-to-date`, or + /// `all` (upstream 0.67.0). Missing or unreadable values read as `rolling:30`. + #[serde(default)] + pub cost_reporting_period: CostReportingPeriod, } fn default_window_scale_percent() -> u16 { @@ -595,6 +601,7 @@ impl Default for Settings { cost_summary_display_style: CostSummaryDisplayStyle::default(), open_codex_usage_logs_enabled: false, hide_native_codex_cost_when_open_codex_present: false, + cost_reporting_period: CostReportingPeriod::default(), } } } diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..6cd64e957f 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -183,6 +183,8 @@ pub(super) struct RawSettings { open_codex_usage_logs_enabled: bool, #[serde(default)] hide_native_codex_cost_when_open_codex_present: bool, + #[serde(default)] + cost_reporting_period: CostReportingPeriod, } impl Default for RawSettings { @@ -290,6 +292,7 @@ impl Default for RawSettings { open_codex_usage_logs_enabled: s.open_codex_usage_logs_enabled, hide_native_codex_cost_when_open_codex_present: s .hide_native_codex_cost_when_open_codex_present, + cost_reporting_period: s.cost_reporting_period, } } } @@ -600,6 +603,7 @@ impl From for Settings { hide_native_codex_cost_when_open_codex_present: raw .hide_native_codex_cost_when_open_codex_present, codex_external_oauth_sources_allowed: raw.codex_external_oauth_sources_allowed, + cost_reporting_period: raw.cost_reporting_period, } } } diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..531ab3052f 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -115,6 +115,42 @@ fn open_codex_usage_logs_default_off_and_round_trip() { assert!(loaded.hide_native_codex_cost_when_open_codex_present); } +#[test] +fn cost_reporting_period_defaults_to_thirty_days_and_round_trips() { + let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) + .expect("missing cost_reporting_period defaults"); + assert_eq!( + defaulted.cost_reporting_period, + CostReportingPeriod::Rolling(30) + ); + + for period in [ + CostReportingPeriod::Rolling(90), + CostReportingPeriod::MonthToDate, + CostReportingPeriod::AllAvailable, + ] { + let settings = Settings { + cost_reporting_period: period, + ..Settings::default() + }; + let json = serde_json::to_string(&settings).expect("serialize cost period"); + assert!(json.contains(&format!(r#""cost_reporting_period":"{}""#, period.raw()))); + let loaded: Settings = serde_json::from_str(&json).expect("deserialize cost period"); + assert_eq!(loaded.cost_reporting_period, period); + } +} + +#[test] +fn unreadable_cost_reporting_period_loads_as_the_default() { + let loaded: Settings = + serde_json::from_str(r#"{ "enabled_providers": [], "cost_reporting_period": "weekly" }"#) + .expect("an unreadable period must not fail the whole settings file"); + assert_eq!( + loaded.cost_reporting_period, + CostReportingPeriod::Rolling(30) + ); +} + #[test] fn notification_sound_paths_round_trip_and_default_for_existing_settings() { let settings = Settings { From 5c37174ce4b0feb2f953f12688c6af95affbc6a0 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:57:32 +0700 Subject: [PATCH 115/407] Port upstream 0.66.0: import Kimi web access tokens from Chromium local storage --- docs/COOKIES.md | 4 + rust/src/providers/kimi/local_storage.rs | 186 +++++++++++++++++++++++ rust/src/providers/kimi/mod.rs | 6 +- rust/src/providers/kimi/web.rs | 79 +++++++++- 4 files changed, 272 insertions(+), 3 deletions(-) create mode 100644 rust/src/providers/kimi/local_storage.rs diff --git a/docs/COOKIES.md b/docs/COOKIES.md index 804562d077..22b6c73dc7 100644 --- a/docs/COOKIES.md +++ b/docs/COOKIES.md @@ -37,6 +37,10 @@ If automatic extraction fails (for example, Chromium App-Bound Encryption is act 4. Copy the `Cookie` header value from **Request Headers** 5. In CodexBar Settings → provider detail → **Browser Cookies**, paste the value +## Kimi local-storage tokens + +With the Kimi cookie source set to automatic, CodexBar also reads `access_token` from Chromium browsers' `Local Storage` for the selected Kimi region (`www.kimi.com` or `www.kimi.ai`), after the Kimi Desktop session and browser cookies. Local storage is not App-Bound encrypted, so this can work when cookie decryption is blocked. Only unexpired three-segment JWTs are used; refresh tokens are never read. A manual Cookie header always wins, and Cookie source Off or Manual skips this step. Open Kimi in the browser to renew an expired session. Firefox and Safari local storage are not read, and only the `Default` and `Profile N` profiles are scanned. + ## Troubleshooting - **"Chromium App-Bound Encryption"**: Modern Chrome, Edge, Brave, and other Chromium-based profiles can protect cookies with ABE. Closing the browser does not remove ABE; use a manual Cookie header or Firefox for the same login diff --git a/rust/src/providers/kimi/local_storage.rs b/rust/src/providers/kimi/local_storage.rs new file mode 100644 index 0000000000..9314082b13 --- /dev/null +++ b/rust/src/providers/kimi/local_storage.rs @@ -0,0 +1,186 @@ +//! Kimi web access tokens from Chromium local storage (upstream #3923). +//! +//! Upstream `KimiCookieImporter.localStorageTokens(region:)` reads `access_token` from each +//! Chromium profile's `Local Storage/leveldb` for the selected region's web origin and keeps only +//! current three-segment ASCII JWTs. It never reads or refreshes refresh tokens, and it runs after +//! cookie discovery, so a manual credential or a cookie session always takes precedence. + +use std::collections::HashSet; +use std::path::PathBuf; +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, +}; +use crate::codex_accounts::api::jwt_payload; + +const ACCESS_TOKEN_KEY: &str = "access_token"; + +/// 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()); + tokens_from_profiles(&chromium_profile_dirs(), region.web_base_url(), now) +} + +fn chromium_profile_dirs() -> Vec { + BrowserDetector::detect_all() + .into_iter() + .filter(|browser| browser.browser_type.is_chromium_based()) + .flat_map(|browser| browser.profiles) + .map(|profile| profile.path) + .collect() +} + +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 dir = local_storage_dir(profile); + if !dir.is_dir() { + continue; + } + match read_local_storage_entries(&dir, origin) { + Ok(entries) => { + for token in access_tokens(&entries, now_unix) { + if seen.insert(token.clone()) { + tokens.push(token); + } + } + } + Err(error) => tracing::debug!(%error, "Kimi local storage is not readable"), + } + } + tokens +} + +fn access_tokens(entries: &[LocalStorageEntry], now_unix: f64) -> impl Iterator { + entries + .iter() + .filter(|entry| entry.key == ACCESS_TOKEN_KEY) + .map(|entry| normalized_value(&entry.value)) + .filter(move |token| is_current_jwt(token, now_unix)) +} + +/// Web storage may hold the token as a JSON string (`"eyJ..."`) or as bare text. +fn normalized_value(value: &str) -> String { + serde_json::from_str::(value).unwrap_or_else(|_| value.trim().to_string()) +} + +fn is_current_jwt(token: &str, now_unix: f64) -> bool { + token.split('.').count() == 3 + && token + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')) + && jwt_payload(token) + .and_then(|payload| payload.get("exp").and_then(serde_json::Value::as_f64)) + .is_some_and(|exp| exp.is_finite() && exp > now_unix) +} + +#[cfg(test)] +mod tests { + use super::*; + + const NOW: f64 = 1_800_000_000.0; + // Fixture from upstream `KimiLocalStorageTests`: payload {"exp":1800003600}. + const TOKEN: &str = "eyJhbGciOiJIUzI1NiJ9.eyJleHAiOjE4MDAwMDM2MDB9.signature"; + // Payload {"exp":1}. + const EXPIRED: &str = "eyJhbGciOiJIUzI1NiJ9.eyJleHAiOjF9.signature"; + + fn entry(key: &str, value: &str) -> LocalStorageEntry { + LocalStorageEntry { + key: key.into(), + value: value.into(), + } + } + + fn tokens(entries: &[LocalStorageEntry], now: f64) -> Vec { + access_tokens(entries, now).collect() + } + + #[test] + fn only_current_access_tokens_are_kept() { + let entries = [ + entry("refresh_token", TOKEN), + entry("access_token", EXPIRED), + entry("access_token", "not-a-token"), + entry("access_token", "a.b.c"), + entry("access_token", &format!("{TOKEN}; kimi-auth=x")), + entry("access_token", &format!(" {TOKEN}\n")), + entry("access_token", &format!("\"{TOKEN}\"")), + ]; + + assert_eq!(tokens(&entries, NOW), vec![TOKEN, TOKEN]); + assert!(tokens(&entries, NOW + 3600.0).is_empty()); + } + + #[test] + fn non_finite_or_missing_expiry_is_rejected() { + // Payloads {"exp":"soon"} and {"sub":"x"}. + for payload in ["eyJleHAiOiJzb29uIn0", "eyJzdWIiOiJ4In0"] { + let token = format!("eyJhbGciOiJIUzI1NiJ9.{payload}.signature"); + assert!(tokens(&[entry("access_token", &token)], NOW).is_empty()); + } + } + + /// Build a one-record LevelDB write-ahead log holding one Local Storage `access_token`. + fn write_log(origin: &str, value: &str) -> Vec { + let mut key = format!("_{origin}\0").into_bytes(); + key.push(1); + key.extend_from_slice(ACCESS_TOKEN_KEY.as_bytes()); + let mut value_bytes = vec![1]; + value_bytes.extend_from_slice(value.as_bytes()); + + let mut batch = 1u64.to_le_bytes().to_vec(); + batch.extend_from_slice(&1u32.to_le_bytes()); + batch.push(1); + batch.push(u8::try_from(key.len()).unwrap()); + batch.extend_from_slice(&key); + batch.push(u8::try_from(value_bytes.len()).unwrap()); + batch.extend_from_slice(&value_bytes); + + let mut record = vec![0, 0, 0, 0]; + record.extend_from_slice(&u16::try_from(batch.len()).unwrap().to_le_bytes()); + record.push(1); + record.extend_from_slice(&batch); + record + } + + fn profile_with(origin: &str, value: &str) -> tempfile::TempDir { + let profile = tempfile::tempdir().unwrap(); + let dir = local_storage_dir(profile.path()); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("000003.log"), write_log(origin, value)).unwrap(); + profile + } + + #[test] + fn profiles_are_read_for_the_selected_region_only_and_deduplicated() { + let first = profile_with(KimiRegion::International.web_base_url(), TOKEN); + let second = profile_with(KimiRegion::International.web_base_url(), TOKEN); + let china = profile_with(KimiRegion::China.web_base_url(), TOKEN); + let empty = tempfile::tempdir().unwrap(); + let profiles = [ + empty.path().to_path_buf(), + first.path().to_path_buf(), + second.path().to_path_buf(), + china.path().to_path_buf(), + ]; + + assert_eq!( + tokens_from_profiles(&profiles, KimiRegion::International.web_base_url(), NOW), + vec![TOKEN] + ); + assert_eq!( + tokens_from_profiles(&profiles, KimiRegion::China.web_base_url(), NOW), + vec![TOKEN] + ); + assert!( + tokens_from_profiles(&profiles[..3], KimiRegion::China.web_base_url(), NOW).is_empty() + ); + } +} diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index e326afa3d6..bb54db7363 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -5,15 +5,19 @@ //! Provider policies are centralized in the submodules: //! - [`web`]: `kimi.com` cookie auth, browser-import gate (Cookie Source Off, //! upstream #2623), and the shared web-token resolution chain -//! (manual cookie → Kimi Desktop session → browser import). +//! (manual cookie → Kimi Desktop session → browser import → Chromium +//! local-storage `access_token`, upstream #3923). //! - [`code_api`]: Kimi Code API auth/endpoint/CLI-credential policy, plus the //! upstream 0.48.0 monthly-membership enrichment of Code API + CLI usage //! from a signed-in Kimi Desktop session (#2622). +//! - [`local_storage`]: current `access_token` JWTs from Chromium local storage +//! for the selected region, read through the shared LevelDB reader. //! - [`desktop_token`]: read-only, WAL-safe reader for the Kimi Desktop //! (Electron) Chromium cookie store. mod code_api; pub mod desktop_token; +mod local_storage; mod region; mod web; diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index c5ed5b8305..3e9cd3e4ca 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -3,13 +3,15 @@ //! Upstream 0.48.0 policies ported here (#2623 / `KimiBrowserImportPolicy`): //! browser cookie import — and reading the Kimi Desktop session store — are //! disabled when the Kimi cookie source is `off`. The shared token chain is -//! manual cookie header → Kimi Desktop session → browser import (upstream -//! `KimiWebEnrichmentTokenResolver`), used both by the web fetch itself and +//! manual cookie header → Kimi Desktop session → browser cookie import → +//! Chromium local-storage `access_token` (upstream `KimiWebEnrichmentTokenResolver` +//! and #3923), used both by the web fetch itself and //! by the Code-API/CLI monthly enrichment. use reqwest::Client; use super::desktop_token::KimiDesktopAuthToken; +use super::local_storage::local_storage_tokens; use super::{ KIMI_SUBSCRIPTION_SERVICE, KIMI_SUBSCRIPTION_STATS_SERVICE, KIMI_WEB_USAGE_SERVICE, KimiProvider, KimiRegion, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, @@ -46,6 +48,7 @@ fn browser_import_error(cookie_source: &str) -> ProviderError { /// 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). +/// 4. Chromium local-storage `access_token` for the region (automatic source only). pub(crate) fn web_auth_tokens(manual_header: Option<&str>, region: KimiRegion) -> Vec { resolve_web_tokens(WebTokenInput { manual_header, @@ -53,6 +56,7 @@ pub(crate) fn web_auth_tokens(manual_header: Option<&str>, region: KimiRegion) - region, desktop_token: KimiDesktopAuthToken::load_for_region, browser_token: browser_auth_token, + local_storage_tokens, }) .into_iter() .map(|candidate| candidate.token) @@ -65,6 +69,7 @@ struct WebTokenInput<'a> { region: KimiRegion, desktop_token: fn(KimiRegion) -> Option, browser_token: fn(KimiRegion) -> Option, + local_storage_tokens: fn(KimiRegion) -> Vec, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -72,6 +77,7 @@ enum WebTokenSource { Manual, Desktop, Browser, + LocalStorage, } #[derive(Debug, PartialEq, Eq)] @@ -111,6 +117,14 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { source: WebTokenSource::Browser, }); } + for token in (input.local_storage_tokens)(input.region) { + if seen.insert(token.clone()) { + candidates.push(WebTokenCandidate { + token, + source: WebTokenSource::LocalStorage, + }); + } + } candidates } @@ -173,6 +187,18 @@ pub(crate) async fn fetch_via_web( } } + // Local-storage tokens are read last, only after every cookie source was rejected. + for token in local_storage_tokens(region) { + if !seen.insert(token.clone()) { + continue; + } + match fetch_via_web_token(&client, &token, region).await { + Ok(usage) => return Ok(usage), + Err(ProviderError::AuthRequired) => {} + Err(error) => return Err(error), + } + } + Err(ProviderError::AuthRequired) } @@ -339,6 +365,14 @@ mod tests { None } + fn no_local_storage(_: KimiRegion) -> Vec { + Vec::new() + } + + fn static_local_storage(_: KimiRegion) -> Vec { + vec!["browser-token".to_string(), "storage-token".to_string()] + } + fn input<'a>( manual_header: Option<&'a str>, cookie_source: &'a str, @@ -351,6 +385,7 @@ mod tests { region: KimiRegion::China, desktop_token, browser_token, + local_storage_tokens: no_local_storage, } } @@ -440,6 +475,46 @@ mod tests { ); } + #[test] + fn local_storage_tokens_follow_cookie_sources_without_duplicates() { + let candidates = resolve_web_tokens(WebTokenInput { + local_storage_tokens: static_local_storage, + ..input(None, "auto", static_desktop, static_browser) + }); + let ordered: Vec<_> = candidates + .iter() + .map(|candidate| (candidate.source, candidate.token.as_str())) + .collect(); + assert_eq!( + ordered, + vec![ + (WebTokenSource::Desktop, "desktop-token"), + (WebTokenSource::Browser, "browser-token"), + (WebTokenSource::LocalStorage, "storage-token"), + ] + ); + } + + #[test] + fn local_storage_is_not_read_for_manual_credentials_or_blocked_sources() { + let manual = resolve_web_tokens(WebTokenInput { + local_storage_tokens: static_local_storage, + ..input(Some("kimi-auth=manual-token"), "auto", no_token, no_token) + }); + assert_eq!(manual.len(), 1); + assert_eq!(manual[0].source, WebTokenSource::Manual); + + for source in ["off", "manual"] { + assert!( + resolve_web_tokens(WebTokenInput { + local_storage_tokens: static_local_storage, + ..input(None, source, no_token, no_token) + }) + .is_empty() + ); + } + } + #[test] fn browser_import_gate_is_case_insensitive() { assert!(!browser_import_allowed("OFF")); From 179ce1f59f688346201dde80167de67ca3830eb2 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 18:57:45 +0700 Subject: [PATCH 116/407] Port upstream 0.66.0: llmman provider Add the llmman provider: memory in use plus loaded and stored models from a local or LAN llmman daemon (GET /llmman/node, best-effort /api/version). Optional LLMMAN_API_KEY, base URL from provider extras or LLMMAN_HOST, and the https-or-private-network-http endpoint policy. --- README.md | 1 + .../src/commands/provider_settings.rs | 6 + .../providers/icons/ProviderIcon-llmman.svg | 4 + .../src/components/providers/providerIcons.ts | 3 + apps/desktop-tauri/src/i18n/keys.ts | 4 + apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 2 +- .../sections/CredentialsDispatcher.tsx | 1 + .../sections/credentials/OpenAiExtras.tsx | 8 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + docs/PROVIDERS.md | 13 + rust/src/cli/serve/dashboard/icons.rs | 4 + .../dashboard/icons/ProviderIcon-llmman.svg | 4 + rust/src/core/provider.rs | 10 +- rust/src/core/provider_factory.rs | 19 +- rust/src/core/token_accounts.rs | 1 + rust/src/locale.rs | 4 + rust/src/locale/en-US.ftl | 4 + rust/src/providers/llmman/endpoint.rs | 140 ++++++ rust/src/providers/llmman/mod.rs | 343 +++++++++++++ rust/src/providers/llmman/tests.rs | 469 ++++++++++++++++++ rust/src/providers/mod.rs | 2 + rust/src/settings/api_keys.rs | 11 + rust/src/settings/provider_workspace.rs | 31 ++ 23 files changed, 1074 insertions(+), 11 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-llmman.svg create mode 100644 rust/src/cli/serve/dashboard/icons/ProviderIcon-llmman.svg create mode 100644 rust/src/providers/llmman/endpoint.rs create mode 100644 rust/src/providers/llmman/mod.rs create mode 100644 rust/src/providers/llmman/tests.rs diff --git a/README.md b/README.md index 2e052ad2f9..0a83e19ae7 100755 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Deepgram | API Key | Project Usage | | Groq | API Key | Enterprise Metrics | | LLM Proxy | API Key | Quota Stats | +| llmman | Local daemon / optional API Key | Memory in use, loaded and stored models | 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 f9344fcfba..e9d32abdcb 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -323,6 +323,7 @@ fn workspace_provider(provider_id: &str) -> Option { Some(match provider_id { "openaiapi" => ProviderId::OpenAIApi, "litellm" => ProviderId::LiteLLM, + "llmman" => ProviderId::LLMMan, "devin" => ProviderId::Devin, "opencodego" => ProviderId::OpenCodeGo, "zed" => ProviderId::Zed, @@ -394,6 +395,11 @@ mod tests { ); } + #[test] + fn maps_llmman_workspace_provider() { + assert_eq!(workspace_provider("llmman"), Some(ProviderId::LLMMan)); + } + #[test] fn litellm_endpoint_change_requires_reentering_saved_key() { assert!( diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-llmman.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-llmman.svg new file mode 100644 index 0000000000..a087c006ca --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-llmman.svg @@ -0,0 +1,4 @@ + + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 2fdc4380d7..3d64590f5e 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -38,6 +38,7 @@ import jetbrains from "./icons/ProviderIcon-jetbrains.svg?raw"; import kilo from "./icons/ProviderIcon-kilo.svg?raw"; import kimi from "./icons/ProviderIcon-kimi.svg?raw"; import kiro from "./icons/ProviderIcon-kiro.svg?raw"; +import llmman from "./icons/ProviderIcon-llmman.svg?raw"; import llmproxy from "./icons/ProviderIcon-llmproxy.svg?raw"; import manus from "./icons/ProviderIcon-manus.svg?raw"; import meta from "./icons/ProviderIcon-meta.svg?raw"; @@ -126,6 +127,7 @@ const RAW: Record = { kilo: tint(kilo), kimi: tint(kimi), kiro: tint(kiro), + llmman: tint(llmman), llmproxy: tint(llmproxy), manus: tint(manus), meta: tint(meta), @@ -197,6 +199,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { kimi: { id: "kimi", brandColor: "#fe603c", fallbackLetter: "☽", svgPath: RAW.kimi }, kimik2: { id: "kimik2", brandColor: "#4c00ff", fallbackLetter: "☽", svgPath: RAW.kimi }, kiro: { id: "kiro", brandColor: "#ff9900", fallbackLetter: "K", svgPath: RAW.kiro }, + llmman: { id: "llmman", brandColor: "#6CC5B0", fallbackLetter: "L", svgPath: RAW.llmman }, llmproxy: { id: "llmproxy", brandColor: "#4f46e5", fallbackLetter: "L", svgPath: RAW.llmproxy }, minimax: { id: "minimax", brandColor: "#fe603c", fallbackLetter: "M", svgPath: RAW.minimax }, mistral: { id: "mistral", brandColor: "#ff500f", fallbackLetter: "M", svgPath: RAW.mistral }, diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..25fea5766b 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -887,6 +887,10 @@ export const ALL_LOCALE_KEYS = [ "Sub2ApiBaseUrlLabel", "Sub2ApiBaseUrlPlaceholder", "Sub2ApiBaseUrlHelp", + "LLMManApiTitle", + "LLMManBaseUrlLabel", + "LLMManBaseUrlPlaceholder", + "LLMManBaseUrlHelp", // Tray icon visibility (Windows 11 hidden-icons overflow) "PromoteTrayIconLabel", diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..e359cf9ce3 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -36,7 +36,7 @@ const HAS_DASHBOARD = new Set([ "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", "opencode", "opencodego", "openrouter", "perplexity", "qoder", "codebuddy", "sakana", "stepfun", "t3chat", "venice", "vertexai", "warp", "windsurf", - "xai", "zai", "fireworks", "meta", "muse", "nous", + "xai", "zai", "fireworks", "meta", "muse", "nous", "llmman", ]); /** Provider IDs that have a status page URL in the backend */ const HAS_STATUS_PAGE = new Set([ diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/CredentialsDispatcher.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CredentialsDispatcher.tsx index 76c2652ecd..05591a7c47 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/CredentialsDispatcher.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/CredentialsDispatcher.tsx @@ -35,6 +35,7 @@ export function CredentialsDispatcher({ providerId, t }: Props) { case "openaiapi": return ; case "litellm": + case "llmman": case "devin": case "opencodego": case "zed": diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/OpenAiExtras.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/OpenAiExtras.tsx index 3ed6842890..2cda52ddb5 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/OpenAiExtras.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/OpenAiExtras.tsx @@ -110,6 +110,7 @@ export function OpenAiExtras({ providerId = "codex", t }: Props) { const WORKSPACE_EXTRA_IDS: Record = { openaiapi: true, litellm: true, + llmman: true, devin: true, opencodego: true, zed: true, @@ -134,6 +135,13 @@ function extraConfig(providerId: string, t: Props["t"]) { placeholder: t("LiteLlmBaseUrlPlaceholder"), help: t("LiteLlmBaseUrlHelp"), }; + case "llmman": + return { + title: t("LLMManApiTitle"), + label: t("LLMManBaseUrlLabel"), + placeholder: t("LLMManBaseUrlPlaceholder"), + help: t("LLMManBaseUrlHelp"), + }; case "devin": return { title: t("DevinApiTitle"), diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index dda4c24441..46dba195d9 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -77,4 +77,5 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["meta", "Meta"], ["muse", "Muse Code"], ["nous", "Nous Portal"], + ["llmman", "llmman"], ]; diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 237f7dab3d..61409f9624 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -55,6 +55,19 @@ spend and, when the optional balance request succeeds, prepaid credit balance. It keeps those values in the cost/detail surfaces and does not invent a quota percentage or use a Replicate API token as a website credential. +### llmman daemon + +llmman reads a local (or LAN) llmman daemon: `GET {base}/llmman/node` for the +memory budget plus loaded and stored models, and a best-effort +`GET {base}/api/version`. The base URL comes from provider extras or +`LLMMAN_HOST` and defaults to `http://127.0.0.1:17434`; a trailing `/v1` is +accepted and dropped per request. `LLMMAN_API_KEY` is optional, because a +daemon without configured keys is open, and is sent as a Bearer token only when +set. Plain HTTP is accepted only for localhost, loopback, private-network +(10/8, 172.16/12, 192.168/16, 169.254/16, fc00::/7, fe80::/10) and `.local` +hosts; any public host must use HTTPS. Embedded credentials, queries, and +fragments are rejected. + ## Listing what is enabled ```powershell diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 6ac230cc63..7e3ec15502 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -193,6 +193,10 @@ static ICONS: &[(&str, &[u8])] = &[ "ProviderIcon-litellm", include_bytes!("icons/ProviderIcon-litellm.svg"), ), + ( + "ProviderIcon-llmman", + include_bytes!("icons/ProviderIcon-llmman.svg"), + ), ( "ProviderIcon-llmproxy", include_bytes!("icons/ProviderIcon-llmproxy.svg"), diff --git a/rust/src/cli/serve/dashboard/icons/ProviderIcon-llmman.svg b/rust/src/cli/serve/dashboard/icons/ProviderIcon-llmman.svg new file mode 100644 index 0000000000..a087c006ca --- /dev/null +++ b/rust/src/cli/serve/dashboard/icons/ProviderIcon-llmman.svg @@ -0,0 +1,4 @@ + + + + diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..b897c8ef95 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -93,6 +93,7 @@ pub enum ProviderId { Muse, Replicate, Nous, + LLMMan, } impl ProviderId { @@ -179,6 +180,7 @@ impl ProviderId { ProviderId::Muse, ProviderId::Replicate, ProviderId::Nous, + ProviderId::LLMMan, ] } @@ -226,6 +228,7 @@ impl ProviderId { ProviderId::Meta => "meta", ProviderId::Muse => "muse", ProviderId::Nous => "nous", + ProviderId::LLMMan => "llmman", ProviderId::AiAnd => "aiand", ProviderId::Windsurf => "windsurf", ProviderId::Manus => "manus", @@ -313,6 +316,7 @@ impl ProviderId { ProviderId::Meta => "Meta", ProviderId::Muse => "Muse Code", ProviderId::Nous => "Nous Portal", + ProviderId::LLMMan => "llmman", ProviderId::AiAnd => "ai&", ProviderId::Windsurf => "Windsurf", ProviderId::Manus => "Manus", @@ -415,6 +419,7 @@ impl ProviderId { ProviderId::Meta => None, ProviderId::Muse => None, ProviderId::Nous => None, + ProviderId::LLMMan => None, ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, @@ -495,6 +500,7 @@ impl ProviderId { "fireworks" | "fireworks-ai" | "fw" => Some(ProviderId::Fireworks), "muse" | "muse-code" | "muse code" => Some(ProviderId::Muse), "nous" | "nous-portal" | "nous portal" | "hermes" => Some(ProviderId::Nous), + "llmman" => Some(ProviderId::LLMMan), "meta" | "metaspark" | "meta-spark" | "muse-spark" | "musespark" | "muse spark" | "meta muse spark" => Some(ProviderId::Meta), "aiand" | "ai&" | "ai-and" | "ai and" => Some(ProviderId::AiAnd), @@ -1106,6 +1112,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Muse => "#0668E1", ProviderId::Replicate => "#000000", ProviderId::Nous => "#D6A55C", + ProviderId::LLMMan => "#6CC5B0", } } @@ -1120,7 +1127,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 80); + assert_eq!(all.len(), 81); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1181,6 +1188,7 @@ mod tests { assert!(all.contains(&ProviderId::Replicate)); assert!(all.contains(&ProviderId::Muse)); assert!(all.contains(&ProviderId::Nous)); + assert!(all.contains(&ProviderId::LLMMan)); } #[test] diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 712c1d4fb9..a2aa8885aa 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -14,15 +14,15 @@ use crate::providers::{ CursorProvider, DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, DoubaoProvider, ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, GrokProvider, GroqProvider, HelmcodeProvider, 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, TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, - WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, - ZoomMateProvider, + JetBrainsProvider, KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMManProvider, + 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, TypeSafeProvider, V0Provider, VeniceProvider, + VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, + ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -111,6 +111,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Meta => Box::new(MetaProvider::new()), ProviderId::Muse => Box::new(MuseProvider::new()), ProviderId::Nous => Box::new(NousProvider::new()), + ProviderId::LLMMan => Box::new(LLMManProvider::new()), } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 6dedc0fbde..e924d26f35 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -377,6 +377,7 @@ impl TokenAccountSupport { | ProviderId::Fireworks | ProviderId::Meta | ProviderId::Nous + | ProviderId::LLMMan | ProviderId::Muse => None, } } diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..2cb81291da 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -1167,6 +1167,10 @@ locale_keys! { Sub2ApiBaseUrlLabel, Sub2ApiBaseUrlPlaceholder, Sub2ApiBaseUrlHelp, + LLMManApiTitle, + LLMManBaseUrlLabel, + LLMManBaseUrlPlaceholder, + LLMManBaseUrlHelp, // Tray icon visibility (Windows 11 hidden-icons overflow) PromoteTrayIconLabel, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..c8a315ce5c 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -797,6 +797,10 @@ Sub2ApiTitle = sub2api Sub2ApiBaseUrlLabel = Base URL Sub2ApiBaseUrlPlaceholder = https://sub2api.example.com Sub2ApiBaseUrlHelp = Deployment URL for GET /v1/usage. Use HTTPS, or loopback HTTP (127.0.0.1 / localhost) for local development. +LLMManApiTitle = llmman +LLMManBaseUrlLabel = Base URL +LLMManBaseUrlPlaceholder = http://127.0.0.1:17434 +LLMManBaseUrlHelp = Optional. Leave blank for the local daemon at http://127.0.0.1:17434. Use HTTPS, or HTTP only for localhost and private-network hosts. The API key is optional for an open daemon. PromoteTrayIconLabel = Pin to Taskbar (Windows 11) PromoteTrayIconHelper = Always show the CodexBar icon on the taskbar instead of inside the overflow chevron PromoteTrayIconUnsupportedHint = Not available on this operating system. Drag the icon out of the ^ overflow area in Windows Settings > Taskbar to pin it manually. diff --git a/rust/src/providers/llmman/endpoint.rs b/rust/src/providers/llmman/endpoint.rs new file mode 100644 index 0000000000..161989f60f --- /dev/null +++ b/rust/src/providers/llmman/endpoint.rs @@ -0,0 +1,140 @@ +//! llmman daemon endpoint normalization and policy. +//! +//! Upstream ships the `https-or-private-network-http` endpoint policy: HTTPS +//! to any host, plain HTTP only to loopback, private-network, and `.local` +//! hosts, so an API key never crosses the public internet unencrypted. + +use std::net::{Ipv4Addr, Ipv6Addr}; + +use reqwest::Url; + +pub(super) const DEFAULT_BASE_URL: &str = "http://127.0.0.1:17434"; +const DEFAULT_PORT: u16 = 17434; + +const POLICY_MESSAGE: &str = "llmman base URL must use HTTPS, or HTTP for localhost and private-network hosts, without embedded credentials, query, or fragment."; + +/// Normalize a user-supplied llmman base URL, or fail with a friendly message. +/// +/// Empty input resolves to the default local daemon. A scheme-less input is +/// HTTP and gets the daemon port unless one is given; `0.0.0.0` (a listen +/// address, not a destination) is rewritten to loopback. The result has no +/// trailing slash; a trailing `/v1` is kept here and dropped per request. +pub(crate) fn validated_llmman_base_url(raw: &str) -> Result { + let trimmed = raw.trim().trim_matches(|c| c == '"' || c == '\'').trim(); + if trimmed.is_empty() { + return Ok(DEFAULT_BASE_URL.to_string()); + } + if trimmed.chars().any(|c| c.is_control() || c.is_whitespace()) { + return Err(POLICY_MESSAGE.to_string()); + } + + let has_scheme = trimmed.contains("://"); + let candidate = if has_scheme { + trimmed.to_string() + } else { + format!("http://{trimmed}") + }; + let url = Url::parse(&candidate).map_err(|_| POLICY_MESSAGE.to_string())?; + + let scheme = url.scheme(); + if !matches!(scheme, "http" | "https") + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err(POLICY_MESSAGE.to_string()); + } + let host = url + .host_str() + .map(Host::parse) + .ok_or_else(|| POLICY_MESSAGE.to_string())?; + if scheme == "http" && !is_private_network_host(&host) { + return Err(POLICY_MESSAGE.to_string()); + } + + let host_text = match host { + Host::Ipv4(ip) if ip.is_unspecified() => Ipv4Addr::LOCALHOST.to_string(), + Host::Ipv4(ip) => ip.to_string(), + Host::Ipv6(ip) => format!("[{ip}]"), + Host::Domain(domain) => domain, + }; + let port = if !has_scheme && !has_explicit_port(&candidate) { + Some(DEFAULT_PORT) + } else { + url.port() + }; + let port_text = port.map(|port| format!(":{port}")).unwrap_or_default(); + let path = url.path().trim_end_matches('/'); + Ok(format!("{scheme}://{host_text}{port_text}{path}")) +} + +/// The base URL requests are built from: `/v1` (the OpenAI-style base agents +/// are given) is not part of the daemon's own routes. +pub(super) fn request_base(base: &str) -> &str { + base.strip_suffix("/v1").unwrap_or(base) +} + +/// Whether the authority of `http://...` text names a port. `Url::port` cannot +/// tell, because it drops the scheme's default port. +fn has_explicit_port(candidate: &str) -> bool { + let authority = candidate + .trim_start_matches("http://") + .split(['/', '?', '#']) + .next() + .unwrap_or_default(); + let host_and_port = authority + .rsplit_once('@') + .map_or(authority, |(_, rest)| rest); + match host_and_port.rsplit_once(']') { + Some((_, after_bracket)) => after_bracket.starts_with(':'), + None => host_and_port.contains(':'), + } +} + +fn is_private_network_host(host: &Host) -> bool { + match host { + Host::Ipv4(ip) => is_private_ipv4(*ip), + Host::Ipv6(ip) => is_private_ipv6(*ip), + Host::Domain(domain) => { + let domain = domain.trim_end_matches('.'); + domain == "localhost" + || domain + .strip_suffix(".local") + .is_some_and(|label| !label.is_empty()) + } + } +} + +enum Host { + Ipv4(Ipv4Addr), + Ipv6(Ipv6Addr), + Domain(String), +} + +impl Host { + /// Classify the serialized host of a parsed URL (IPv6 keeps its brackets). + fn parse(host: &str) -> Self { + if let Some(inner) = host.strip_prefix('[').and_then(|h| h.strip_suffix(']')) + && let Ok(ip) = inner.parse() + { + return Self::Ipv6(ip); + } + match host.parse() { + Ok(ip) => Self::Ipv4(ip), + Err(_) => Self::Domain(host.to_string()), + } + } +} + +fn is_private_ipv4(ip: Ipv4Addr) -> bool { + // Loopback 127/8, 10/8, 172.16/12, 192.168/16, link-local 169.254/16, and + // 0.0.0.0 (rewritten to loopback by the caller). + ip.is_loopback() || ip.is_private() || ip.is_link_local() || ip.is_unspecified() +} + +fn is_private_ipv6(ip: Ipv6Addr) -> bool { + let first = ip.segments()[0]; + // ::1, unique local fc00::/7, link-local fe80::/10. + ip.is_loopback() || first & 0xfe00 == 0xfc00 || first & 0xffc0 == 0xfe80 +} diff --git a/rust/src/providers/llmman/mod.rs b/rust/src/providers/llmman/mod.rs new file mode 100644 index 0000000000..96a5e5b392 --- /dev/null +++ b/rust/src/providers/llmman/mod.rs @@ -0,0 +1,343 @@ +//! llmman provider: memory and model usage of a local (or LAN) llmman daemon. +//! +//! Ported from upstream CodexBar 0.66.0 (`llmman` plugin). The daemon reports +//! its memory budget plus the loaded and stored models at `GET /llmman/node`; +//! the version comes best-effort from `GET /api/version`. The API key is +//! optional because a daemon without configured keys is open. + +mod endpoint; +#[cfg(test)] +mod tests; + +use std::collections::BTreeMap; +use std::time::Duration; + +use async_trait::async_trait; +use reqwest::{Client, Response, StatusCode}; +use serde::Deserialize; + +use crate::core::{ + FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, ProviderStateKind, RateWindow, SourceMode, UsageSnapshot, +}; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +pub(crate) use endpoint::validated_llmman_base_url; +use endpoint::{DEFAULT_BASE_URL, request_base}; + +const CREDENTIAL_TARGET: &str = "codexbar-llmman"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const MAX_BODY_BYTES: usize = 1024 * 1024; +const MAX_LOADED_MODEL_ROWS: usize = 24; +/// JavaScript's `Number.isSafeInteger` bound, which the upstream plugin uses. +const MAX_SAFE_BYTES: u64 = (1 << 53) - 1; +const NOT_REACHABLE_PREFIX: &str = "llmman is not reachable at "; + +pub struct LLMManProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl LLMManProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::LLMMan, + display_name: "llmman", + session_label: "Memory", + weekly_label: "Models", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some(DEFAULT_BASE_URL), + status_page_url: None, + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .build() + .unwrap_or_else(|_| Client::new()), + } + } +} + +impl Default for LLMManProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for LLMManProvider { + fn id(&self) -> ProviderId { + ProviderId::LLMMan + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => { + let base = resolve_base_url(ctx.workspace_id.as_deref())?; + let key = crate::providers::resolve_api_key( + ctx.api_key.as_deref(), + CREDENTIAL_TARGET, + &["LLMMAN_API_KEY"], + ) + .ok(); + fetch_daemon(&self.client, &base, key.as_deref()).await + } + SourceMode::Web | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::OAuth] + } + + /// An unreachable daemon is a local runtime that is not running, not a + /// credential problem. + fn error_state_kind(&self, error: &ProviderError) -> ProviderStateKind { + match error { + ProviderError::Other(message) if message.starts_with(NOT_REACHABLE_PREFIX) => { + ProviderStateKind::LocalRuntimeOffline + } + _ => error.state_kind(), + } + } +} + +/// Saved base URL, else `LLMMAN_HOST`, else the default local daemon. +fn resolve_base_url(saved: Option<&str>) -> Result { + let configured = saved + .filter(|value| !value.trim().is_empty()) + .map(str::to_string) + .or_else(|| std::env::var("LLMMAN_HOST").ok()) + .unwrap_or_default(); + validated_llmman_base_url(&configured).map_err(ProviderError::Other) +} + +async fn fetch_daemon( + client: &Client, + base: &str, + key: Option<&str>, +) -> Result { + let base = request_base(base); + let response = get(client, base, "/llmman/node", key).await?; + check_status(response.status(), base, key.is_some())?; + let node = parse_node(&read_body(response).await?)?; + let version = fetch_version(client, base, key).await; + Ok(build_result(&node, version, key.is_some())) +} + +async fn get( + client: &Client, + base: &str, + path: &str, + key: Option<&str>, +) -> Result { + let mut request = client + .get(format!("{base}{path}")) + .header("Accept", "application/json") + .timeout(REQUEST_TIMEOUT); + if let Some(key) = key { + request = request.bearer_auth(key); + } + request.send().await.map_err(|_| not_reachable(base)) +} + +fn not_reachable(base: &str) -> ProviderError { + ProviderError::Other(format!( + "{NOT_REACHABLE_PREFIX}{base}. Start it with llmman serve." + )) +} + +fn check_status(status: StatusCode, base: &str, has_key: bool) -> Result<(), ProviderError> { + let code = status.as_u16(); + match code { + 200 => Ok(()), + 401 | 403 if !has_key => Err(ProviderError::NotInstalled( + "llmman requires an API key. Set one in Settings or LLMMAN_API_KEY.".into(), + )), + 401 | 403 => Err(ProviderError::OAuthExpired(format!( + "llmman rejected the API key (HTTP {code})." + ))), + 429 => Err(ProviderError::Other("llmman is busy.".into())), + 500.. => Err(ProviderError::Other(format!("llmman error: HTTP {code}"))), + _ => Err(ProviderError::Other(format!( + "{base} is not an llmman daemon (HTTP {code})." + ))), + } +} + +async fn read_body(response: Response) -> Result, ProviderError> { + read_bounded_response(response, MAX_BODY_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => unrecognized_node(), + BoundedBodyError::Read(_) => ProviderError::Other( + "llmman closed the connection before sending the full response.".into(), + ), + }) +} + +fn unrecognized_node() -> ProviderError { + ProviderError::Parse("llmman returned an unrecognized /llmman/node response.".into()) +} + +/// Best effort: the version is display-only and older daemons may answer +/// differently, so every failure is ignored. +async fn fetch_version(client: &Client, base: &str, key: Option<&str>) -> Option { + #[derive(Deserialize)] + struct VersionReply { + version: String, + } + + let response = get(client, base, "/api/version", key).await.ok()?; + if response.status() != StatusCode::OK { + return None; + } + let body = read_bounded_response(response, MAX_BODY_BYTES).await.ok()?; + serde_json::from_slice::(&body) + .ok() + .map(|reply| reply.version) +} + +/// A byte count that must be a non-negative safe integer. +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(try_from = "u64")] +struct ByteCount(u64); + +impl TryFrom for ByteCount { + type Error = String; + + fn try_from(value: u64) -> Result { + if value <= MAX_SAFE_BYTES { + Ok(Self(value)) + } else { + Err(format!("byte count {value} is not a safe integer")) + } + } +} + +#[derive(Debug, Deserialize)] +struct NodeBody { + memory: ByteCount, + loaded: BTreeMap, + stored: BTreeMap, +} + +/// A validated `/llmman/node` reply: sizes in bytes, models largest first +/// with ties broken by name. +#[derive(Debug)] +struct Node { + memory: u64, + loaded: Vec<(String, u64)>, + stored: Vec<(String, u64)>, +} + +fn parse_node(body: &[u8]) -> Result { + let value: serde_json::Value = serde_json::from_slice(body).map_err(|_| unrecognized_node())?; + // A derived struct would also accept a positional JSON array. + if !value.is_object() { + return Err(unrecognized_node()); + } + let node = NodeBody::deserialize(value).map_err(|_| unrecognized_node())?; + Ok(Node { + memory: node.memory.0, + loaded: largest_first(node.loaded), + stored: largest_first(node.stored), + }) +} + +fn largest_first(models: BTreeMap) -> Vec<(String, u64)> { + // BTreeMap iterates by name, and the sort is stable, so ties stay by name. + let mut models: Vec<_> = models + .into_iter() + .map(|(name, size)| (name, size.0)) + .collect(); + models.sort_by_key(|model| std::cmp::Reverse(model.1)); + models +} + +fn total(models: &[(String, u64)]) -> u64 { + models.iter().map(|(_, size)| size).sum() +} + +/// Decimal units, as llmman's own list/ps print them. +fn format_size(bytes: u64) -> String { + if bytes < 1_000 { + return format!("{bytes} B"); + } + let (unit, scale) = match bytes { + 1_000_000_000.. => ("GB", 1e9), + 1_000_000.. => ("MB", 1e6), + _ => ("kB", 1e3), + }; + format!("{:.1} {unit}", bytes as f64 / scale) +} + +fn count_and_size(models: &[(String, u64)]) -> String { + format!("{} · {}", models.len(), format_size(total(models))) +} + +fn build_result(node: &Node, version: Option, has_key: bool) -> ProviderFetchResult { + let in_use = total(&node.loaded); + let primary = if node.memory > 0 { + let mut window = RateWindow::new((in_use as f64 / node.memory as f64 * 100.0).min(100.0)); + window.reset_description = Some(format!( + "{} of {}", + format_size(in_use), + format_size(node.memory) + )); + window + } else { + RateWindow::informational("Memory limit not reported") + }; + let login_method = if has_key { "API key" } else { "Local daemon" }; + let usage = UsageSnapshot::new(primary).with_login_method(login_method); + + let mut result = ProviderFetchResult::new(usage, "api") + .with_display_detail(ProviderDisplayDetail::new( + "loaded", + "Loaded", + count_and_size(&node.loaded), + )) + .with_display_detail(ProviderDisplayDetail::new( + "stored", + "Stored", + count_and_size(&node.stored), + )) + .with_display_detail( + version.and_then(|version| ProviderDisplayDetail::new("version", "Version", version)), + ); + for row in loaded_model_rows(node) { + result = result.with_display_detail(Some(row)); + } + result +} + +/// Up to 24 loaded models, largest first; names that cannot be shown are +/// skipped before the cap is applied. +fn loaded_model_rows(node: &Node) -> Vec { + let mut rows = Vec::new(); + for (name, weight) in &node.loaded { + if rows.len() == MAX_LOADED_MODEL_ROWS { + break; + } + let id = format!("loaded-model-{}", rows.len()); + let Some(row) = ProviderDisplayDetail::new(id, name, format_size(*weight)) else { + continue; + }; + let row = match node.memory { + 0 => Some(row), + memory => row.with_progress(weight.min(&memory).to_owned() as f64, memory as f64), + }; + rows.extend(row); + } + rows +} diff --git a/rust/src/providers/llmman/tests.rs b/rust/src/providers/llmman/tests.rs new file mode 100644 index 0000000000..998bf3cb71 --- /dev/null +++ b/rust/src/providers/llmman/tests.rs @@ -0,0 +1,469 @@ +use super::*; +use mockito::Matcher; + +const NODE_BODY: &str = r#"{"memory":40000000000,"loaded":{"fixture/small:Q4_K_M":2000000000,"fixture/large:Q4_K_M":8000000000},"stored":{"fixture/small:Q4_K_M":2000000000,"fixture/large:Q4_K_M":8000000000,"fixture/idle":500000000}}"#; +const VERSION_BODY: &str = r#"{"version":"0.9.0","pid":1}"#; + +fn client() -> Client { + Client::builder().no_proxy().build().unwrap() +} + +fn row<'a>(result: &'a ProviderFetchResult, id: &str) -> &'a ProviderDisplayDetail { + result + .display_details() + .iter() + .find(|row| row.id() == id) + .unwrap_or_else(|| panic!("missing detail row {id}")) +} + +// -- Endpoint policy --------------------------------------------------------- + +#[test] +fn accepted_endpoints_normalize() { + for (input, expected) in [ + ("localhost", "http://localhost:17434"), + ("127.0.0.1", "http://127.0.0.1:17434"), + ("0.0.0.0", "http://127.0.0.1:17434"), + ("0.0.0.0:18000", "http://127.0.0.1:18000"), + ("http://0.0.0.0:18000", "http://127.0.0.1:18000"), + ("192.168.1.10", "http://192.168.1.10:17434"), + ("daemon.local", "http://daemon.local:17434"), + ("[::1]", "http://[::1]:17434"), + ("127.0.0.1:18000", "http://127.0.0.1:18000"), + ("http://localhost", "http://localhost"), + ("https://llmman.example.com", "https://llmman.example.com"), + ("http://127.0.0.1:17434///", "http://127.0.0.1:17434"), + ("http://127.0.0.1:17434/v1/", "http://127.0.0.1:17434/v1"), + (" \"localhost\" ", "http://localhost:17434"), + ("", "http://127.0.0.1:17434"), + ] { + assert_eq!( + validated_llmman_base_url(input).as_deref(), + Ok(expected), + "input {input:?}" + ); + } +} + +#[test] +fn private_network_http_hosts_are_accepted() { + for host in [ + "http://10.1.2.3", + "http://172.16.0.5", + "http://172.31.255.255", + "http://192.168.0.1", + "http://169.254.10.10", + "http://127.9.9.9", + "http://[fd12:3456::1]", + "http://[fe80::1]", + "http://workstation.local", + ] { + assert!(validated_llmman_base_url(host).is_ok(), "{host}"); + } +} + +#[test] +fn rejected_endpoints_fail_the_policy() { + for input in [ + "http://127.0.0.1:17434?probe=1", + "http://127.0.0.1:17434#x", + "localhost?probe=1", + "127.0.0.1#x", + "http://127.0.0.1:17434?", + "http://127.0.0.1:17434#", + "http://public.example.com", + "public.example.com", + "http://user:password@127.0.0.1:17434", + "https://user:password@example.com", + "file:///tmp/llmman", + "ftp://127.0.0.1", + "http://172.32.0.1", + "http://8.8.8.8", + "http://[2001:db8::1]", + "http://.local", + "http://exa mple.local", + ] { + assert!( + validated_llmman_base_url(input).is_err(), + "input {input:?} should be rejected" + ); + } +} + +#[test] +fn requests_drop_a_trailing_v1() { + assert_eq!( + request_base("http://127.0.0.1:17434/v1"), + "http://127.0.0.1:17434" + ); + assert_eq!( + request_base("http://127.0.0.1:17434"), + "http://127.0.0.1:17434" + ); + assert_eq!( + request_base("https://h.example.com/proxy/v1"), + "https://h.example.com/proxy" + ); +} + +// -- Body mapping ------------------------------------------------------------ + +#[test] +fn node_fixture_maps_memory_and_models() { + let node = parse_node(NODE_BODY.as_bytes()).unwrap(); + let result = build_result(&node, Some("0.9.0".into()), false); + + let primary = &result.usage.primary; + assert_eq!(primary.used_percent, 25.0); + assert_eq!( + primary.reset_description.as_deref(), + Some("10.0 GB of 40.0 GB") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("Local daemon")); + + assert_eq!(row(&result, "loaded").value(), "2 · 10.0 GB"); + assert_eq!(row(&result, "stored").value(), "3 · 10.5 GB"); + assert_eq!(row(&result, "version").value(), "0.9.0"); + + let large = row(&result, "loaded-model-0"); + assert_eq!(large.title(), "fixture/large:Q4_K_M"); + assert_eq!(large.value(), "8.0 GB"); + let progress = large.progress().unwrap(); + assert_eq!(progress.used() / progress.total(), 0.2); + assert_eq!( + row(&result, "loaded-model-1").title(), + "fixture/small:Q4_K_M" + ); +} + +#[test] +fn api_key_sets_the_login_method() { + let node = parse_node(NODE_BODY.as_bytes()).unwrap(); + let result = build_result(&node, None, true); + assert_eq!(result.usage.login_method.as_deref(), Some("API key")); + assert!( + result + .display_details() + .iter() + .all(|row| row.id() != "version") + ); +} + +#[test] +fn idle_daemon_still_renders() { + let node = parse_node(br#"{"memory":0,"loaded":{},"stored":{}}"#).unwrap(); + let result = build_result(&node, None, false); + assert!(result.usage.primary.is_informational); + assert_eq!(row(&result, "loaded").value(), "0 · 0 B"); + assert_eq!(row(&result, "stored").value(), "0 · 0 B"); + assert!( + result + .display_details() + .iter() + .all(|row| !row.id().starts_with("loaded-model")) + ); +} + +#[test] +fn loaded_models_without_a_memory_budget_have_no_progress() { + let node = parse_node(br#"{"memory":0,"loaded":{"a":5},"stored":{}}"#).unwrap(); + let result = build_result(&node, None, false); + let model = row(&result, "loaded-model-0"); + assert_eq!(model.value(), "5 B"); + assert!(model.progress().is_none()); +} + +#[test] +fn loaded_models_are_largest_first_with_ties_by_name_and_capped() { + let mut loaded = serde_json::Map::new(); + for index in 0..30 { + loaded.insert(format!("model-{index:02}"), serde_json::json!(100)); + } + loaded.insert("big".into(), serde_json::json!(900)); + loaded.insert("bad\nname".into(), serde_json::json!(1_000)); + let body = serde_json::json!({"memory": 10_000, "loaded": loaded, "stored": {}}); + let node = parse_node(body.to_string().as_bytes()).unwrap(); + let rows = loaded_model_rows(&node); + + // The unshowable name is skipped before the 24-row cap is applied. + assert_eq!(rows.len(), 24); + assert_eq!(rows[0].title(), "big"); + assert_eq!(rows[1].title(), "model-00"); + assert_eq!(rows[2].title(), "model-01"); + assert_eq!(rows[23].title(), "model-22"); +} + +#[test] +fn progress_is_capped_at_the_memory_budget() { + let node = parse_node(br#"{"memory":10,"loaded":{"a":50},"stored":{}}"#).unwrap(); + let result = build_result(&node, None, false); + assert_eq!(result.usage.primary.used_percent, 100.0); + let progress = row(&result, "loaded-model-0").progress().unwrap(); + assert_eq!(progress.used(), progress.total()); +} + +#[test] +fn sizes_use_decimal_units() { + assert_eq!(format_size(0), "0 B"); + assert_eq!(format_size(999), "999 B"); + assert_eq!(format_size(1_000), "1.0 kB"); + assert_eq!(format_size(1_500_000), "1.5 MB"); + assert_eq!(format_size(2_000_000_000), "2.0 GB"); + assert_eq!(format_size(40_000_000_000), "40.0 GB"); +} + +#[test] +fn malformed_bodies_are_parse_failures() { + for body in [ + "not json", + "[]", + "[1,{},{}]", + "null", + r#"{"memory":-1,"loaded":{},"stored":{}}"#, + r#"{"memory":1.5,"loaded":{},"stored":{}}"#, + r#"{"memory":9007199254740992,"loaded":{},"stored":{}}"#, + r#"{"memory":"40","loaded":{},"stored":{}}"#, + r#"{"memory":1,"loaded":[],"stored":{}}"#, + r#"{"memory":1,"loaded":{"a":"1"},"stored":{}}"#, + r#"{"memory":1,"loaded":{"a":-1},"stored":{}}"#, + r#"{"memory":1,"loaded":{}}"#, + r#"{"loaded":{},"stored":{}}"#, + ] { + assert!( + matches!(parse_node(body.as_bytes()), Err(ProviderError::Parse(_))), + "body {body:?} should not parse" + ); + } +} + +// -- HTTP behavior ----------------------------------------------------------- + +#[tokio::test] +async fn fetches_node_and_version_with_a_bearer_key() { + let mut server = mockito::Server::new_async().await; + let node = server + .mock("GET", "/llmman/node") + .match_header("authorization", "Bearer sk-llmman") + .with_body(NODE_BODY) + .create_async() + .await; + let version = server + .mock("GET", "/api/version") + .match_header("authorization", "Bearer sk-llmman") + .with_body(VERSION_BODY) + .create_async() + .await; + + let result = fetch_daemon(&client(), &server.url(), Some("sk-llmman")) + .await + .unwrap(); + + node.assert_async().await; + version.assert_async().await; + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!(row(&result, "version").value(), "0.9.0"); + assert_eq!(result.usage.login_method.as_deref(), Some("API key")); +} + +#[tokio::test] +async fn open_daemon_is_queried_without_an_authorization_header() { + let mut server = mockito::Server::new_async().await; + let node = server + .mock("GET", "/llmman/node") + .match_header("authorization", Matcher::Missing) + .with_body(NODE_BODY) + .create_async() + .await; + server + .mock("GET", "/api/version") + .with_status(404) + .create_async() + .await; + + let result = fetch_daemon(&client(), &server.url(), None).await.unwrap(); + + node.assert_async().await; + assert_eq!(result.usage.login_method.as_deref(), Some("Local daemon")); + // The version request failed, which never fails the fetch. + assert!( + result + .display_details() + .iter() + .all(|row| row.id() != "version") + ); +} + +#[tokio::test] +async fn a_v1_base_is_dropped_from_request_paths() { + let mut server = mockito::Server::new_async().await; + let node = server + .mock("GET", "/llmman/node") + .with_body(NODE_BODY) + .create_async() + .await; + server + .mock("GET", "/api/version") + .with_body(VERSION_BODY) + .create_async() + .await; + + fetch_daemon(&client(), &format!("{}/v1", server.url()), None) + .await + .unwrap(); + + node.assert_async().await; +} + +#[tokio::test] +async fn unusable_version_replies_are_ignored() { + for reply in [ + "not json", + r#"{"version":7}"#, + r#"{"version":""}"#, + r#"{"version":"bad\nlabel"}"#, + r#"{"pid":1}"#, + ] { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/llmman/node") + .with_body(NODE_BODY) + .create_async() + .await; + server + .mock("GET", "/api/version") + .with_body(reply) + .create_async() + .await; + + let result = fetch_daemon(&client(), &server.url(), None).await.unwrap(); + assert!( + result + .display_details() + .iter() + .all(|row| row.id() != "version"), + "reply {reply:?}" + ); + } +} + +async fn node_status_error(status: usize, key: Option<&str>) -> ProviderError { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/llmman/node") + .with_status(status) + .with_body("sk-echoed-secret") + .create_async() + .await; + fetch_daemon(&client(), &server.url(), key) + .await + .unwrap_err() +} + +#[tokio::test] +async fn rejected_credentials_are_classified_by_key_presence() { + for status in [401, 403] { + let error = node_status_error(status, None).await; + assert!( + matches!(&error, ProviderError::NotInstalled(message) + if message == "llmman requires an API key. Set one in Settings or LLMMAN_API_KEY."), + "{error:?}" + ); + + let error = node_status_error(status, Some("sk-secret")).await; + assert!( + matches!(&error, ProviderError::OAuthExpired(message) + if *message == format!("llmman rejected the API key (HTTP {status}).")), + "{error:?}" + ); + assert!(!format!("{error:?}").contains("sk-secret")); + } +} + +#[tokio::test] +async fn other_statuses_map_to_friendly_errors() { + let busy = node_status_error(429, None).await; + assert!(matches!(&busy, ProviderError::Other(m) if m == "llmman is busy.")); + + let down = node_status_error(503, None).await; + assert!(matches!(&down, ProviderError::Other(m) if m == "llmman error: HTTP 503")); + + let foreign = node_status_error(404, None).await; + assert!( + matches!(&foreign, ProviderError::Other(m) + if m.starts_with("http://127.0.0.1:") && m.ends_with(" is not an llmman daemon (HTTP 404).")), + "{foreign:?}" + ); +} + +#[tokio::test] +async fn unrecognized_body_is_a_parse_error() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/llmman/node") + .with_body("not json") + .create_async() + .await; + let error = fetch_daemon(&client(), &server.url(), None) + .await + .unwrap_err(); + assert!(matches!( + error, + ProviderError::Parse(ref m) if m == "llmman returned an unrecognized /llmman/node response." + )); +} + +#[tokio::test] +async fn unreachable_daemon_reports_offline_runtime() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + drop(listener); + let base = format!("http://127.0.0.1:{port}"); + + let error = fetch_daemon(&client(), &base, None).await.unwrap_err(); + let ProviderError::Other(message) = &error else { + panic!("expected Other, got {error:?}"); + }; + assert_eq!( + *message, + format!("llmman is not reachable at {base}. Start it with llmman serve.") + ); + assert_eq!( + LLMManProvider::new().error_state_kind(&error), + ProviderStateKind::LocalRuntimeOffline + ); +} + +#[test] +fn only_the_unreachable_error_is_an_offline_runtime() { + let provider = LLMManProvider::new(); + for error in [ + ProviderError::Other("llmman is busy.".into()), + ProviderError::NotInstalled("llmman requires an API key.".into()), + ] { + assert_ne!( + provider.error_state_kind(&error), + ProviderStateKind::LocalRuntimeOffline + ); + } +} + +#[test] +fn saved_base_url_is_validated() { + assert_eq!( + resolve_base_url(Some("192.168.1.10")).as_deref().ok(), + Some("http://192.168.1.10:17434") + ); + assert!(resolve_base_url(Some("http://public.example.com")).is_err()); +} + +#[test] +fn metadata_matches_upstream() { + let provider = LLMManProvider::new(); + let metadata = provider.metadata(); + assert_eq!(metadata.display_name, "llmman"); + assert_eq!(metadata.session_label, "Memory"); + assert_eq!(metadata.weekly_label, "Models"); + assert!(!metadata.default_enabled); + assert!(!metadata.supports_credits); + assert_eq!(metadata.dashboard_url, Some("http://127.0.0.1:17434")); +} diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index de9d265f6e..6c59e74e21 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -49,6 +49,7 @@ pub mod kimi; pub mod kimik2; pub mod kiro; pub mod litellm; +pub mod llmman; pub mod llmproxy; pub mod longcat; pub mod manus; @@ -132,6 +133,7 @@ pub use kimi::KimiProvider; pub use kimik2::KimiK2Provider; pub use kiro::KiroProvider; pub use litellm::LiteLLMProvider; +pub use llmman::LLMManProvider; pub use llmproxy::LLMProxyProvider; pub use longcat::LongCatProvider; pub use manus::ManusProvider; diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index c2adf3565c..d3b12eca2d 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -555,6 +555,17 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: None, }, + ProviderConfigInfo { + id: ProviderId::LLMMan, + name: "llmman", + requires_api_key: false, + api_key_env_var: Some("LLMMAN_API_KEY + LLMMAN_HOST"), + api_key_help: Some( + "Optional: an open local daemon needs no key. Set the base URL in provider extras or LLMMAN_HOST (default http://127.0.0.1:17434).", + ), + config_file_path: None, + dashboard_url: None, + }, ProviderConfigInfo { id: ProviderId::Poe, name: "Poe", diff --git a/rust/src/settings/provider_workspace.rs b/rust/src/settings/provider_workspace.rs index 338747c8ef..5701e2b0a4 100644 --- a/rust/src/settings/provider_workspace.rs +++ b/rust/src/settings/provider_workspace.rs @@ -57,6 +57,8 @@ pub fn validate_provider_workspace_value( } ProviderId::LiteLLM => validate_token_endpoint(trimmed, "LiteLLM base URL", |_| true), ProviderId::Sub2Api => validate_sub2api_base_url(trimmed), + ProviderId::LLMMan => crate::providers::llmman::validated_llmman_base_url(trimmed) + .map_err(|err| err.to_string()), _ => Ok(trimmed.to_string()), } } @@ -219,6 +221,35 @@ mod tests { } } + #[test] + fn validates_and_normalizes_llmman_base_url() { + for (value, expected) in [ + ("localhost", "http://localhost:17434"), + ("0.0.0.0:18000", "http://127.0.0.1:18000"), + ("https://llmman.example.com/", "https://llmman.example.com"), + ] { + assert_eq!( + validate_provider_workspace_value(ProviderId::LLMMan, value).unwrap(), + expected + ); + } + assert_eq!( + validate_provider_workspace_value(ProviderId::LLMMan, " ").unwrap(), + "" + ); + for value in [ + "http://public.example.com", + "http://user:password@127.0.0.1:17434", + "http://127.0.0.1:17434?probe=1", + "file:///tmp/llmman", + ] { + assert!( + validate_provider_workspace_value(ProviderId::LLMMan, value).is_err(), + "accepted {value}" + ); + } + } + #[test] fn validates_sub2api_base_url() { assert_eq!( From 75159922a3d454251d286c05308968ce90b86630 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:08:07 +0700 Subject: [PATCH 117/407] Port upstream 0.66.0: discover MiniMax browser storage across Chromium profiles --- rust/src/browser/mod.rs | 1 + rust/src/browser/storage_discovery.rs | 139 +++++++++++ rust/src/browser/storage_discovery/tests.rs | 158 +++++++++++++ rust/src/providers/minimax/local_storage.rs | 243 ++++++++++++-------- 4 files changed, 443 insertions(+), 98 deletions(-) create mode 100644 rust/src/browser/storage_discovery.rs create mode 100644 rust/src/browser/storage_discovery/tests.rs diff --git a/rust/src/browser/mod.rs b/rust/src/browser/mod.rs index 035fb6e6d2..b15cbdf3c5 100755 --- a/rust/src/browser/mod.rs +++ b/rust/src/browser/mod.rs @@ -4,6 +4,7 @@ pub mod cookie_cache; pub mod cookies; pub mod detection; pub mod leveldb; +pub mod storage_discovery; pub mod watchdog; pub mod wsl_paths; diff --git a/rust/src/browser/storage_discovery.rs b/rust/src/browser/storage_discovery.rs new file mode 100644 index 0000000000..49d3a63d53 --- /dev/null +++ b/rust/src/browser/storage_discovery.rs @@ -0,0 +1,139 @@ +//! Locates raw Chromium profile stores (Local Storage, Session Storage, IndexedDB) across the +//! installed Chromium-family browsers. +//! +//! Each consumer supplies its own decoder; this module only answers "which directories could hold +//! the data". It never opens, locks, or decrypts anything, so it is safe to run while the browser +//! is open and needs no credential-store access. +//! +//! The browser catalog is [`super::detection::BrowserType`] (via [`BrowserDetector`]); a +//! Chromium-family browser that is not in that catalog is not visited. + +use std::path::{Path, PathBuf}; + +use super::detection::BrowserDetector; +use super::leveldb::local_storage::local_storage_dir; + +const INDEXED_DB_DIR: &str = "IndexedDB"; +const SESSION_STORAGE_DIR: &str = "Session Storage"; +const INDEXED_DB_SUFFIX: &str = ".indexeddb.leveldb"; + +/// Which per-profile store to locate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StorageKind { + /// `/Local Storage/leveldb` + LocalStorage, + /// `/Session Storage` + SessionStorage, + /// `/IndexedDB/.indexeddb.leveldb` for every database whose directory name + /// starts with one of `origin_prefixes` (for example `https_platform.minimax.io_`). + IndexedDb { + origin_prefixes: &'static [&'static str], + }, +} + +impl StorageKind { + fn label_suffix(self) -> &'static str { + match self { + Self::LocalStorage => "", + Self::SessionStorage => " (Session Storage)", + Self::IndexedDb { .. } => " (IndexedDB)", + } + } +} + +/// One directory that may hold data of the requested [`StorageKind`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StorageCandidate { + /// Human-readable source, such as `Google Chrome Default (Session Storage)`. + pub label: String, + pub path: PathBuf, +} + +/// Candidates for `kind` in every installed Chromium-family browser, in catalog order. +pub fn discover(kind: StorageKind) -> Vec { + BrowserDetector::detect_all() + .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, + ) + }) + .collect() +} + +/// Candidates for `kind` in one browser `User Data` directory, profiles sorted by name. +/// +/// Only `Default`, `Profile *`, and `user-*` directories count as profiles; guest and system +/// profiles are ignored. +pub 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) + .into_iter() + .map(move |path| StorageCandidate { + label: label.clone(), + path, + }) + }) + .collect() +} + +fn profile_store_paths(profile_dir: &Path, kind: StorageKind) -> Vec { + match kind { + StorageKind::LocalStorage => existing(local_storage_dir(profile_dir)), + StorageKind::SessionStorage => existing(profile_dir.join(SESSION_STORAGE_DIR)), + StorageKind::IndexedDb { origin_prefixes } => { + sorted_child_dirs(&profile_dir.join(INDEXED_DB_DIR), |name| { + name.ends_with(INDEXED_DB_SUFFIX) + && origin_prefixes + .iter() + .any(|prefix| name.starts_with(prefix)) + }) + .into_iter() + .map(|(_, path)| path) + .collect() + } + } +} + +fn existing(path: PathBuf) -> Vec { + if path.exists() { + 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)> { + let Ok(entries) = std::fs::read_dir(dir) else { + return Vec::new(); + }; + let mut children: Vec<(String, PathBuf)> = entries + .flatten() + .filter_map(|entry| { + let name = entry.file_name().into_string().ok()?; + let is_dir = entry.file_type().is_ok_and(|kind| kind.is_dir()); + (is_dir && !name.starts_with('.') && keep(&name)).then(|| (name, entry.path())) + }) + .collect(); + children.sort_by(|left, right| left.0.cmp(&right.0)); + children +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/browser/storage_discovery/tests.rs b/rust/src/browser/storage_discovery/tests.rs new file mode 100644 index 0000000000..e6f2239024 --- /dev/null +++ b/rust/src/browser/storage_discovery/tests.rs @@ -0,0 +1,158 @@ +use super::*; +use std::fs; + +const PREFIXES: &[&str] = &[ + "https_platform.minimax.io_", + "https_www.minimax.io_", + "https_minimax.io_", + "https_platform.minimaxi.com_", + "https_minimaxi.com_", + "https_www.minimaxi.com_", +]; + +fn indexed_db() -> StorageKind { + StorageKind::IndexedDb { + origin_prefixes: PREFIXES, + } +} + +fn mkdir(root: &Path, relative: &str) { + fs::create_dir_all(root.join(relative)).unwrap(); +} + +fn paths(root: &Path, candidates: &[StorageCandidate]) -> Vec { + candidates + .iter() + .map(|candidate| { + candidate + .path + .strip_prefix(root) + .unwrap() + .to_string_lossy() + .replace('\\', "/") + }) + .collect() +} + +#[test] +fn each_storage_kind_resolves_to_its_profile_directory() { + let root = tempfile::tempdir().unwrap(); + let database = "IndexedDB/https_platform.minimax.io_0.indexeddb.leveldb"; + for path in ["Local Storage/leveldb", "Session Storage", database] { + mkdir(root.path(), &format!("Default/{path}")); + } + + let cases = [ + ( + StorageKind::LocalStorage, + "Default/Local Storage/leveldb", + "", + ), + ( + StorageKind::SessionStorage, + "Default/Session Storage", + " (Session Storage)", + ), + (indexed_db(), &format!("Default/{database}"), " (IndexedDB)"), + ]; + for (kind, expected_path, suffix) in cases { + let found = candidates_in_user_data_dir(root.path(), "Google Chrome", kind); + assert_eq!(paths(root.path(), &found), [expected_path]); + assert_eq!(found[0].label, format!("Google Chrome Default{suffix}")); + } +} + +#[test] +fn missing_stores_and_missing_user_data_dir_yield_nothing() { + let root = tempfile::tempdir().unwrap(); + mkdir(root.path(), "Default"); + for kind in [ + StorageKind::LocalStorage, + StorageKind::SessionStorage, + indexed_db(), + ] { + assert!(candidates_in_user_data_dir(root.path(), "Chrome", kind).is_empty()); + assert!( + candidates_in_user_data_dir(&root.path().join("absent"), "Chrome", kind).is_empty() + ); + } +} + +#[test] +fn only_default_profile_and_user_profiles_are_visited_in_name_order() { + let root = tempfile::tempdir().unwrap(); + for profile in [ + "user-work", + "Profile 2", + "Default", + "Guest Profile", + "System Profile", + ".hidden", + "Profile1", + ] { + mkdir(root.path(), &format!("{profile}/Local Storage/leveldb")); + } + // A file named like a profile is not a profile. + fs::write(root.path().join("Profile 9"), b"x").unwrap(); + + let found = candidates_in_user_data_dir(root.path(), "Edge", StorageKind::LocalStorage); + assert_eq!( + found + .iter() + .map(|candidate| candidate.label.as_str()) + .collect::>(), + ["Edge Default", "Edge Profile 2", "Edge user-work"] + ); +} + +#[test] +fn indexed_db_keeps_only_allowed_origin_databases_that_are_directories() { + let root = tempfile::tempdir().unwrap(); + let allowed = [ + "https_platform.minimax.io_0", + "https_www.minimax.io_0", + "https_minimax.io_0", + "https_platform.minimaxi.com_0", + "https_www.minimaxi.com_0", + "https_minimaxi.com_0", + ] + .map(|origin| format!("{origin}.indexeddb.leveldb")); + let excluded = [ + "https_other.example_0.indexeddb.leveldb", + "https_minimax.io.evil_0.indexeddb.leveldb", + "https_minimax.io_0.indexeddb.blob", + ".https_minimax.io_0.indexeddb.leveldb", + ]; + for profile in ["Default", "Profile 2", "user-work", "Guest Profile"] { + for database in allowed.iter().map(String::as_str).chain(excluded) { + mkdir(root.path(), &format!("{profile}/IndexedDB/{database}")); + } + } + fs::write( + root.path() + .join("Default/IndexedDB/https_minimax.io_1.indexeddb.leveldb"), + b"x", + ) + .unwrap(); + + let found = candidates_in_user_data_dir(root.path(), "Fixture", indexed_db()); + + assert_eq!(found.len(), 3 * allowed.len()); + let mut expected_names: Vec<&str> = allowed.iter().map(String::as_str).collect(); + expected_names.sort_unstable(); + for (profile, group) in ["Default", "Profile 2", "user-work"] + .iter() + .zip(found.chunks(allowed.len())) + { + let names: Vec<_> = group + .iter() + .map(|candidate| candidate.path.file_name().unwrap().to_str().unwrap()) + .collect(); + assert_eq!(names, expected_names); + assert!( + group + .iter() + .all(|candidate| candidate.label == format!("Fixture {profile} (IndexedDB)")) + ); + } +} diff --git a/rust/src/providers/minimax/local_storage.rs b/rust/src/providers/minimax/local_storage.rs index a4a3d0617a..41ef8c8d29 100755 --- a/rust/src/providers/minimax/local_storage.rs +++ b/rust/src/providers/minimax/local_storage.rs @@ -1,10 +1,12 @@ //! MiniMax LocalStorage Importer //! -//! Extracts session data from browser localStorage for MiniMax platform. -//! Supports Chrome, Edge, Firefox, and Brave browsers. +//! Extracts session data from Chromium browser storage for the MiniMax platform. +//! 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::storage_discovery::{self, StorageCandidate, StorageKind}; use serde::{Deserialize, Serialize}; -use std::path::PathBuf; +use std::path::Path; /// Session data extracted from MiniMax localStorage #[derive(Debug, Clone, Serialize, Deserialize)] @@ -40,111 +42,59 @@ impl std::fmt::Display for ImportError { impl std::error::Error for ImportError {} +/// Origin prefixes of the MiniMax IndexedDB databases (`__`). +const INDEXED_DB_ORIGIN_PREFIXES: &[&str] = &[ + "https_platform.minimax.io_", + "https_www.minimax.io_", + "https_minimax.io_", + "https_platform.minimaxi.com_", + "https_minimaxi.com_", + "https_www.minimaxi.com_", +]; + +/// Stores tried in order; a later store is read only when earlier ones yield no session. +const STORAGE_ORDER: [StorageKind; 3] = [ + StorageKind::LocalStorage, + StorageKind::SessionStorage, + StorageKind::IndexedDb { + origin_prefixes: INDEXED_DB_ORIGIN_PREFIXES, + }, +]; + /// MiniMax localStorage importer pub struct MiniMaxLocalStorageImporter; impl MiniMaxLocalStorageImporter { - /// Import MiniMax session from browser localStorage + /// Import a MiniMax session from Chromium browser storage. + /// + /// Visits every installed Chromium-family browser and profile: Local Storage first, then + /// Session Storage, then MiniMax-origin IndexedDB when the earlier stores yield nothing. pub fn import_session() -> Result { - // Try browsers in order of preference - let browsers = Self::get_browser_paths(); - - for (browser_name, ls_path) in browsers { - if let Ok(session) = Self::extract_from_path(&ls_path, &browser_name) { - return Ok(session); - } - } - - Err(ImportError::BrowserNotFound) + Self::import_with(storage_discovery::discover) } - /// Get paths to browser localStorage databases - fn get_browser_paths() -> Vec<(String, PathBuf)> { - #[allow( - unused_mut, - reason = "mutability needed for conditional initialization that the compiler cannot prove" - )] - let mut paths = Vec::new(); - - #[cfg(target_os = "windows")] - { - if let Some(local_data) = dirs::data_local_dir() { - // Chrome - let chrome_path = local_data - .join("Google") - .join("Chrome") - .join("User Data") - .join("Default") - .join("Local Storage") - .join("leveldb"); - if chrome_path.exists() { - paths.push(("Chrome".to_string(), chrome_path)); - } - - // Edge - let edge_path = local_data - .join("Microsoft") - .join("Edge") - .join("User Data") - .join("Default") - .join("Local Storage") - .join("leveldb"); - if edge_path.exists() { - paths.push(("Edge".to_string(), edge_path)); - } - - // Brave - let brave_path = local_data - .join("BraveSoftware") - .join("Brave-Browser") - .join("User Data") - .join("Default") - .join("Local Storage") - .join("leveldb"); - if brave_path.exists() { - paths.push(("Brave".to_string(), brave_path)); - } - } - } - - #[cfg(target_os = "macos")] - { - if let Some(home) = dirs::home_dir() { - // Chrome - let chrome_path = home - .join("Library") - .join("Application Support") - .join("Google") - .join("Chrome") - .join("Default") - .join("Local Storage") - .join("leveldb"); - if chrome_path.exists() { - paths.push(("Chrome".to_string(), chrome_path)); - } - - // Edge - let edge_path = home - .join("Library") - .join("Application Support") - .join("Microsoft Edge") - .join("Default") - .join("Local Storage") - .join("leveldb"); - if edge_path.exists() { - paths.push(("Edge".to_string(), edge_path)); + fn import_with( + discover: impl Fn(StorageKind) -> Vec, + ) -> Result { + let mut found_any_store = false; + 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); } } } - paths + Err(if found_any_store { + ImportError::StorageNotFound + } else { + ImportError::BrowserNotFound + }) } - /// Extract session from a localStorage path - fn extract_from_path( - path: &PathBuf, - browser_name: &str, - ) -> Result { + /// 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()))?; @@ -168,7 +118,7 @@ impl MiniMaxLocalStorageImporter { } match minimax_data { - Some(json) => Self::parse_session_from_json(&json, browser_name), + Some(json) => Self::parse_session_from_json(&json, source_label), None => Err(ImportError::StorageNotFound), } } @@ -230,7 +180,7 @@ impl MiniMaxLocalStorageImporter { /// Parse session from extracted JSON fn parse_session_from_json( json: &serde_json::Value, - browser_name: &str, + source_label: &str, ) -> Result { let access_token = json .get("access_token") @@ -290,7 +240,7 @@ impl MiniMaxLocalStorageImporter { email, phone, plan_type, - source_label: format!("{} localStorage", browser_name), + source_label: source_label.to_string(), }) } } @@ -325,4 +275,101 @@ mod tests { let result = MiniMaxLocalStorageImporter::parse_session_from_json(&json, "Chrome"); assert!(result.is_err()); } + + fn candidate(label: &str, path: &Path) -> StorageCandidate { + StorageCandidate { + label: label.to_string(), + path: path.to_path_buf(), + } + } + + fn write_store(dir: &Path, token: Option<&str>) -> std::path::PathBuf { + std::fs::create_dir_all(dir).unwrap(); + let body = match token { + Some(token) => { + format!("minimax_user{{\"access_token\":\"{token}\",\"user_id\":\"1\"}}") + } + None => "unrelated".to_string(), + }; + std::fs::write(dir.join("000003.log"), body).unwrap(); + dir.to_path_buf() + } + + fn discover_from( + local: Vec, + session: Vec, + indexed: Vec, + ) -> impl Fn(StorageKind) -> Vec { + move |kind| match kind { + StorageKind::LocalStorage => local.clone(), + StorageKind::SessionStorage => session.clone(), + StorageKind::IndexedDb { origin_prefixes } => { + assert_eq!(origin_prefixes, INDEXED_DB_ORIGIN_PREFIXES); + indexed.clone() + } + } + } + + #[test] + fn import_prefers_local_storage_over_later_stores() { + let root = tempfile::tempdir().unwrap(); + let local = write_store(&root.path().join("local"), Some("from-local")); + let session = write_store(&root.path().join("session"), Some("from-session")); + + let found = MiniMaxLocalStorageImporter::import_with(discover_from( + vec![candidate("Chrome Default", &local)], + vec![candidate("Chrome Default (Session Storage)", &session)], + Vec::new(), + )) + .unwrap(); + + assert_eq!(found.access_token.as_deref(), Some("from-local")); + assert_eq!(found.source_label, "Chrome Default"); + } + + #[test] + fn import_falls_back_to_session_then_indexed_db_when_earlier_stores_are_empty() { + let root = tempfile::tempdir().unwrap(); + let local = write_store(&root.path().join("local"), None); + let session = write_store(&root.path().join("session"), None); + let indexed = write_store(&root.path().join("indexed"), Some("from-indexed")); + + let found = MiniMaxLocalStorageImporter::import_with(discover_from( + vec![candidate("Edge Default", &local)], + vec![candidate("Edge Default (Session Storage)", &session)], + vec![candidate("Edge Default (IndexedDB)", &indexed)], + )) + .unwrap(); + assert_eq!(found.access_token.as_deref(), Some("from-indexed")); + assert_eq!(found.source_label, "Edge Default (IndexedDB)"); + + let session_token = write_store(&root.path().join("session2"), Some("from-session")); + let found = MiniMaxLocalStorageImporter::import_with(discover_from( + vec![candidate("Edge Default", &local)], + vec![candidate("Edge Default (Session Storage)", &session_token)], + vec![candidate("Edge Default (IndexedDB)", &indexed)], + )) + .unwrap(); + assert_eq!(found.access_token.as_deref(), Some("from-session")); + } + + #[test] + fn import_distinguishes_no_browser_from_no_session() { + let root = tempfile::tempdir().unwrap(); + let empty = write_store(&root.path().join("empty"), None); + + let none = MiniMaxLocalStorageImporter::import_with(discover_from( + Vec::new(), + Vec::new(), + Vec::new(), + )); + assert!(matches!(none, Err(ImportError::BrowserNotFound))); + + let no_session = MiniMaxLocalStorageImporter::import_with(discover_from( + vec![candidate("Brave Default", &empty)], + Vec::new(), + Vec::new(), + )); + assert!(matches!(no_session, Err(ImportError::StorageNotFound))); + } } From ed916e3b5ec1afbc5008082350d9823193c92e3c Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:13:20 +0700 Subject: [PATCH 118/407] Port upstream 0.67.0: read Grok bearer credits from the CLI credits proxy --- rust/src/providers/grok/credits_proxy.rs | 250 ++++++++++++ .../src/providers/grok/credits_proxy_tests.rs | 355 ++++++++++++++++++ rust/src/providers/grok/mod.rs | 34 +- rust/src/providers/grok/tests.rs | 14 +- 4 files changed, 644 insertions(+), 9 deletions(-) create mode 100644 rust/src/providers/grok/credits_proxy.rs create mode 100644 rust/src/providers/grok/credits_proxy_tests.rs diff --git a/rust/src/providers/grok/credits_proxy.rs b/rust/src/providers/grok/credits_proxy.rs new file mode 100644 index 0000000000..dd89e57668 --- /dev/null +++ b/rust/src/providers/grok/credits_proxy.rs @@ -0,0 +1,250 @@ +//! Grok CLI credits-proxy source (`cli-chat-proxy.grok.com/v1/billing?format=credits`). +//! +//! The grok.com gRPC-web billing endpoint requires a browser-held Web Key +//! Exchange keypair, so bearer logins read credits from the CLI proxy first +//! and only fall back to gRPC-web when the proxy is unavailable. Wire shape +//! and precedence follow upstream `GrokCreditsProxyFetcher` (v0.67.0). + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use reqwest::StatusCode; +use serde::Deserialize; + +use crate::core::ProviderError; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +use super::billing::GrokBillingSnapshot; +use super::{GrokCredentials, GrokProvider, grok_plan_display_name}; + +pub(super) const CREDITS_PROXY_ENDPOINT: &str = + "https://cli-chat-proxy.grok.com/v1/billing?format=credits"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +/// Budget for the grok.com retry that supplies a percent a period-only +/// credits answer lacks. Period-only payloads recur on every refresh, so a +/// grok.com outage must not delay the credits answer already in hand. +const PERCENT_RETRY_TIMEOUT: Duration = Duration::from_secs(6); +const MAX_RESPONSE_BYTES: usize = 256 * 1024; + +/// Bearer billing result plus the plan name the credits payload carried. +#[derive(Debug, Clone)] +pub(super) struct BearerBilling { + pub(super) billing: GrokBillingSnapshot, + pub(super) subscription_tier: Option, +} + +#[derive(Deserialize)] +struct CreditsResponse { + config: Option, + #[serde(rename = "subscriptionTier")] + subscription_tier: Option, +} + +#[derive(Deserialize)] +struct CreditsConfig { + #[serde(rename = "creditUsagePercent")] + credit_usage_percent: Option, + #[serde(rename = "currentPeriod")] + current_period: Option, + #[serde(rename = "billingPeriodStart")] + billing_period_start: Option, + #[serde(rename = "billingPeriodEnd")] + billing_period_end: Option, + #[serde(rename = "onDemandCap")] + on_demand_cap: Option, + #[serde(rename = "onDemandUsed")] + on_demand_used: Option, + #[serde(rename = "subscriptionTier")] + subscription_tier: Option, +} + +#[derive(Deserialize)] +struct CurrentPeriod { + start: Option, + end: Option, +} + +/// The proxy reports amounts as `{ "val": }`; fractional values are +/// accepted so an unusual cap/used shape cannot fail an otherwise valid answer. +#[derive(Deserialize)] +struct CreditsAmount { + val: Option, +} + +fn parse_timestamp(raw: Option<&str>) -> Option> { + DateTime::parse_from_rfc3339(raw?) + .ok() + .map(|parsed| parsed.with_timezone(&Utc)) +} + +/// Full billing-cycle length; `None` unless the start precedes both the end +/// and `now`, so a reset timestamp alone never invents a cadence. +fn window_minutes( + start: Option<&str>, + end: Option>, + now: DateTime, +) -> Option { + let start = parse_timestamp(start)?; + let end = end?; + if end <= start || start > now { + return None; + } + u32::try_from((end - start).num_minutes()) + .ok() + .filter(|minutes| *minutes > 0) +} + +pub(super) fn parse_credits_response( + data: &[u8], + now: DateTime, +) -> Result { + let malformed = || ProviderError::Parse("Grok credits response is malformed".to_string()); + let response: CreditsResponse = serde_json::from_slice(data).map_err(|_| malformed())?; + let config = response.config.ok_or_else(malformed)?; + + let subscription_tier = grok_plan_display_name( + config + .subscription_tier + .or(response.subscription_tier) + .as_deref(), + ); + let current_end = parse_timestamp( + config + .current_period + .as_ref() + .and_then(|period| period.end.as_deref()), + ); + let resets_at = current_end.or_else(|| parse_timestamp(config.billing_period_end.as_deref())); + // Match the start to the selected end; never combine different billing periods. + let period_start = if current_end.is_some() { + config + .current_period + .as_ref() + .and_then(|period| period.start.as_deref()) + } else { + config.billing_period_start.as_deref() + }; + let window_minutes = window_minutes(period_start, resets_at, now); + + let used_percent = config + .credit_usage_percent + .filter(|percent| percent.is_finite()) + .or_else(|| { + let cap = config.on_demand_cap?.val.filter(|cap| *cap > 0.0)?; + Some(config.on_demand_used?.val? / cap * 100.0) + }) + .filter(|percent| percent.is_finite()) + .map(|percent| percent.clamp(0.0, 100.0)); + if used_percent.is_none() && resets_at.is_none() { + return Err(malformed()); + } + Ok(BearerBilling { + billing: GrokBillingSnapshot { + used_percent, + // A percent taken from this payload is published on the wire; a + // period-only answer leaves usage unknown rather than zero. + used_percent_is_wire_published: used_percent.is_some(), + used_percent_is_implicit_zero: false, + resets_at, + window_minutes, + }, + subscription_tier, + }) +} + +/// Keep the proxy's authoritative period and plan, adopting the grok.com +/// reading only when it actually published a percent (or an implicit zero). +fn adopt_grpc_percent(proxy: BearerBilling, grpc: GrokBillingSnapshot) -> BearerBilling { + let publishes_percent = grpc.used_percent.is_some() + && (grpc.used_percent_is_wire_published || grpc.used_percent_is_implicit_zero); + if !publishes_percent { + return proxy; + } + BearerBilling { + billing: GrokBillingSnapshot { + used_percent: grpc.used_percent, + used_percent_is_wire_published: grpc.used_percent_is_wire_published, + used_percent_is_implicit_zero: grpc.used_percent_is_implicit_zero, + ..proxy.billing + }, + subscription_tier: proxy.subscription_tier, + } +} + +impl GrokProvider { + /// Bearer billing: credits proxy first, then the grok.com gRPC-web path. + pub(super) async fn fetch_bearer_billing( + &self, + credentials: &GrokCredentials, + ) -> Result { + // Expired tokens are never sent to either endpoint. + if credentials.is_expired(Utc::now()) { + return Err(ProviderError::AuthRequired); + } + let authorization = format!("Bearer {}", credentials.access_token); + let proxy = match self.fetch_credits_proxy(&authorization).await { + Ok(proxy) => proxy, + Err(error) => { + tracing::debug!("Grok credits proxy failed, trying grok.com billing: {error}"); + let billing = self.fetch_billing(Some(authorization), None).await?; + return Ok(BearerBilling { + billing, + subscription_tier: None, + }); + } + }; + if proxy.billing.used_percent.is_some() { + return Ok(proxy); + } + match tokio::time::timeout( + PERCENT_RETRY_TIMEOUT, + self.fetch_billing(Some(authorization), None), + ) + .await + { + Ok(Ok(grpc)) => Ok(adopt_grpc_percent(proxy, grpc)), + Ok(Err(error)) => { + tracing::debug!("Grok grok.com percent retry failed: {error}"); + Ok(proxy) + } + Err(_) => { + tracing::debug!("Grok grok.com percent retry timed out"); + Ok(proxy) + } + } + } + + async fn fetch_credits_proxy( + &self, + authorization: &str, + ) -> Result { + let response = self + .client + .get(&self.credits_proxy_endpoint) + .timeout(REQUEST_TIMEOUT) + .header("Authorization", authorization) + .header("x-xai-token-auth", "xai-grok-cli") + .header("Accept", "application/json") + .header("User-Agent", "CodexBar") + .send() + .await?; + let status = response.status(); + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + if status != StatusCode::OK { + return Err(ProviderError::Other(format!( + "Grok credits proxy returned status {status}" + ))); + } + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => { + ProviderError::Parse("Grok credits response is too large".to_string()) + } + BoundedBodyError::Read(error) => ProviderError::Network(error), + })?; + parse_credits_response(&body, Utc::now()) + } +} diff --git a/rust/src/providers/grok/credits_proxy_tests.rs b/rust/src/providers/grok/credits_proxy_tests.rs new file mode 100644 index 0000000000..3af65e187c --- /dev/null +++ b/rust/src/providers/grok/credits_proxy_tests.rs @@ -0,0 +1,355 @@ +use super::credits_proxy::{BearerBilling, parse_credits_response}; +use super::tests::billing_response_with_percent; +use super::*; + +/// Upstream live capture shape (`GrokCreditsProxyFetcherTests`), including +/// fields the port does not read. +const WEEKLY_CREDITS_BODY: &str = r#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-09-20T18:42:45.537749+00:00","end":"2026-09-27T18:42:45.537749+00:00"},"creditUsagePercent":1.0,"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"productUsage":[{"product":"GrokBuild","usagePercent":1.0}],"isUnifiedBillingUser":true,"prepaidBalance":{"val":0},"topUpMethod":"TOP_UP_METHOD_SAVED_PAYMENT_METHOD","billingPeriodStart":"2026-09-20T18:42:45.537749+00:00","billingPeriodEnd":"2026-09-27T18:42:45.537749+00:00"}}"#; +const PERIOD_ONLY_BODY: &str = r#"{"config":{"currentPeriod":{"start":"2026-08-06T00:00:00Z","end":"2026-08-13T00:00:00Z"}},"subscriptionTier":"SUPERGROK_HEAVY"}"#; + +fn at(raw: &str) -> DateTime { + DateTime::parse_from_rfc3339(raw) + .unwrap() + .with_timezone(&Utc) +} + +fn parse(body: &str) -> BearerBilling { + parse_credits_response(body.as_bytes(), at("2026-08-12T00:00:00Z")).unwrap() +} + +fn provider_for(server: &mockito::ServerGuard) -> GrokProvider { + GrokProvider::new() + .with_billing_endpoint_for_tests(format!("{}/billing", server.url())) + .with_credits_proxy_endpoint_for_tests(format!("{}/credits", server.url())) +} + +#[test] +fn parses_the_live_weekly_capture() { + let parsed = + parse_credits_response(WEEKLY_CREDITS_BODY.as_bytes(), at("2026-09-25T00:00:00Z")).unwrap(); + assert_eq!(parsed.billing.used_percent, Some(1.0)); + assert!(parsed.billing.used_percent_is_wire_published); + assert_eq!( + parsed.billing.resets_at, + Some(at("2026-09-27T18:42:45.537749+00:00")) + ); + assert_eq!( + parsed.billing.window_minutes, + Some(crate::core::WEEKLY_WINDOW_MINUTES) + ); + assert_eq!(parsed.subscription_tier, None); +} + +#[test] +fn measures_only_matching_valid_period_bounds() { + let cases: [(&str, Option); 10] = [ + ( + r#""billingPeriodStart":"2026-08-06T00:00:00Z","billingPeriodEnd":"2026-08-13T00:00:00Z""#, + Some(10080), + ), + ( + r#""currentPeriod":{"start":"2026-07-13T00:00:00Z","end":"2026-08-13T00:00:00Z"}"#, + Some(44640), + ), + (r#""currentPeriod":{"end":"2026-08-13T00:00:00Z"}"#, None), + ( + r#""currentPeriod":{"end":"2026-08-13T00:00:00Z"},"billingPeriodStart":"2026-07-13T00:00:00Z","billingPeriodEnd":"2026-08-14T00:00:00Z""#, + None, + ), + ( + r#""currentPeriod":{"start":"invalid","end":"2026-08-13T00:00:00Z"}"#, + None, + ), + ( + r#""currentPeriod":{"start":"2026-08-14T00:00:00Z","end":"2026-08-21T00:00:00Z"}"#, + None, + ), + ( + r#""currentPeriod":{"start":"2026-08-11T00:00:00Z","end":"2026-08-10T00:00:00Z"}"#, + None, + ), + ( + r#""currentPeriod":{"start":"2026-08-11T00:00:00Z","end":"2026-08-11T00:00:00Z"}"#, + None, + ), + ( + r#""currentPeriod":{"start":"2026-08-11T00:00:00Z","end":"2026-08-11T00:00:30Z"}"#, + None, + ), + ( + r#""currentPeriod":{"start":"2026-07-01T00:00:00Z","end":"invalid"},"billingPeriodStart":"2026-08-06T00:00:00Z","billingPeriodEnd":"2026-08-13T00:00:00Z""#, + Some(10080), + ), + ]; + for (period, expected_minutes) in cases { + let parsed = parse(&format!( + r#"{{"config":{{"creditUsagePercent":90,{period}}}}}"# + )); + assert_eq!(parsed.billing.used_percent, Some(90.0), "{period}"); + assert_eq!(parsed.billing.window_minutes, expected_minutes, "{period}"); + } +} + +#[test] +fn derives_percent_from_on_demand_cap_and_usage() { + let parsed = parse(r#"{"config":{"onDemandCap":{"val":1000.0},"onDemandUsed":{"val":250.5}}}"#); + assert_eq!(parsed.billing.used_percent, Some(25.05)); + assert!(parsed.billing.used_percent_is_wire_published); + assert_eq!(parsed.billing.resets_at, None); +} + +#[test] +fn clamps_out_of_range_credit_usage() { + let over = parse( + r#"{"config":{"creditUsagePercent":104.2,"billingPeriodEnd":"2026-08-13T00:00:00Z"}}"#, + ); + let under = parse(r#"{"config":{"creditUsagePercent":-3.5}}"#); + assert_eq!(over.billing.used_percent, Some(100.0)); + assert_eq!(over.billing.resets_at, Some(at("2026-08-13T00:00:00Z"))); + assert_eq!(under.billing.used_percent, Some(0.0)); +} + +#[test] +fn period_without_usage_is_unknown_not_zero() { + let parsed = parse( + r#"{"config":{"currentPeriod":{"end":"2026-08-13T00:00:00.123Z"},"billingPeriodEnd":"2026-08-14T00:00:00Z"}}"#, + ); + assert_eq!(parsed.billing.used_percent, None); + assert!(!parsed.billing.used_percent_is_wire_published); + assert_eq!( + parsed.billing.resets_at, + Some(at("2026-08-13T00:00:00.123Z")) + ); + assert_eq!(parsed.subscription_tier, None); +} + +#[test] +fn billing_period_end_backs_up_a_missing_current_period_end() { + let parsed = parse( + r#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY"},"billingPeriodEnd":"2026-08-13T00:00:00Z"}}"#, + ); + assert_eq!(parsed.billing.used_percent, None); + assert_eq!(parsed.billing.resets_at, Some(at("2026-08-13T00:00:00Z"))); +} + +#[test] +fn subscription_tier_prefers_config_over_the_envelope() { + let parsed = parse( + r#"{"config":{"creditUsagePercent":8,"billingPeriodEnd":"2026-08-13T00:00:00Z","subscriptionTier":"SuperGrok Heavy"},"subscriptionTier":"SuperGrok"}"#, + ); + assert_eq!(parsed.subscription_tier.as_deref(), Some("SuperGrok Heavy")); + let envelope = parse( + r#"{"config":{"creditUsagePercent":57,"billingPeriodEnd":"2026-08-13T00:00:00Z"},"subscriptionTier":"SUPERGROK"}"#, + ); + assert_eq!(envelope.subscription_tier.as_deref(), Some("SuperGrok")); + let heavy = parse( + r#"{"config":{"subscriptionTier":"SUPERGROK_HEAVY","currentPeriod":{"end":"2026-08-13T00:00:00Z"}}}"#, + ); + assert_eq!(heavy.billing.used_percent, None); + assert_eq!(heavy.subscription_tier.as_deref(), Some("SuperGrok Heavy")); +} + +#[test] +fn rejects_responses_with_neither_usage_nor_period() { + let now = at("2026-08-12T00:00:00Z"); + for body in [ + r#"{"config":{}}"#, + r#"{"config":{"onDemandCap":{"val":0}},"subscriptionTier":"supergrok_heavy"}"#, + r#"{"subscriptionTier":"SuperGrok"}"#, + r#"{"config":{"creditUsagePercent":"12"}}"#, + "not json", + ] { + assert!( + matches!( + parse_credits_response(body.as_bytes(), now), + Err(ProviderError::Parse(_)) + ), + "{body}" + ); + } +} + +#[tokio::test] +async fn bearer_billing_sends_the_cli_proxy_request_and_skips_grpc() { + let mut server = mockito::Server::new_async().await; + let proxy = server + .mock("GET", "/credits") + .match_header("authorization", "Bearer token-123") + .match_header("x-xai-token-auth", "xai-grok-cli") + .match_header("accept", "application/json") + .match_header("user-agent", "CodexBar") + .with_status(200) + .with_body( + r#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"start":"2020-01-01T00:00:00Z","end":"2020-01-08T00:00:00Z"}}}"#, + ) + .expect(1) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .with_status(200) + .expect(0) + .create_async() + .await; + + let bearer = provider_for(&server) + .fetch_bearer_billing(&GrokCredentials::from_bearer("token-123")) + .await + .unwrap(); + + proxy.assert_async().await; + grpc.assert_async().await; + assert_eq!(bearer.billing.used_percent, Some(12.5)); + assert_eq!( + bearer.billing.window_minutes, + Some(crate::core::WEEKLY_WINDOW_MINUTES) + ); +} + +#[tokio::test] +async fn proxy_failure_falls_back_to_grpc_billing() { + for status in [401, 500] { + let mut server = mockito::Server::new_async().await; + let proxy = server + .mock("GET", "/credits") + .with_status(status) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .match_header("authorization", "Bearer token-123") + .with_status(200) + .with_body(billing_response_with_percent(37.0)) + .create_async() + .await; + + let bearer = provider_for(&server) + .fetch_bearer_billing(&GrokCredentials::from_bearer("token-123")) + .await + .unwrap(); + + proxy.assert_async().await; + grpc.assert_async().await; + assert_eq!(bearer.billing.used_percent, Some(37.0), "status {status}"); + assert_eq!(bearer.subscription_tier, None); + } +} + +#[tokio::test] +async fn period_only_answer_adopts_the_grpc_percent_and_keeps_proxy_metadata() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body(PERIOD_ONLY_BODY) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .with_status(200) + .with_body(billing_response_with_percent(90.0)) + .expect(1) + .create_async() + .await; + + let bearer = provider_for(&server) + .fetch_bearer_billing(&GrokCredentials::from_bearer("token-123")) + .await + .unwrap(); + + grpc.assert_async().await; + assert_eq!(bearer.billing.used_percent, Some(90.0)); + assert!(bearer.billing.used_percent_is_wire_published); + assert_eq!(bearer.billing.resets_at, Some(at("2026-08-13T00:00:00Z"))); + assert_eq!(bearer.subscription_tier.as_deref(), Some("SuperGrok Heavy")); +} + +#[tokio::test] +async fn period_only_answer_stays_unknown_when_grpc_has_no_percent_or_fails() { + // An empty gRPC-web frame carries no percent; a 401 is a failed retry. + for (status, body) in [(200, vec![0, 0, 0, 0, 0]), (401, Vec::new())] { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body(PERIOD_ONLY_BODY) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .with_status(status) + .with_body(body) + .create_async() + .await; + + let bearer = provider_for(&server) + .fetch_bearer_billing(&GrokCredentials::from_bearer("token-123")) + .await + .unwrap(); + + grpc.assert_async().await; + assert_eq!(bearer.billing.used_percent, None, "status {status}"); + assert_eq!(bearer.billing.resets_at, Some(at("2026-08-13T00:00:00Z"))); + assert_eq!(bearer.subscription_tier.as_deref(), Some("SuperGrok Heavy")); + } +} + +#[tokio::test] +async fn expired_credentials_are_never_sent() { + let mut server = mockito::Server::new_async().await; + let proxy = server + .mock("GET", "/credits") + .expect(0) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .expect(0) + .create_async() + .await; + let mut credentials = GrokCredentials::from_bearer("stale-token"); + credentials.expires_at = Some(Utc::now() - chrono::Duration::minutes(1)); + + let error = provider_for(&server) + .fetch_bearer_billing(&credentials) + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::AuthRequired)); + proxy.assert_async().await; + grpc.assert_async().await; +} + +#[tokio::test] +async fn credits_plan_labels_oauth_results_without_settings_lookup() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body( + r#"{"config":{"creditUsagePercent":40,"billingPeriodEnd":"2026-08-13T00:00:00Z"},"subscriptionTier":"SUPERGROK_HEAVY"}"#, + ) + .create_async() + .await; + let context = FetchContext { + include_credits: false, + ..FetchContext::default() + }; + + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context, + ) + .await + .unwrap(); + + assert_eq!(result.source_label, "grok-oauth"); + assert_eq!(result.usage.primary.used_percent, 40.0); + assert!(!result.usage.primary.is_informational); + assert_eq!( + result.usage.login_method.as_deref(), + Some("SuperGrok Heavy") + ); +} diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 7451f281ef..0db1fdcc7a 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -1,10 +1,12 @@ //! Grok provider implementation. //! -//! Uses the grok.com billing gRPC-web endpoint via either browser cookies or -//! `~/.grok/auth.json` produced by `grok login`. +//! Bearer logins (`~/.grok/auth.json` from `grok login`) read credits from the +//! Grok CLI credits proxy first, then fall back to the grok.com billing +//! gRPC-web endpoint. Browser cookies use the gRPC-web endpoint only. pub mod accounts; mod billing; +mod credits_proxy; pub mod local_sessions; use async_trait::async_trait; @@ -24,6 +26,7 @@ use crate::core::{ use self::accounts::{GrokAuthKind, ParsedGrokAuthFile}; use self::billing::GrokBillingSnapshot; +use self::credits_proxy::{BearerBilling, CREDITS_PROXY_ENDPOINT}; const BILLING_ENDPOINT: &str = "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"; const BILLING_REQUEST_BODY: [u8; 7] = [0, 0, 0, 0, 2, 0x08, 0]; @@ -37,6 +40,7 @@ pub struct GrokProvider { metadata: ProviderMetadata, client: Client, billing_endpoint: String, + credits_proxy_endpoint: String, } impl GrokProvider { @@ -60,6 +64,7 @@ impl GrokProvider { .build() .unwrap_or_else(|_| Client::new()), billing_endpoint: BILLING_ENDPOINT.to_string(), + credits_proxy_endpoint: CREDITS_PROXY_ENDPOINT.to_string(), } } @@ -83,6 +88,12 @@ impl GrokProvider { self } + #[cfg(test)] + fn with_credits_proxy_endpoint_for_tests(mut self, endpoint: String) -> Self { + self.credits_proxy_endpoint = endpoint; + self + } + fn load_credentials(kind: GrokAuthKind) -> Result { let path = Self::auth_file_path() .ok_or_else(|| ProviderError::NotInstalled("Grok auth path not found".to_string()))?; @@ -123,11 +134,11 @@ impl GrokProvider { None, self.client.clone(), ); - let billing = match self - .fetch_billing(Some(format!("Bearer {}", credentials.access_token)), None) - .await - { - Ok(billing) => billing, + let BearerBilling { + billing, + subscription_tier, + } = match self.fetch_bearer_billing(credentials).await { + Ok(bearer) => bearer, Err(error) => { reset_lookup.abort(); return Err(error); @@ -138,6 +149,7 @@ impl GrokProvider { } else { None } + .or(subscription_tier) .or_else(|| credentials.login_method()); let reset_credits = reset_lookup .join(ctx.requires_optional_usage_completeness) @@ -670,6 +682,10 @@ impl GrokCredentials { }) } + fn is_expired(&self, now: DateTime) -> bool { + self.expires_at.is_some_and(|expires_at| expires_at <= now) + } + fn login_method(&self) -> Option { match self.auth_mode.as_deref().map(str::to_lowercase).as_deref() { Some("oidc") => Some("SuperGrok".to_string()), @@ -817,3 +833,7 @@ fn cookie_refresh_action( #[cfg(test)] #[path = "tests.rs"] mod tests; + +#[cfg(test)] +#[path = "credits_proxy_tests.rs"] +mod credits_proxy_tests; diff --git a/rust/src/providers/grok/tests.rs b/rust/src/providers/grok/tests.rs index 586c20c637..585c0255a8 100644 --- a/rust/src/providers/grok/tests.rs +++ b/rust/src/providers/grok/tests.rs @@ -1,7 +1,7 @@ use super::*; use chrono::TimeZone; -fn billing_response_with_percent(percent: f32) -> Vec { +pub(super) fn billing_response_with_percent(percent: f32) -> Vec { let mut payload = vec![0x0a, 0x05, 0x0d]; payload.extend(percent.to_le_bytes()); @@ -57,7 +57,16 @@ async fn oauth_billing_auth_failure_falls_back_to_configured_local_token() { .with_body(billing_response_with_percent(42.0)) .create_async() .await; - let provider = GrokProvider::new().with_billing_endpoint_for_tests(endpoint); + // The credits proxy is unavailable, so both tokens use the gRPC-web path. + let proxy = server + .mock("GET", "/credits") + .with_status(500) + .expect_at_least(1) + .create_async() + .await; + let provider = GrokProvider::new() + .with_billing_endpoint_for_tests(endpoint) + .with_credits_proxy_endpoint_for_tests(format!("{}/credits", server.url())); let credentials = GrokCredentials::from_bearer("expired-token"); let context = FetchContext { include_credits: false, @@ -70,6 +79,7 @@ async fn oauth_billing_auth_failure_falls_back_to_configured_local_token() { .await .unwrap(); + proxy.assert_async().await; rejected.assert_async().await; fallback.assert_async().await; assert_eq!(result.source_label, "grok-oauth"); From 67c1fe9149850d2e1050acf08a6fe3ddc778f376 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:16:02 +0700 Subject: [PATCH 119/407] Port upstream 0.67.0: add portable preferences export and import --- .../src-tauri/src/commands/mod.rs | 2 + .../src/commands/preferences_transfer.rs | 75 +++ apps/desktop-tauri/src-tauri/src/main.rs | 2 + apps/desktop-tauri/src/i18n/keys.ts | 6 + apps/desktop-tauri/src/lib/tauri.ts | 10 + .../surfaces/settings/tabs/AdvancedTab.tsx | 3 + .../tabs/PreferencesTransferSection.test.tsx | 123 +++++ .../tabs/PreferencesTransferSection.tsx | 106 ++++ docs/CONFIGURATION.md | 11 + rust/src/cli/config.rs | 72 ++- rust/src/locale.rs | 6 + rust/src/locale/en-US.ftl | 6 + rust/src/settings.rs | 2 + rust/src/settings/preferences_document.rs | 364 ++++++++++++++ .../settings/preferences_document/tests.rs | 474 ++++++++++++++++++ 15 files changed, 1261 insertions(+), 1 deletion(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/preferences_transfer.rs create mode 100644 apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.test.tsx create mode 100644 apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.tsx create mode 100644 rust/src/settings/preferences_document.rs create mode 100644 rust/src/settings/preferences_document/tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index fdfaeda967..3867c9716c 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -41,6 +41,7 @@ mod credentials; mod diagnostics; mod grok_accounts; mod locale_cmd; +mod preferences_transfer; mod provider_detail; mod provider_refresh; mod provider_settings; @@ -63,6 +64,7 @@ pub use credentials::*; pub use diagnostics::*; pub use grok_accounts::*; pub use locale_cmd::*; +pub use preferences_transfer::*; pub use provider_detail::*; #[cfg(test)] pub(crate) use provider_refresh::is_provider_cache_fresh; diff --git a/apps/desktop-tauri/src-tauri/src/commands/preferences_transfer.rs b/apps/desktop-tauri/src-tauri/src/commands/preferences_transfer.rs new file mode 100644 index 0000000000..1d7694f833 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/preferences_transfer.rs @@ -0,0 +1,75 @@ +//! Export and import of the portable preferences document. +//! +//! The document format, allowlist and validation live in +//! `codexbar::settings::PreferencesDocument`; this module only adds the +//! desktop side effects that a settings save normally triggers. + +use std::path::Path; +use std::sync::Mutex; + +use codexbar::settings::{PreferencesDocument, Settings}; +use tauri::{Emitter, Manager}; + +use super::{SettingsSnapshot, language_label}; +use crate::events; +use crate::state::AppState; + +fn checked_path(path: &str) -> Result<&Path, String> { + let path = path.trim(); + if path.is_empty() { + return Err("Preferences path must not be empty".to_string()); + } + Ok(Path::new(path)) +} + +/// Write the current portable preferences to `path`; returns how many were written. +#[tauri::command] +pub fn export_preferences(path: String) -> Result { + let path = checked_path(&path)?; + let document = + PreferencesDocument::from_settings(&Settings::load()).map_err(|error| error.to_string())?; + document + .write_file(path) + .map_err(|error| error.to_string())?; + tracing::info!(count = document.len(), "exported portable preferences"); + Ok(document.len()) +} + +/// Apply the preferences file at `path` to this machine's settings. +/// +/// Nothing is saved when the file is rejected. On success the same live +/// updates as `update_settings` follow: locale, provider cache, float bar, +/// tray and dependent windows. +#[tauri::command] +pub async fn import_preferences( + app: tauri::AppHandle, + path: String, +) -> Result { + let path = checked_path(&path)?; + let document = PreferencesDocument::read_file(path).map_err(|error| error.to_string())?; + let mut settings = Settings::load(); + let previous_language = settings.ui_language; + let applied = document + .apply_to(&mut settings) + .map_err(|error| error.to_string())?; + settings.save().map_err(|error| error.to_string())?; + tracing::info!(applied, "imported portable preferences"); + + if settings.ui_language != previous_language { + let _ = app.emit(events::LOCALE_CHANGED, language_label(settings.ui_language)); + } + let enabled_ids = settings.get_enabled_provider_ids(); + let state = app.state::>(); + let _ = super::invalidate_provider_refresh_and_prune_disabled(&state, &enabled_ids); + crate::floatbar::notify_settings_changed(&app); + crate::floatbar::apply_state(&app, &settings); + crate::tray_bridge::rebuild_tray_menu(&app); + crate::tray_bridge::refresh_tray_presentation(&app); + events::emit_settings_changed(&app); + + let refresh_app = app.clone(); + tauri::async_runtime::spawn(async move { + let _ = super::do_refresh_providers(&refresh_app).await; + }); + Ok(SettingsSnapshot::from(settings)) +} diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 02f48672b6..8bb06fe1af 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -230,6 +230,8 @@ fn main() { commands::get_provider_local_usage_summary, commands::get_usage_spend_summary, commands::write_usage_spend_export, + commands::export_preferences, + commands::import_preferences, commands::get_spend_contract, commands::get_codex_workspaces_snapshot, commands::reorder_providers, diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..2bbb371f3d 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -530,6 +530,12 @@ export const ALL_LOCALE_KEYS = [ "CriticalUsageWarningHelper", "GlobalShortcutFieldLabel", "GlobalShortcutToggleHelper", + "SectionPreferencesTransfer", + "PreferencesTransferCaption", + "PreferencesExportButton", + "PreferencesImportButton", + "PreferencesExportSuccess", + "PreferencesImportSuccess", "ShortcutRecordButton", "ShortcutRecordingLabel", "ShortcutRecordingHint", diff --git a/apps/desktop-tauri/src/lib/tauri.ts b/apps/desktop-tauri/src/lib/tauri.ts index 6ccd63c0de..f53ae41f41 100644 --- a/apps/desktop-tauri/src/lib/tauri.ts +++ b/apps/desktop-tauri/src/lib/tauri.ts @@ -295,6 +295,16 @@ export function writeUsageSpendExport(path: string, payload: string): Promise("write_usage_spend_export", { path, payload }); } +/** Write portable preferences to `path`; resolves with how many were exported. */ +export function exportPreferences(path: string): Promise { + return invoke("export_preferences", { path }); +} + +/** Apply the preferences file at `path`; the shell validates before saving anything. */ +export function importPreferences(path: string): Promise { + return invoke("import_preferences", { path }); +} + export function getSpendContract( providerId: string, options?: { historyDays?: number; includeOpenCodex?: boolean }, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx index fcf98d09c3..be048a5dec 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx @@ -8,6 +8,7 @@ import { import { ShortcutCapture } from "../../../components/ShortcutCapture"; import { Field, Toggle } from "../../../components/FormControls"; import type { TabProps } from "../settingsTabs"; +import PreferencesTransferSection from "./PreferencesTransferSection"; function formatCodexSessionsDirs(paths: string[]): string { return paths.join("; "); @@ -372,6 +373,8 @@ export default function AdvancedTab({ settings, set, saving }: TabProps) { + + {/* ── Diagnostics ──────────────────────────────────────────── */}

diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.test.tsx new file mode 100644 index 0000000000..7e98389cc7 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.test.tsx @@ -0,0 +1,123 @@ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const mocks = vi.hoisted(() => ({ + open: vi.fn(), + save: vi.fn(), + exportPreferences: vi.fn(), + importPreferences: vi.fn(), +})); + +vi.mock("@tauri-apps/plugin-dialog", () => ({ + open: mocks.open, + save: mocks.save, +})); +vi.mock("../../../lib/tauri", () => ({ + exportPreferences: mocks.exportPreferences, + importPreferences: mocks.importPreferences, +})); +vi.mock("../../../hooks/useLocale", () => ({ + useLocale: () => ({ t: (key: string) => key }), +})); + +import PreferencesTransferSection from "./PreferencesTransferSection"; + +describe("PreferencesTransferSection", () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it("exports to the chosen path and reports success", async () => { + mocks.save.mockResolvedValue("C:\\temp\\prefs.json"); + mocks.exportPreferences.mockResolvedValue(12); + render(); + + fireEvent.click(screen.getByText("PreferencesExportButton")); + + await waitFor(() => + expect(mocks.exportPreferences).toHaveBeenCalledWith("C:\\temp\\prefs.json"), + ); + expect(await screen.findByRole("status")).toHaveTextContent( + "PreferencesExportSuccess", + ); + expect(mocks.save).toHaveBeenCalledWith( + expect.objectContaining({ defaultPath: "codexbar-preferences.json" }), + ); + }); + + it("does nothing when the save dialog is cancelled", async () => { + mocks.save.mockResolvedValue(null); + render(); + + fireEvent.click(screen.getByText("PreferencesExportButton")); + + await waitFor(() => expect(mocks.save).toHaveBeenCalled()); + expect(mocks.exportPreferences).not.toHaveBeenCalled(); + expect(screen.queryByRole("status")).toBeNull(); + expect(screen.queryByRole("alert")).toBeNull(); + }); + + it("imports the chosen file and reports success", async () => { + mocks.open.mockResolvedValue("C:\\temp\\prefs.json"); + mocks.importPreferences.mockResolvedValue({}); + render(); + + fireEvent.click(screen.getByText("PreferencesImportButton")); + + await waitFor(() => + expect(mocks.importPreferences).toHaveBeenCalledWith("C:\\temp\\prefs.json"), + ); + expect(await screen.findByRole("status")).toHaveTextContent( + "PreferencesImportSuccess", + ); + expect(mocks.open).toHaveBeenCalledWith( + expect.objectContaining({ multiple: false }), + ); + }); + + it("does nothing when the open dialog is cancelled", async () => { + mocks.open.mockResolvedValue(null); + render(); + + fireEvent.click(screen.getByText("PreferencesImportButton")); + + await waitFor(() => expect(mocks.open).toHaveBeenCalled()); + expect(mocks.importPreferences).not.toHaveBeenCalled(); + }); + + it("shows the shell's rejection and no success message", async () => { + mocks.open.mockResolvedValue("C:\\temp\\bad.json"); + mocks.importPreferences.mockRejectedValue( + "Invalid or non-portable preference: refresh_interval_secs", + ); + render(); + + fireEvent.click(screen.getByText("PreferencesImportButton")); + + expect(await screen.findByRole("alert")).toHaveTextContent( + "Invalid or non-portable preference: refresh_interval_secs", + ); + expect(screen.queryByRole("status")).toBeNull(); + }); + + it("disables both buttons while a transfer is running", async () => { + let finish: (value: string | null) => void = () => {}; + mocks.save.mockReturnValue( + new Promise((resolve) => { + finish = resolve; + }), + ); + render(); + + fireEvent.click(screen.getByText("PreferencesExportButton")); + + await waitFor(() => + expect(screen.getByText("PreferencesImportButton")).toBeDisabled(), + ); + expect(screen.getByText("PreferencesExportButton")).toBeDisabled(); + finish(null); + await waitFor(() => + expect(screen.getByText("PreferencesImportButton")).toBeEnabled(), + ); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.tsx new file mode 100644 index 0000000000..1d195291a7 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/PreferencesTransferSection.tsx @@ -0,0 +1,106 @@ +import { useCallback, useState } from "react"; +import { open, save } from "@tauri-apps/plugin-dialog"; +import { useLocale } from "../../../hooks/useLocale"; +import { exportPreferences, importPreferences } from "../../../lib/tauri"; + +const JSON_FILTERS = [{ name: "JSON", extensions: ["json"] }]; + +type TransferOutcome = + | { kind: "success"; message: string } + | { kind: "error"; message: string }; + +/** + * Export or import the portable preferences document. The shell validates the + * whole file before saving anything, so a rejected import changes nothing. + */ +export default function PreferencesTransferSection() { + const { t } = useLocale(); + const [busy, setBusy] = useState(false); + const [outcome, setOutcome] = useState(null); + + const run = useCallback( + async (action: () => Promise) => { + setBusy(true); + setOutcome(null); + try { + const message = await action(); + if (message) setOutcome({ kind: "success", message }); + } catch (cause: unknown) { + setOutcome({ + kind: "error", + message: cause instanceof Error ? cause.message : String(cause), + }); + } finally { + setBusy(false); + } + }, + [], + ); + + const onExport = useCallback( + () => + run(async () => { + const path = await save({ + defaultPath: "codexbar-preferences.json", + filters: JSON_FILTERS, + }); + if (!path) return null; + await exportPreferences(path); + return t("PreferencesExportSuccess"); + }), + [run, t], + ); + + const onImport = useCallback( + () => + run(async () => { + const path = await open({ multiple: false, filters: JSON_FILTERS }); + if (typeof path !== "string") return null; + await importPreferences(path); + return t("PreferencesImportSuccess"); + }), + [run, t], + ); + + return ( +
+

+ {t("SectionPreferencesTransfer")} +

+

+ {t("PreferencesTransferCaption")} +

+
+ + +
+ {outcome?.kind === "success" && ( +

+ {outcome.message} +

+ )} + {outcome?.kind === "error" && ( +

+ {outcome.message} +

+ )} +
+ ); +} diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 99e470487b..03b6f7221f 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -49,8 +49,19 @@ codexbar config path # API key via stdin (example) printf '%s' $env:OPENROUTER_API_KEY | codexbar config set-api-key -p openrouter --stdin + +# Portable preferences (also under Settings > Advanced) +codexbar config preferences export prefs.json +codexbar config preferences import prefs.json ``` +Portable preferences are a versioned JSON document (`{"version": 1, "preferences": {...}}`) +limited to display, refresh, notification, provider-order and float-bar choices. API keys, cookies, +token accounts, folders, SSH hosts, proxy settings and update settings are never exported, and +an import that contains any other key, or any invalid value, is rejected without changing settings. +The CLI import writes `settings.json` only; restart a running CodexBar to pick it up. The desktop +Import button applies the file live. + Notes: - `enable` / `disable` are **persistent** (same idea as upstream). diff --git a/rust/src/cli/config.rs b/rust/src/cli/config.rs index b49be7ae2c..d7bb8c1e41 100755 --- a/rust/src/cli/config.rs +++ b/rust/src/cli/config.rs @@ -7,7 +7,7 @@ use serde::de::DeserializeOwned; use std::path::Path; use crate::core::{ProviderId, TokenAccountStore, instantiate_provider}; -use crate::settings::{ApiKeys, ManualCookies, Settings}; +use crate::settings::{ApiKeys, ManualCookies, PreferencesDocument, Settings}; /// Arguments for the config command #[derive(Parser, Debug)] @@ -57,6 +57,25 @@ pub enum ConfigCommand { }, /// Show configuration file paths Path, + /// Export or import portable preferences (no secrets, no machine state) + Preferences { + #[command(subcommand)] + action: PreferencesAction, + }, +} + +#[derive(Subcommand, Debug)] +pub enum PreferencesAction { + /// Write portable preferences to a JSON file + Export { + /// Destination file + path: std::path::PathBuf, + }, + /// Apply a preferences file; restart a running CodexBar afterwards + Import { + /// Preferences file to read + path: std::path::PathBuf, + }, } /// Run the config command @@ -77,9 +96,36 @@ pub async fn run(args: ConfigArgs) -> anyhow::Result<()> { no_enable, } => set_api_key(&provider, api_key.as_deref(), stdin, !no_enable).await, ConfigCommand::Path => show_paths().await, + ConfigCommand::Preferences { action } => transfer_preferences(action), } } +/// Export or import the portable preferences document. +fn transfer_preferences(action: PreferencesAction) -> anyhow::Result<()> { + match action { + PreferencesAction::Export { path } => { + let document = PreferencesDocument::from_settings(&Settings::load())?; + document.write_file(&path)?; + println!( + "Config: exported {} preferences to {}", + document.len(), + path.display() + ); + } + PreferencesAction::Import { path } => { + let document = PreferencesDocument::read_file(&path)?; + let mut settings = Settings::load(); + let applied = document.apply_to(&mut settings)?; + settings.save()?; + println!( + "Config: imported {applied} preferences from {}. Restart CodexBar to apply them to a running app.", + path.display() + ); + } + } + Ok(()) +} + /// Validate configuration files async fn validate_config() -> anyhow::Result<()> { let mut report = ValidationReport::default(); @@ -484,6 +530,30 @@ mod tests { use crate::settings::ManualCookies; use serde_json::json; + #[test] + fn preferences_subcommands_parse_a_path() { + use super::{ConfigArgs, ConfigCommand, PreferencesAction}; + use clap::Parser; + + let export = ConfigArgs::try_parse_from(["config", "preferences", "export", "p.json"]) + .expect("export parses"); + assert!(matches!( + export.command, + ConfigCommand::Preferences { + action: PreferencesAction::Export { .. } + } + )); + let import = ConfigArgs::try_parse_from(["config", "preferences", "import", "p.json"]) + .expect("import parses"); + assert!(matches!( + import.command, + ConfigCommand::Preferences { + action: PreferencesAction::Import { .. } + } + )); + assert!(ConfigArgs::try_parse_from(["config", "preferences", "import"]).is_err()); + } + #[cfg(windows)] #[test] fn validates_dpapi_protected_manual_cookies_without_plaintext() { diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..70ee10bf53 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -794,6 +794,12 @@ locale_keys! { CriticalUsageWarningHelper, GlobalShortcutFieldLabel, GlobalShortcutToggleHelper, + SectionPreferencesTransfer, + PreferencesTransferCaption, + PreferencesExportButton, + PreferencesImportButton, + PreferencesExportSuccess, + PreferencesImportSuccess, ShortcutRecordButton, ShortcutRecordingLabel, ShortcutRecordingHint, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..fb58e2c59e 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -467,6 +467,12 @@ HighUsageWarningHelper = Show a warning when usage exceeds this percentage. CriticalUsageWarningHelper = Show a critical alert when usage exceeds this percentage. GlobalShortcutFieldLabel = Global shortcut GlobalShortcutToggleHelper = Key combination to toggle the tray panel. +SectionPreferencesTransfer = Portable preferences +PreferencesTransferCaption = Copy display, refresh, notification and provider choices to another PC. API keys, cookies, folders and hosts are never included. +PreferencesExportButton = Export preferences +PreferencesImportButton = Import preferences +PreferencesExportSuccess = Preferences exported. +PreferencesImportSuccess = Preferences imported. ShortcutRecordButton = Record ShortcutRecordingLabel = Recording… ShortcutRecordingHint = Press modifiers + a key. Esc cancels, Backspace clears. diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..4afc4ec2f0 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -27,6 +27,7 @@ pub const CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID: &str = "metric:extra-claude-routi mod api_keys; mod manual_cookies; +mod preferences_document; mod provider_workspace; mod raw; mod status; @@ -34,6 +35,7 @@ mod types; pub use api_keys::*; pub use manual_cookies::*; +pub use preferences_document::*; pub use provider_workspace::*; use raw::RawSettings; pub use status::*; diff --git a/rust/src/settings/preferences_document.rs b/rust/src/settings/preferences_document.rs new file mode 100644 index 0000000000..c11089bcc1 --- /dev/null +++ b/rust/src/settings/preferences_document.rs @@ -0,0 +1,364 @@ +//! Portable UI preferences document (upstream 0.67.0 "portable preferences"). +//! +//! A versioned JSON file, `{"version":1,"preferences":{...}}`, that moves the +//! look-and-behavior preferences between machines. It is an explicit +//! allowlist, independent of the on-disk `settings.json` shape: +//! +//! - Unknown keys, unsupported versions, wrong types and out-of-range values +//! are rejected before anything is applied. +//! - Keys missing from a document leave the receiving machine unchanged. +//! - `null` restores the documented default (`Settings::default()`). +//! +//! Secrets and machine-specific or consent settings are never portable. The +//! [`EXCLUDED_KEYS`] list names every other [`Settings`] field, and a test +//! fails when a new field is added without being classified, so a new field +//! is excluded until someone opts it in. + +use std::collections::{BTreeMap, HashSet}; +use std::io::Read; +use std::path::Path; + +use serde::de::DeserializeOwned; +use serde_json::{Map, Value}; + +use super::{MetricPreference, Settings}; +use crate::core::ProviderId; + +/// Current (and only) document version. +pub const PREFERENCES_DOCUMENT_VERSION: u64 = 1; + +/// Upper bound for a document read from disk. A real document is a few KiB. +pub const MAX_PREFERENCES_DOCUMENT_BYTES: u64 = 256 * 1024; + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum PreferencesError { + #[error("Preferences file is not a valid preferences document")] + Malformed, + #[error("Unsupported preferences version")] + UnsupportedVersion, + #[error("Invalid or non-portable preference: {0}")] + InvalidPreference(String), + #[error("Preferences file is larger than {} KiB", MAX_PREFERENCES_DOCUMENT_BYTES / 1024)] + TooLarge, + #[error("Could not read or write the preferences file: {0}")] + Io(String), + #[error("Could not apply preferences: {0}")] + Apply(String), +} + +type Validator = fn(&Value) -> bool; + +/// Portable keys, named after the `Settings` fields, with their validators. +const ALLOWED_PREFERENCES: &[(&str, Validator)] = &[ + ("refresh_interval_secs", refresh_interval_secs), + ("adaptive_refresh", typed::), + ("refresh_all_providers_on_menu_open", typed::), + ( + "low_power_mode_preference", + typed::, + ), + ("show_notifications", typed::), + ("sound_enabled", typed::), + ("high_usage_threshold", percent), + ("critical_usage_threshold", percent), + ("provider_usage_thresholds", provider_usage_thresholds), + ("predictive_pace_warning_enabled", typed::), + ("show_pace", typed::), + ("show_as_used", typed::), + ("reset_time_relative", typed::), + ("show_reset_when_exhausted", typed::), + ("hide_personal_info", typed::), + ("menu_bar_shows_highest_usage", typed::), + ("menu_bar_shows_percent", typed::), + ("tray_icon_mode", typed::), + ("switcher_shows_icons", typed::), + ("overview_layout", |v| one_of(v, &["compact", "detailed"])), + ("menu_bar_display_mode", |v| { + one_of(v, &["minimal", "compact", "detailed"]) + }), + ("provider_order", provider_id_list), + ("enabled_providers", provider_id_list), + ("provider_metrics", provider_metrics), + ("theme", typed::), + ("ui_language", typed::), + ("window_scale_percent", uint_in::<100, 250>), + ("tray_scale_percent", uint_in::<100, 200>), + ("float_bar_opacity", uint_in::<30, 100>), + ("float_bar_scale", uint_in::<75, 200>), + ("float_bar_orientation", |v| { + one_of(v, &["horizontal", "vertical"]) + }), + ("float_bar_style", |v| one_of(v, &["floating", "taskbar"])), + ("float_bar_dark_text", typed::), + ("float_bar_show_reset_inline", typed::), + ("float_bar_show_cost", typed::), +]; + +/// Every other `Settings` field. Secrets (`provider_configs` holds API +/// tokens and manual cookie headers, `http_proxy_password`), machine paths and +/// hosts, and consent or side-effect toggles must stay out of a portable file. +pub const EXCLUDED_KEYS: &[&str] = &[ + // Secrets and credential-bearing config. + "provider_configs", + "http_proxy_enabled", + "http_proxy_url", + "http_proxy_username", + "http_proxy_password", + // Machine-specific paths and hosts. + "notification_sound_paths", + "codex_custom_sessions_dirs", + "agent_session_ssh_hosts", + // Consent, autostart and other side-effect toggles. + "start_at_login", + "start_minimized", + "hooks_enabled", + "global_shortcut", + "agent_sessions_enabled", + "disable_keychain_access", + "claude_allow_reading_claude_code_credentials", + "codex_external_oauth_sources_allowed", + "open_codex_usage_logs_enabled", + "hide_native_codex_cost_when_open_codex_present", + "powertoys_status_pipe_enabled", + "float_bar_enabled", + "float_bar_click_through", + "tray_panel_always_on_top", + "promote_tray_icon", + // Update settings. + "update_channel", + "auto_download_updates", + "install_updates_on_quit", + // Not part of the v1 allowlist. + "notification_sound_theme", + "enable_animations", + "show_all_token_accounts_in_menu", + "float_bar_provider_ids", + "claude_daily_routines_usage_visible", + "weekly_progress_work_days", + "alibaba_token_plan_region", + "cost_summary_display_style", + // No reader anywhere in the app. + "merge_tray_icons", +]; + +fn typed(value: &Value) -> bool { + serde_json::from_value::(value.clone()).is_ok() +} + +fn one_of(value: &Value, choices: &[&str]) -> bool { + value.as_str().is_some_and(|raw| choices.contains(&raw)) +} + +fn uint_in(value: &Value) -> bool { + value + .as_u64() + .is_some_and(|number| (MIN..=MAX).contains(&number)) +} + +/// `0` is manual refresh; otherwise the interval must be at least a minute +/// (the shortest option offered in Settings) and at most a day. +fn refresh_interval_secs(value: &Value) -> bool { + value + .as_u64() + .is_some_and(|secs| secs == 0 || (60..=86_400).contains(&secs)) +} + +fn percent(value: &Value) -> bool { + value.is_number() + && value + .as_f64() + .is_some_and(|number| number.is_finite() && (0.0..=100.0).contains(&number)) +} + +fn known_provider_ids() -> HashSet<&'static str> { + ProviderId::all().iter().map(|id| id.cli_name()).collect() +} + +fn provider_id_list(value: &Value) -> bool { + let (Some(items), known) = (value.as_array(), known_provider_ids()) else { + return false; + }; + let mut seen = HashSet::new(); + items.iter().all(|item| { + item.as_str() + .is_some_and(|id| known.contains(id) && seen.insert(id)) + }) +} + +fn provider_metrics(value: &Value) -> bool { + let known = known_provider_ids(); + value.as_object().is_some_and(|map| { + map.iter().all(|(provider, metric)| { + known.contains(provider.as_str()) && typed::(metric) + }) + }) +} + +/// Keys are `` or `:session|weekly`; values carry optional +/// `high` / `critical` percentages. +fn provider_usage_thresholds(value: &Value) -> bool { + let known = known_provider_ids(); + let Some(map) = value.as_object() else { + return false; + }; + map.iter().all(|(key, entry)| { + let (provider, window) = key + .split_once(':') + .map_or((key.as_str(), None), |(provider, window)| { + (provider, Some(window)) + }); + known.contains(provider) + && window.is_none_or(|window| matches!(window, "session" | "weekly")) + && entry.as_object().is_some_and(|fields| { + fields.iter().all(|(field, number)| { + matches!(field.as_str(), "high" | "critical") + && (number.is_null() || percent(number)) + }) + }) + }) +} + +fn invalid(key: &str) -> PreferencesError { + PreferencesError::InvalidPreference(key.to_string()) +} + +/// A validated preferences document. +#[derive(Debug, Clone, PartialEq)] +pub struct PreferencesDocument { + preferences: BTreeMap, +} + +impl PreferencesDocument { + /// Parse and fully validate a document. Nothing is applied here. + pub fn from_json(text: &str) -> Result { + if text.len() as u64 > MAX_PREFERENCES_DOCUMENT_BYTES { + return Err(PreferencesError::TooLarge); + } + let root: Value = serde_json::from_str(text.trim_start_matches('\u{feff}')) + .map_err(|_| PreferencesError::Malformed)?; + let Value::Object(mut root) = root else { + return Err(PreferencesError::Malformed); + }; + if root.get("version").and_then(Value::as_u64) != Some(PREFERENCES_DOCUMENT_VERSION) { + return Err(PreferencesError::UnsupportedVersion); + } + let Some(Value::Object(preferences)) = root.remove("preferences") else { + return Err(PreferencesError::Malformed); + }; + if root.len() != 1 { + return Err(PreferencesError::Malformed); + } + Self::validated(preferences) + } + + fn validated(preferences: Map) -> Result { + for (key, value) in &preferences { + let Some((_, valid)) = ALLOWED_PREFERENCES.iter().find(|(name, _)| name == key) else { + return Err(invalid(key)); + }; + if !value.is_null() && !valid(value) { + return Err(invalid(key)); + } + } + Ok(Self { + preferences: preferences.into_iter().collect(), + }) + } + + /// Snapshot the portable preferences of `settings`. + pub fn from_settings(settings: &Settings) -> Result { + let Value::Object(mut all) = serde_json::to_value(settings) + .map_err(|error| PreferencesError::Apply(error.to_string()))? + else { + return Err(PreferencesError::Apply( + "settings did not serialize to an object".to_string(), + )); + }; + let mut preferences = Map::new(); + for (key, _) in ALLOWED_PREFERENCES { + if let Some(value) = all.remove(*key) { + preferences.insert((*key).to_string(), value); + } + } + if let Some(Value::Array(ids)) = preferences.get_mut("enabled_providers") { + // `enabled_providers` is a set; keep exports deterministic. + ids.sort_by(|a, b| a.as_str().cmp(&b.as_str())); + } + Self::validated(preferences) + } + + /// Pretty-printed JSON with a trailing newline. + pub fn to_json(&self) -> String { + let document = serde_json::json!({ + "version": PREFERENCES_DOCUMENT_VERSION, + "preferences": self.preferences, + }); + let mut text = serde_json::to_string_pretty(&document).unwrap_or_default(); + text.push('\n'); + text + } + + /// Number of preferences the document sets. + pub fn len(&self) -> usize { + self.preferences.len() + } + + pub fn is_empty(&self) -> bool { + self.preferences.is_empty() + } + + /// Apply to `settings`. Absent keys are untouched; `null` restores the + /// default. Returns the number of preferences applied. `settings` is left + /// unchanged on error. + pub fn apply_to(&self, settings: &mut Settings) -> Result { + let apply_error = |error: serde_json::Error| PreferencesError::Apply(error.to_string()); + let Value::Object(mut current) = serde_json::to_value(&*settings).map_err(apply_error)? + else { + return Err(PreferencesError::Apply( + "settings did not serialize to an object".to_string(), + )); + }; + let Value::Object(defaults) = + serde_json::to_value(Settings::default()).map_err(apply_error)? + else { + return Err(PreferencesError::Apply( + "default settings did not serialize to an object".to_string(), + )); + }; + for (key, value) in &self.preferences { + let replacement = if value.is_null() { + defaults.get(key).cloned().ok_or_else(|| invalid(key))? + } else { + value.clone() + }; + current.insert(key.clone(), replacement); + } + // Re-read through the normal settings loader so clamping and + // normalization match what `settings.json` gets on load. + *settings = serde_json::from_value(Value::Object(current)).map_err(apply_error)?; + Ok(self.preferences.len()) + } + + /// Read a document from `path`, bounded to [`MAX_PREFERENCES_DOCUMENT_BYTES`]. + pub fn read_file(path: &Path) -> Result { + let file = + std::fs::File::open(path).map_err(|error| PreferencesError::Io(error.to_string()))?; + let mut bytes = Vec::new(); + file.take(MAX_PREFERENCES_DOCUMENT_BYTES + 1) + .read_to_end(&mut bytes) + .map_err(|error| PreferencesError::Io(error.to_string()))?; + if bytes.len() as u64 > MAX_PREFERENCES_DOCUMENT_BYTES { + return Err(PreferencesError::TooLarge); + } + let text = String::from_utf8(bytes).map_err(|_| PreferencesError::Malformed)?; + Self::from_json(&text) + } + + pub fn write_file(&self, path: &Path) -> Result<(), PreferencesError> { + std::fs::write(path, self.to_json()) + .map_err(|error| PreferencesError::Io(error.to_string())) + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/settings/preferences_document/tests.rs b/rust/src/settings/preferences_document/tests.rs new file mode 100644 index 0000000000..a21ae3ef90 --- /dev/null +++ b/rust/src/settings/preferences_document/tests.rs @@ -0,0 +1,474 @@ +use super::*; +use crate::settings::{ + Language, ProviderConfig, ThemePreference, TrayIconMode, UsageThresholdOverride, +}; +use serde_json::json; + +fn document(preferences: Value) -> String { + json!({ "version": 1, "preferences": preferences }).to_string() +} + +fn parse(preferences: Value) -> Result { + PreferencesDocument::from_json(&document(preferences)) +} + +fn max_document_bytes() -> usize { + usize::try_from(MAX_PREFERENCES_DOCUMENT_BYTES).expect("limit fits in usize") +} + +/// Settings as JSON with the set-valued `enabled_providers` in stable order. +fn snapshot(settings: &Settings) -> Value { + let mut value = serde_json::to_value(settings).expect("serialize settings"); + if let Some(Value::Array(ids)) = value.get_mut("enabled_providers") { + ids.sort_by(|a, b| a.as_str().cmp(&b.as_str())); + } + value +} + +fn customized_settings() -> Settings { + let mut settings = Settings { + refresh_interval_secs: 900, + adaptive_refresh: true, + refresh_all_providers_on_menu_open: true, + low_power_mode_preference: crate::settings::LowPowerModePreference::Automatic, + show_notifications: false, + sound_enabled: false, + high_usage_threshold: 55.5, + critical_usage_threshold: 80.0, + predictive_pace_warning_enabled: true, + show_pace: false, + show_as_used: false, + reset_time_relative: false, + show_reset_when_exhausted: true, + hide_personal_info: true, + menu_bar_shows_highest_usage: true, + menu_bar_shows_percent: true, + tray_icon_mode: TrayIconMode::PerProvider, + switcher_shows_icons: false, + overview_layout: "detailed".to_string(), + menu_bar_display_mode: "minimal".to_string(), + // Loading normalizes the order to the full canonical list. + provider_order: crate::settings::normalize_provider_order(&[ + "codex".to_string(), + "claude".to_string(), + ]), + theme: ThemePreference::Dark, + ui_language: Language::Japanese, + window_scale_percent: 150, + tray_scale_percent: 120, + float_bar_opacity: 60, + float_bar_scale: 120, + float_bar_orientation: "vertical".to_string(), + float_bar_style: "taskbar".to_string(), + float_bar_dark_text: true, + float_bar_show_reset_inline: true, + float_bar_show_cost: true, + ..Settings::default() + }; + settings.enabled_providers = ["claude", "codex", "zai"].map(String::from).into(); + settings.provider_usage_thresholds.insert( + "codex:weekly".to_string(), + UsageThresholdOverride { + high: Some(60.0), + critical: None, + }, + ); + settings.provider_metrics.insert( + "claude".to_string(), + crate::settings::MetricPreference::Weekly, + ); + settings +} + +/// Settings whose secret and machine-specific fields are all populated. +fn secret_settings() -> Settings { + let mut settings = customized_settings(); + settings.provider_configs.insert( + ProviderId::Claude, + ProviderConfig { + api_token: Some("SECRET-API-TOKEN".to_string()), + manual_cookie_header: Some("sessionKey=SECRET-COOKIE".to_string()), + management_api_token: Some("SECRET-MGMT".to_string()), + ..ProviderConfig::default() + }, + ); + settings.http_proxy_enabled = true; + settings.http_proxy_url = "http://proxy.invalid:8080".to_string(); + settings.http_proxy_username = "SECRET-PROXY-USER".to_string(); + settings.http_proxy_password = "SECRET-PROXY-PASSWORD".to_string(); + settings.codex_custom_sessions_dirs = vec!["C:\\secret\\sessions".to_string()]; + settings.agent_session_ssh_hosts = vec!["secret-host.invalid".to_string()]; + settings.notification_sound_paths.high_usage = Some("C:\\secret\\ding.wav".to_string()); + settings +} + +#[test] +fn round_trips_every_allowlisted_preference() { + let source = customized_settings(); + let exported = PreferencesDocument::from_settings(&source).expect("export"); + assert_eq!(exported.len(), ALLOWED_PREFERENCES.len()); + + let reparsed = PreferencesDocument::from_json(&exported.to_json()).expect("reimport"); + assert_eq!(reparsed, exported); + + let mut target = Settings { + start_at_login: true, + global_shortcut: "Ctrl+Alt+K".to_string(), + ..Settings::default() + }; + let applied = reparsed.apply_to(&mut target).expect("apply"); + assert_eq!(applied, ALLOWED_PREFERENCES.len()); + + let again = PreferencesDocument::from_settings(&target).expect("re-export"); + assert_eq!(again, exported); + assert_eq!(target.window_scale_percent, 150); + assert_eq!(target.ui_language, Language::Japanese); + assert_eq!( + target.enabled_providers, + ["claude", "codex", "zai"].map(String::from).into() + ); + // Non-portable state on the receiving machine is untouched. + assert!(target.start_at_login); + assert_eq!(target.global_shortcut, "Ctrl+Alt+K"); +} + +#[test] +fn export_is_deterministic_and_sorted() { + let first = PreferencesDocument::from_settings(&customized_settings()) + .expect("export") + .to_json(); + let second = PreferencesDocument::from_settings(&customized_settings()) + .expect("export") + .to_json(); + assert_eq!(first, second); + assert!(first.ends_with("}\n")); + let claude = first.find("\"claude\"").expect("claude listed"); + let codex = first.find("\"codex\"").expect("codex listed"); + let zai = first.find("\"zai\"").expect("zai listed"); + assert!(claude < codex && codex < zai); +} + +#[test] +fn rejection_matrix_names_the_offending_key() { + let cases = [ + ( + "unknown key", + json!({ "not_a_preference": true }), + "not_a_preference", + ), + ( + "provider_configs smuggled in", + json!({ "provider_configs": { "claude": { "api_token": "x" } } }), + "provider_configs", + ), + ( + "http proxy password", + json!({ "http_proxy_password": "x" }), + "http_proxy_password", + ), + ( + "start at login", + json!({ "start_at_login": true }), + "start_at_login", + ), + ("bad enum", json!({ "theme": "sepia" }), "theme"), + ( + "bad language", + json!({ "ui_language": "klingon" }), + "ui_language", + ), + ( + "bad layout", + json!({ "overview_layout": "huge" }), + "overview_layout", + ), + ( + "threshold over 100", + json!({ "high_usage_threshold": 100.5 }), + "high_usage_threshold", + ), + ( + "negative threshold", + json!({ "critical_usage_threshold": -1 }), + "critical_usage_threshold", + ), + ( + "threshold as string", + json!({ "high_usage_threshold": "70" }), + "high_usage_threshold", + ), + ( + "bool as string", + json!({ "show_pace": "true" }), + "show_pace", + ), + ("bool as number", json!({ "show_pace": 1 }), "show_pace"), + ( + "scale too small", + json!({ "window_scale_percent": 99 }), + "window_scale_percent", + ), + ( + "scale too large", + json!({ "tray_scale_percent": 201 }), + "tray_scale_percent", + ), + ( + "scale fractional", + json!({ "window_scale_percent": 100.5 }), + "window_scale_percent", + ), + ( + "interval below a minute", + json!({ "refresh_interval_secs": 5 }), + "refresh_interval_secs", + ), + ( + "interval above a day", + json!({ "refresh_interval_secs": 86_401 }), + "refresh_interval_secs", + ), + ( + "unknown provider id", + json!({ "enabled_providers": ["nope"] }), + "enabled_providers", + ), + ( + "duplicate provider id", + json!({ "provider_order": ["codex", "codex"] }), + "provider_order", + ), + ( + "provider alias", + json!({ "enabled_providers": ["openai"] }), + "enabled_providers", + ), + ( + "metric value", + json!({ "provider_metrics": { "codex": "bogus" } }), + "provider_metrics", + ), + ( + "metric provider", + json!({ "provider_metrics": { "nope": "weekly" } }), + "provider_metrics", + ), + ( + "threshold window", + json!({ "provider_usage_thresholds": { "codex:monthly": { "high": 50 } } }), + "provider_usage_thresholds", + ), + ( + "threshold field", + json!({ "provider_usage_thresholds": { "codex": { "low": 50 } } }), + "provider_usage_thresholds", + ), + ( + "threshold range", + json!({ "provider_usage_thresholds": { "codex": { "high": 150 } } }), + "provider_usage_thresholds", + ), + ]; + for (name, preferences, key) in cases { + assert_eq!( + parse(preferences).expect_err(name), + PreferencesError::InvalidPreference(key.to_string()), + "{name}" + ); + } +} + +#[test] +fn rejects_bad_envelopes() { + for text in [ + r#"{"version":2,"preferences":{}}"#, + r#"{"version":"1","preferences":{}}"#, + r#"{"preferences":{}}"#, + ] { + assert_eq!( + PreferencesDocument::from_json(text), + Err(PreferencesError::UnsupportedVersion), + "{text}" + ); + } + for text in [ + "not json", + "[]", + r#"{"version":1}"#, + r#"{"version":1,"preferences":[]}"#, + r#"{"version":1,"preferences":{},"extra":true}"#, + ] { + assert_eq!( + PreferencesDocument::from_json(text), + Err(PreferencesError::Malformed), + "{text}" + ); + } + let oversized = format!( + r#"{{"version":1,"preferences":{{}},"pad":"{}"}}"#, + "x".repeat(max_document_bytes()) + ); + assert_eq!( + PreferencesDocument::from_json(&oversized), + Err(PreferencesError::TooLarge) + ); +} + +#[test] +fn errors_never_echo_values() { + let error = parse(json!({ "theme": "TOP-SECRET-VALUE" })).expect_err("bad theme"); + assert!(!error.to_string().contains("TOP-SECRET-VALUE")); + let error = PreferencesDocument::from_json("TOP-SECRET-VALUE").expect_err("not json"); + assert!(!error.to_string().contains("TOP-SECRET-VALUE")); +} + +#[test] +fn missing_keys_leave_settings_unchanged() { + let empty = parse(json!({})).expect("empty document"); + let mut settings = customized_settings(); + // The first apply normalizes state the way a settings load does. + empty.apply_to(&mut settings).expect("apply"); + let before = snapshot(&settings); + assert_eq!(empty.apply_to(&mut settings).expect("apply"), 0); + assert_eq!(snapshot(&settings), before); + + parse(json!({ "theme": "light" })) + .expect("theme") + .apply_to(&mut settings) + .expect("apply"); + assert_eq!(settings.theme, ThemePreference::Light); + assert_eq!(settings.window_scale_percent, 150); +} + +#[test] +fn null_restores_the_default() { + let mut settings = customized_settings(); + let applied = parse(json!({ + "window_scale_percent": null, + "theme": null, + "enabled_providers": null, + "provider_usage_thresholds": null, + "provider_metrics": null, + "provider_order": null, + "show_pace": null, + })) + .expect("nulls are valid") + .apply_to(&mut settings) + .expect("apply"); + assert_eq!(applied, 7); + + let defaults = Settings::default(); + assert_eq!(settings.window_scale_percent, defaults.window_scale_percent); + assert_eq!(settings.theme, defaults.theme); + assert_eq!(settings.enabled_providers, defaults.enabled_providers); + assert!(settings.provider_usage_thresholds.is_empty()); + assert!(settings.provider_metrics.is_empty()); + assert!(settings.provider_order.is_empty()); + assert!(settings.show_pace); + // Untouched preferences keep their customized values. + assert_eq!(settings.tray_scale_percent, 120); +} + +#[test] +fn allowed_and_excluded_keys_classify_every_settings_field() { + let Value::Object(serialized) = snapshot(&secret_settings()) else { + panic!("settings must serialize to an object"); + }; + let allowed: HashSet<&str> = ALLOWED_PREFERENCES.iter().map(|(key, _)| *key).collect(); + let excluded: HashSet<&str> = EXCLUDED_KEYS.iter().copied().collect(); + + assert_eq!( + allowed.len(), + ALLOWED_PREFERENCES.len(), + "duplicate allowlist key" + ); + assert_eq!( + excluded.len(), + EXCLUDED_KEYS.len(), + "duplicate excluded key" + ); + assert!( + allowed.is_disjoint(&excluded), + "a key is both allowed and excluded" + ); + + for key in serialized.keys() { + assert!( + allowed.contains(key.as_str()) || excluded.contains(key.as_str()), + "Settings field `{key}` is not classified: add it to ALLOWED_PREFERENCES \ + (portable) or EXCLUDED_KEYS (never exported)" + ); + } + for key in allowed.iter().chain(excluded.iter()) { + assert!( + serialized.contains_key(*key), + "`{key}` is classified but is not a Settings field" + ); + } +} + +#[test] +fn export_never_contains_secrets_or_machine_state() { + let text = PreferencesDocument::from_settings(&secret_settings()) + .expect("export") + .to_json(); + for secret in [ + "SECRET", + "proxy.invalid", + "secret-host", + "secret\\\\", + "sessionKey", + "provider_configs", + "http_proxy", + "start_at_login", + "global_shortcut", + ] { + assert!(!text.contains(secret), "export leaked `{secret}`"); + } +} + +#[test] +fn applying_keeps_everything_else_on_the_receiving_machine() { + let empty = parse(json!({})).expect("empty"); + let mut target = secret_settings(); + empty.apply_to(&mut target).expect("apply"); + let before = snapshot(&target); + // The whole-struct JSON round trip used by apply is lossless once normalized. + empty.apply_to(&mut target).expect("apply"); + assert_eq!(snapshot(&target), before); + + parse(json!({ "theme": "light" })) + .expect("theme") + .apply_to(&mut target) + .expect("apply"); + let config = target + .provider_configs + .get(&ProviderId::Claude) + .expect("config"); + assert_eq!(config.api_token.as_deref(), Some("SECRET-API-TOKEN")); + assert_eq!(target.http_proxy_password, "SECRET-PROXY-PASSWORD"); +} + +#[test] +fn file_round_trip_and_size_bound() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("preferences.json"); + let exported = PreferencesDocument::from_settings(&customized_settings()).expect("export"); + exported.write_file(&path).expect("write"); + assert_eq!( + PreferencesDocument::read_file(&path).expect("read"), + exported + ); + + let big = dir.path().join("big.json"); + std::fs::write(&big, vec![b' '; max_document_bytes() + 10]).expect("write"); + assert_eq!( + PreferencesDocument::read_file(&big), + Err(PreferencesError::TooLarge) + ); + + assert!(matches!( + PreferencesDocument::read_file(&dir.path().join("missing.json")), + Err(PreferencesError::Io(_)) + )); +} From 3de6386c91de49882e4bfcf963d8500f6bcc2000 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:16:46 +0700 Subject: [PATCH 120/407] Port upstream 0.67.0: add cost --period and period-aware serve /cost --- docs/CLI.md | 9 ++ rust/src/cli/cost.rs | 173 ++++++++++++++++++++++----- rust/src/cli/cost_period.rs | 189 ++++++++++++++++++++++++++++++ rust/src/cli/mod.rs | 1 + rust/src/cli/serve/data.rs | 46 +++++--- rust/src/cost_reporting_period.rs | 19 +++ 6 files changed, 391 insertions(+), 46 deletions(-) create mode 100644 rust/src/cli/cost_period.rs diff --git a/docs/CLI.md b/docs/CLI.md index a9c25bcdcc..1106e15a13 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -64,8 +64,17 @@ codexbar cost codexbar cost -p codex -f json --pretty codexbar cost -p codex --remote user@mac-host codexbar cost -p codex --format json --summary-only --provider-native-only --days 30 +codexbar cost --period month-to-date --json +codexbar cost --period all --json +codexbar cost --days 7 --json ``` +**Reporting period.** `--period month-to-date|all` selects the calendar month so far (midnight on the first of the month, in the machine's local zone, through today) or all available local history. `--days N` (1..=365) always selects a rolling window and wins over `--period`. With neither flag the saved `cost_reporting_period` in `settings.json` applies (`rolling:30`, `month-to-date` or `all`; 30 days when unset). Text headings show the period label (`Last 30 days`, `Month to date`, `All`, `Today`). JSON payloads add `reportingPeriod` (the raw value), `historyLabel`, and a `totals` object (`inputTokens`, `outputTokens`, `cachedTokens`, `reasoningTokens`, `totalTokens`, `totalCost`) for the selected window, and `days_scanned` is the resolved day count; the existing `cost`, `tokens`, and `sessions_count` fields keep their meaning. `totals` is `null` when a scan found nothing without establishing a known zero. + +`--remote` and `--summary-only` keep the 1..=365 day protocol: month to date is sent as its current day count, and All is rejected unless you pass an explicit `--days N`. `--group-by session` lists the most recent 365 days at most when All is selected. All reads every available local log, including logs older than a year; missing or deleted logs cannot be recovered and incomplete scans stay marked incomplete. + +`serve` `/cost` reads the saved period on every request, so `days_scanned`, `totals`, the scanner windows and the `daily` chart rows follow it (`daily` covers at most the latest 365 days for All). The dashboard cost line, the dashboard snapshot, and the Prometheus gauge `codexbar_cost_last_30_days_usd` keep their 30-day meaning. + 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. `--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. diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 8e9be16dc2..96d85a8f0a 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -4,11 +4,15 @@ use clap::Args; +use super::cost_period::{ + cost_totals_json, resolve_period, rolling_window_days, stamp_period, window_days, +}; use super::usage::{OutputFormat, ProviderSelection}; use crate::codex_costs::{ CodexHostCostReport, CodexHostCostsArgs, CodexHostOutcome, run_codex_host_costs, }; use crate::core::{CostScanOptions, ProviderId}; +use crate::cost_reporting_period::CostReportingPeriod; use crate::cost_scanner::{CostScanner, CostSummary}; use crate::settings::Settings; use crate::spend_contract::build_local_spend_contract_from_summary; @@ -36,9 +40,14 @@ pub struct CostArgs { #[arg(long)] pub pretty: bool, - /// Number of days to scan (default: 30) - #[arg(short, long, default_value = "30")] - pub days: u32, + /// Cost history window in days (1..=365); always rolling and overrides --period + #[arg(short, long)] + pub days: Option, + + /// Cost period: month-to-date or all. Without --days or --period the + /// saved cost period applies (30 days by default) + #[arg(long)] + pub period: Option, /// A16 (upstream 0.48.0): exclude pi/OMP-compatible agent session mirrors, /// reporting only the provider-native local JSONL logs. When omitted @@ -91,9 +100,18 @@ 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(); - if args.remote.is_some() || args.summary_only { + 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, + host_summary, + )?; + let days = window_days(period); + + if host_summary { return run_codex_host_costs(&CodexHostCostsArgs { - days: args.days, + days, remote: args.remote.clone(), summary_only: args.summary_only, pretty: args.pretty, @@ -115,13 +133,14 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { // 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); + let scanner = CostScanner::for_period(period).with_options(scan_options); tracing::debug!( - "Running cost command: providers={:?}, format={:?}, days={}", + "Running cost command: providers={:?}, format={:?}, period={}, days={}", providers.as_list(), format, - args.days + period, + days ); // Collect cost data for requested providers @@ -170,12 +189,12 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { summary: CostSummary::default(), supported: true, token_history: Some(crate::providers::antigravity::local_sessions::summarize( - args.days, + days, )), }); } ProviderId::Muse => { - let report = crate::providers::muse::local_usage::scan(args.days, None); + let report = crate::providers::muse::local_usage::scan(days, None); results.push(CostResult { provider: provider.cli_name().to_string(), display_name: provider.display_name().to_string(), @@ -199,10 +218,10 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { match format { OutputFormat::Text => { - print_text_output(&results, use_color, args.days, group_by); + print_text_output(&results, use_color, period, group_by); } OutputFormat::Json => { - print_json_output(&results, args.pretty, args.days)?; + print_json_output(&results, args.pretty, period, days)?; } } @@ -273,12 +292,18 @@ struct CostResult { } /// Print text output -fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_by: CostGroupBy) { +fn print_text_output( + results: &[CostResult], + use_color: bool, + period: CostReportingPeriod, + group_by: CostGroupBy, +) { + let label = period.label(); for (i, result) in results.iter().enumerate() { let title = if result.token_history.is_some() { - format!("{} Token History (last {} days)", result.display_name, days) + format!("{} Token History ({label})", result.display_name) } else { - format!("{} Cost (last {} days)", result.display_name, days) + format!("{} Cost ({label})", result.display_name) }; if use_color { println!("\x1b[1m{title}\x1b[0m"); @@ -287,9 +312,9 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b } if let Some(history) = result.token_history { - print_local_token_history(history, days); + print_local_token_history(history); } else if group_by == CostGroupBy::Session && result.provider == "codex" { - print_codex_session_output(result, days); + print_codex_session_output(result, period); } else if group_by == CostGroupBy::Session { println!(" Session grouping is only available for Codex local conversations"); } else if !result.supported { @@ -297,7 +322,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b println!(" (Only Codex and Claude have local logs)"); } else if result.summary.sessions_count == 0 { if result.summary.known_zero { - println!(" No usage in the last {} days (scan complete)", days); + println!(" No usage in the selected period (scan complete)"); } else { println!(" No usage data found"); println!(" Check that you have used {} locally", result.display_name); @@ -372,11 +397,11 @@ 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) { use crate::spend_contract::LocalHistoryCoverage; match history.coverage { LocalHistoryCoverage::Complete if history.total_tokens == 0 => { - println!(" No token usage in the last {days} days (scan complete)"); + println!(" No token usage in the selected period (scan complete)"); } LocalHistoryCoverage::Complete => { println!(" Tokens: {} total", format_number(history.total_tokens)); @@ -389,8 +414,10 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history; dollar costs unavailable"); } -fn print_codex_session_output(result: &CostResult, days: u32) { - let index = crate::codex_workspaces::CodexWorkspacesIndex::new(days); +fn print_codex_session_output(result: &CostResult, period: CostReportingPeriod) { + // The conversation index keeps rolling 1..=365 day windows, so All lists + // the most recent year; the heading below reports the window it used. + let index = crate::codex_workspaces::CodexWorkspacesIndex::new(rolling_window_days(period)); let snapshot = match index.load_snapshot(false, |_| {}) { Ok(snapshot) => snapshot, Err(err) => { @@ -454,13 +481,20 @@ fn short_session_id(value: &str) -> String { } /// Print JSON output -fn build_json_payloads(results: &[CostResult], days: u32) -> Vec { +fn build_json_payloads( + results: &[CostResult], + period: CostReportingPeriod, + days: u32, +) -> Vec { let settings = Settings::load(); results .iter() .map(|r| { if let Some(history) = r.token_history { - return crate::spend_contract::local_token_history_json(&r.provider, history, days); + let mut payload = + crate::spend_contract::local_token_history_json(&r.provider, history, days); + stamp_period(&mut payload, period); + return payload; } if !r.supported { serde_json::json!({ @@ -478,10 +512,11 @@ fn build_json_payloads(results: &[CostResult], days: u32) -> Vec Vec>(), "period": {"start": r.summary.period_start.map(|d| d.to_string()), "end": r.summary.period_end.map(|d| d.to_string())}, "spendContract": spend_contract - }) + }); + stamp_period(&mut payload, period); + payload } }) .collect() } -fn print_json_output(results: &[CostResult], pretty: bool, days: u32) -> anyhow::Result<()> { - let payloads = build_json_payloads(results, days); +fn print_json_output( + results: &[CostResult], + pretty: bool, + period: CostReportingPeriod, + days: u32, +) -> anyhow::Result<()> { + let payloads = build_json_payloads(results, period, days); let output = if pretty { serde_json::to_string_pretty(&payloads)? @@ -646,6 +688,81 @@ mod tests { assert!(partial["tokens"]["total"].is_null()); assert_eq!(partial["historyCoverage"], "partial"); } + #[derive(clap::Parser)] + struct Wrapper { + #[command(flatten)] + args: CostArgs, + } + + fn parse_args(argv: &[&str]) -> CostArgs { + use clap::Parser; + Wrapper::try_parse_from(std::iter::once("cost").chain(argv.iter().copied())) + .expect("cost args parse") + .args + } + + #[test] + fn days_and_period_are_detectable_when_absent() { + let args = parse_args(&[]); + assert_eq!(args.days, None); + assert_eq!(args.period, None); + let args = parse_args(&["--period", "month-to-date", "--days", "7"]); + assert_eq!(args.days, Some(7)); + assert_eq!(args.period.as_deref(), Some("month-to-date")); + } + + #[test] + fn json_payload_carries_period_identity_and_totals() { + let summary = CostSummary { + total_cost_usd: 2.5, + input_tokens: 10, + output_tokens: 5, + sessions_count: 1, + ..Default::default() + }; + let results = [CostResult { + provider: "claude".to_string(), + display_name: "Claude".to_string(), + summary, + supported: true, + token_history: None, + }]; + let payloads = build_json_payloads(&results, CostReportingPeriod::MonthToDate, 12); + let payload = &payloads[0]; + assert_eq!(payload["reportingPeriod"], "month-to-date"); + assert_eq!(payload["historyLabel"], "Month to date"); + assert_eq!(payload["days_scanned"], 12); + assert_eq!(payload["totals"]["totalCost"], 2.5); + assert_eq!(payload["totals"]["totalTokens"], 15); + // Existing fields keep their meaning. + assert_eq!(payload["cost"]["total_usd"], 2.5); + assert_eq!(payload["tokens"]["input"], 10); + + let rolling = build_json_payloads(&results, CostReportingPeriod::Rolling(30), 30); + assert_eq!(rolling[0]["reportingPeriod"], "rolling:30"); + assert_eq!(rolling[0]["historyLabel"], "Last 30 days"); + } + + #[test] + fn token_history_payload_is_stamped_with_the_period() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let results = [CostResult { + provider: "antigravity".to_string(), + display_name: "Antigravity".to_string(), + summary: CostSummary::default(), + supported: true, + token_history: Some(LocalTokenHistorySummary { + total_tokens: 3, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + }), + }]; + let payloads = build_json_payloads(&results, CostReportingPeriod::AllAvailable, 20_000); + assert_eq!(payloads[0]["reportingPeriod"], "all"); + assert_eq!(payloads[0]["historyLabel"], "All"); + assert_eq!(payloads[0]["days_scanned"], 20_000); + } + #[test] fn provider_native_only_flag_default_false() { // Default CostArgs has provider_native_only = false (backward compat). diff --git a/rust/src/cli/cost_period.rs b/rust/src/cli/cost_period.rs new file mode 100644 index 0000000000..6335b5bc37 --- /dev/null +++ b/rust/src/cli/cost_period.rs @@ -0,0 +1,189 @@ +//! Reporting-period selection shared by `codexbar cost` and `serve` `/cost` +//! (upstream 0.67.0 `decodeCostReportingPeriod`, `costTotals`). + +use chrono::Utc; +use serde_json::{Value, json}; + +use crate::cost_reporting_period::{CostReportingPeriod, CostTimeZone, MAX_ROLLING_DAYS}; +use crate::cost_scanner::CostSummary; + +/// Pick the period for one `cost` invocation. +/// +/// `--days N` always selects a rolling window and wins over `--period`; +/// without either flag the saved selection applies. Host summaries +/// (`--remote`, `--summary-only`) speak a 1..=365 day protocol, so All needs an +/// explicit `--days`. +pub(super) fn resolve_period( + days: Option, + period: Option<&str>, + saved: CostReportingPeriod, + host_summary: bool, +) -> anyhow::Result { + if let Some(days) = days { + 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"))?, + None => saved, + }; + if host_summary && period == CostReportingPeriod::AllAvailable { + anyhow::bail!("host summaries support 1...365 days; override All with --days N"); + } + Ok(period) +} + +/// Days covered by `period` today, in the zone the local scanners bucket by. +pub(crate) fn window_days(period: CostReportingPeriod) -> u32 { + period.days(Utc::now(), CostTimeZone::Local, None) +} + +/// Days for consumers that only understand a rolling 1..=365 window (daily +/// chart rows, conversation grouping). Month to date fits; All is capped. +pub(crate) fn rolling_window_days(period: CostReportingPeriod) -> u32 { + window_days(period).min(MAX_ROLLING_DAYS) +} + +/// Add `reportingPeriod` (raw) and `historyLabel` to one provider payload. +pub(crate) fn stamp_period(payload: &mut Value, period: CostReportingPeriod) { + if let Some(object) = payload.as_object_mut() { + object.insert("reportingPeriod".to_string(), json!(period.raw())); + object.insert("historyLabel".to_string(), json!(period.label())); + } +} + +/// Totals for the selected window. `null` when the scan found nothing and did +/// not establish a known zero, so a missing scan is never shown as `$0`. +pub(crate) fn cost_totals_json(summary: &CostSummary) -> Value { + if summary.sessions_count == 0 && !summary.known_zero { + return Value::Null; + } + json!({ + "inputTokens": summary.input_tokens, + "outputTokens": summary.output_tokens, + "cachedTokens": summary.cached_tokens, + "reasoningTokens": summary.reasoning_tokens, + "totalTokens": summary.input_tokens.saturating_add(summary.output_tokens), + "totalCost": summary.total_cost_usd, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const SAVED_MTD: CostReportingPeriod = CostReportingPeriod::MonthToDate; + + #[test] + fn days_beats_period_and_clamps() { + let resolved = resolve_period(Some(7), Some("all"), SAVED_MTD, false).unwrap(); + assert_eq!(resolved, CostReportingPeriod::Rolling(7)); + let clamped = resolve_period(Some(900), None, SAVED_MTD, false).unwrap(); + assert_eq!(clamped, CostReportingPeriod::Rolling(365)); + let floor = resolve_period(Some(0), None, SAVED_MTD, false).unwrap(); + assert_eq!(floor, CostReportingPeriod::Rolling(1)); + } + + #[test] + fn days_wins_even_over_an_invalid_period() { + let resolved = resolve_period(Some(7), Some("nonsense"), SAVED_MTD, false).unwrap(); + assert_eq!(resolved, CostReportingPeriod::Rolling(7)); + } + + #[test] + fn absent_flags_use_the_saved_period() { + let default = CostReportingPeriod::default(); + assert_eq!( + resolve_period(None, None, default, false).unwrap(), + CostReportingPeriod::Rolling(30) + ); + assert_eq!( + resolve_period(None, None, SAVED_MTD, false).unwrap(), + SAVED_MTD + ); + } + + #[test] + fn explicit_period_overrides_saved() { + let saved = CostReportingPeriod::Rolling(7); + assert_eq!( + resolve_period(None, Some("all"), saved, false).unwrap(), + CostReportingPeriod::AllAvailable + ); + assert_eq!( + resolve_period(None, Some("month-to-date"), saved, false).unwrap(), + CostReportingPeriod::MonthToDate + ); + } + + #[test] + fn invalid_period_is_rejected() { + let error = resolve_period(None, Some("weekly"), SAVED_MTD, false).unwrap_err(); + assert_eq!(error.to_string(), "--period must be month-to-date or all"); + assert!(resolve_period(None, Some("rolling:0"), SAVED_MTD, false).is_err()); + } + + #[test] + fn host_summaries_reject_all_without_explicit_days() { + let error = resolve_period(None, Some("all"), SAVED_MTD, true).unwrap_err(); + assert!(error.to_string().contains("override All with --days N")); + let saved_all = resolve_period(None, None, CostReportingPeriod::AllAvailable, true); + assert!(saved_all.is_err()); + assert_eq!( + resolve_period(Some(14), Some("all"), SAVED_MTD, true).unwrap(), + CostReportingPeriod::Rolling(14) + ); + assert_eq!( + resolve_period(None, None, SAVED_MTD, true).unwrap(), + SAVED_MTD + ); + } + + #[test] + fn rolling_window_is_capped_at_a_year_for_all() { + assert_eq!(rolling_window_days(CostReportingPeriod::AllAvailable), 365); + assert_eq!(rolling_window_days(CostReportingPeriod::Rolling(30)), 30); + assert!(window_days(CostReportingPeriod::AllAvailable) > 365); + assert!((1..=31).contains(&window_days(CostReportingPeriod::MonthToDate))); + } + + #[test] + fn stamp_adds_raw_period_and_label() { + let mut payload = json!({"provider": "codex"}); + stamp_period(&mut payload, CostReportingPeriod::MonthToDate); + assert_eq!(payload["reportingPeriod"], "month-to-date"); + assert_eq!(payload["historyLabel"], "Month to date"); + assert_eq!(payload["provider"], "codex"); + + let mut rolling = json!({}); + stamp_period(&mut rolling, CostReportingPeriod::Rolling(30)); + assert_eq!(rolling["reportingPeriod"], "rolling:30"); + assert_eq!(rolling["historyLabel"], "Last 30 days"); + } + + #[test] + fn totals_sum_tokens_and_keep_unknown_distinct_from_zero() { + let summary = CostSummary { + total_cost_usd: 1.5, + input_tokens: 100, + output_tokens: 20, + cached_tokens: 7, + sessions_count: 2, + ..CostSummary::default() + }; + let totals = cost_totals_json(&summary); + assert_eq!(totals["inputTokens"], 100); + assert_eq!(totals["outputTokens"], 20); + assert_eq!(totals["cachedTokens"], 7); + assert_eq!(totals["totalTokens"], 120); + assert_eq!(totals["totalCost"], 1.5); + assert!(totals["reasoningTokens"].is_null()); + + assert!(cost_totals_json(&CostSummary::default()).is_null()); + let known_zero = CostSummary { + known_zero: true, + ..CostSummary::default() + }; + assert_eq!(cost_totals_json(&known_zero)["totalCost"], 0.0); + } +} diff --git a/rust/src/cli/mod.rs b/rust/src/cli/mod.rs index 497581222c..714073ed58 100755 --- a/rust/src/cli/mod.rs +++ b/rust/src/cli/mod.rs @@ -15,6 +15,7 @@ pub mod account; pub mod autostart; pub mod config; pub mod cost; +mod cost_period; pub mod dashboard; pub mod diagnose; pub mod guard; diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index d1841db6bc..b086745d90 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -6,9 +6,11 @@ use serde_json::json; +use crate::cli::cost_period::{cost_totals_json, rolling_window_days, stamp_period, window_days}; use crate::cli::usage::ProviderSelection; use crate::core::{CostScanOptions, FetchContext, ProviderId, SourceMode, instantiate_provider}; use crate::cost_scanner::{self, CostScanner}; +use crate::settings::Settings; use super::json_response; @@ -64,25 +66,29 @@ pub async fn cost_response(provider: Option<&str>) -> String { return json_response(400, json!({ "error": error.to_string() })); } }; - let scanner = CostScanner::new(30).with_options(CostScanOptions::app_driven()); + // The saved selection is read per request, so a change in Settings (or a + // month rollover for month to date) applies without restarting `serve`. + // There is no `/cost` response cache; the scanners' own caches are keyed by + // the resolved day range, so one period never reuses another's entries. + let period = Settings::load().cost_reporting_period; + let days = window_days(period); + let scanner = CostScanner::for_period(period).with_options(CostScanOptions::app_driven()); let mut results = Vec::new(); for provider_id in selection.as_list() { if provider_id == ProviderId::Antigravity { - let history = crate::providers::antigravity::local_sessions::summarize(30); - results.push(crate::spend_contract::local_token_history_json( - "antigravity", - history, - 30, - )); + let history = crate::providers::antigravity::local_sessions::summarize(days); + let mut payload = + crate::spend_contract::local_token_history_json("antigravity", history, days); + stamp_period(&mut payload, period); + results.push(payload); continue; } if provider_id == ProviderId::Muse { - let report = crate::providers::muse::local_usage::scan(30, None); - results.push(crate::spend_contract::local_token_history_json( - "muse", - report.into(), - 30, - )); + let report = crate::providers::muse::local_usage::scan(days, None); + let mut payload = + crate::spend_contract::local_token_history_json("muse", report.into(), days); + stamp_period(&mut payload, period); + results.push(payload); continue; } let (supported, summary) = match provider_id { @@ -94,15 +100,17 @@ pub async fn cost_response(provider: Option<&str>) -> String { 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. + // 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(), - 30, + rolling_window_days(period), )); - results.push(json!({ + let mut payload = json!({ "provider": provider_id.cli_name(), "supported": true, - "days_scanned": 30, + "days_scanned": days, + "totals": cost_totals_json(&summary), "cost": { "total_usd": summary.total_cost_usd, "currency": "USD" @@ -115,7 +123,9 @@ pub async fn cost_response(provider: Option<&str>) -> String { }, "sessions_count": summary.sessions_count, "by_model": summary.by_model, - })); + }); + stamp_period(&mut payload, period); + results.push(payload); } else { results.push(json!({ "provider": provider_id.cli_name(), diff --git a/rust/src/cost_reporting_period.rs b/rust/src/cost_reporting_period.rs index c1721afaa3..e59be25a3a 100644 --- a/rust/src/cost_reporting_period.rs +++ b/rust/src/cost_reporting_period.rs @@ -158,6 +158,17 @@ impl CostReportingPeriod { } } + /// Human label shown next to totals: "Today", "Last N days", "Month to + /// date" or "All". + pub fn label(&self) -> String { + match self { + Self::Rolling(1) => "Today".to_string(), + Self::Rolling(days) => format!("Last {days} days"), + Self::MonthToDate => "Month to date".to_string(), + Self::AllAvailable => "All".to_string(), + } + } + /// Resolve a saved selection, falling back to a legacy day count and then /// to the 30-day default. An unreadable saved value migrates like a /// missing one. @@ -287,6 +298,14 @@ mod tests { } } + #[test] + fn labels_match_upstream_wording() { + assert_eq!(CostReportingPeriod::Rolling(1).label(), "Today"); + assert_eq!(CostReportingPeriod::Rolling(30).label(), "Last 30 days"); + assert_eq!(CostReportingPeriod::MonthToDate.label(), "Month to date"); + assert_eq!(CostReportingPeriod::AllAvailable.label(), "All"); + } + #[test] fn identity_distinguishes_periods_rollover_and_zones() { let january = at("2026-02-01T07:59:59Z"); From c9fc00418fc94f6365588cd8969f47e0f99b879f Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:18:05 +0700 Subject: [PATCH 121/407] Port upstream 0.67.0: use --file for preferences CLI like upstream --- docs/CONFIGURATION.md | 4 ++-- rust/src/cli/config.rs | 20 ++++++++++++-------- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 03b6f7221f..198e2ec513 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -51,8 +51,8 @@ codexbar config path printf '%s' $env:OPENROUTER_API_KEY | codexbar config set-api-key -p openrouter --stdin # Portable preferences (also under Settings > Advanced) -codexbar config preferences export prefs.json -codexbar config preferences import prefs.json +codexbar config preferences export --file prefs.json +codexbar config preferences import --file prefs.json ``` Portable preferences are a versioned JSON document (`{"version": 1, "preferences": {...}}`) diff --git a/rust/src/cli/config.rs b/rust/src/cli/config.rs index d7bb8c1e41..8d07cddb7b 100755 --- a/rust/src/cli/config.rs +++ b/rust/src/cli/config.rs @@ -69,12 +69,14 @@ pub enum PreferencesAction { /// Write portable preferences to a JSON file Export { /// Destination file - path: std::path::PathBuf, + #[arg(long)] + file: std::path::PathBuf, }, /// Apply a preferences file; restart a running CodexBar afterwards Import { /// Preferences file to read - path: std::path::PathBuf, + #[arg(long)] + file: std::path::PathBuf, }, } @@ -103,7 +105,7 @@ pub async fn run(args: ConfigArgs) -> anyhow::Result<()> { /// Export or import the portable preferences document. fn transfer_preferences(action: PreferencesAction) -> anyhow::Result<()> { match action { - PreferencesAction::Export { path } => { + PreferencesAction::Export { file: path } => { let document = PreferencesDocument::from_settings(&Settings::load())?; document.write_file(&path)?; println!( @@ -112,7 +114,7 @@ fn transfer_preferences(action: PreferencesAction) -> anyhow::Result<()> { path.display() ); } - PreferencesAction::Import { path } => { + PreferencesAction::Import { file: path } => { let document = PreferencesDocument::read_file(&path)?; let mut settings = Settings::load(); let applied = document.apply_to(&mut settings)?; @@ -535,16 +537,18 @@ mod tests { use super::{ConfigArgs, ConfigCommand, PreferencesAction}; use clap::Parser; - let export = ConfigArgs::try_parse_from(["config", "preferences", "export", "p.json"]) - .expect("export parses"); + let export = + ConfigArgs::try_parse_from(["config", "preferences", "export", "--file", "p.json"]) + .expect("export parses"); assert!(matches!( export.command, ConfigCommand::Preferences { action: PreferencesAction::Export { .. } } )); - let import = ConfigArgs::try_parse_from(["config", "preferences", "import", "p.json"]) - .expect("import parses"); + let import = + ConfigArgs::try_parse_from(["config", "preferences", "import", "--file", "p.json"]) + .expect("import parses"); assert!(matches!( import.command, ConfigCommand::Preferences { From 81af533ed86ac160321eac16a32746d7ecde6619 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:21:09 +0700 Subject: [PATCH 122/407] Port upstream 0.68.0: read selected Muse browser team quota when login has none Optional dev.meta.ai selected-team quota (upstream #4011). Cookie source is Off by default; the Windows default (manual, no cookie) reads no browser. Adds a team picker, digits-only team id, and the oauth+web source label. --- .../src/commands/provider_settings.rs | 40 + .../src-tauri/src/commands/providers.rs | 18 +- .../src-tauri/src/commands/tests.rs | 68 ++ apps/desktop-tauri/src/i18n/keys.ts | 4 + .../settings/providers/ProviderDetailPane.tsx | 10 + .../sections/MuseBrowserTeamSection.test.tsx | 96 +++ .../sections/MuseBrowserTeamSection.tsx | 109 +++ docs/PROVIDERS.md | 23 + rust/src/cli/diagnose.rs | 1 + rust/src/cli/guard.rs | 1 + rust/src/cli/hooks.rs | 1 + rust/src/cli/serve/dashboard/source.rs | 2 + rust/src/cli/serve/data.rs | 1 + rust/src/cli/usage.rs | 1 + rust/src/core/provider.rs | 20 +- rust/src/locale.rs | 4 + rust/src/locale/en-US.ftl | 4 + rust/src/providers/muse/mod.rs | 60 +- rust/src/providers/muse/web_quota/mod.rs | 665 +++++++++++++++++ rust/src/providers/muse/web_quota/tests.rs | 703 ++++++++++++++++++ rust/src/settings/provider_workspace.rs | 21 + 21 files changed, 1827 insertions(+), 25 deletions(-) create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/MuseBrowserTeamSection.test.tsx create mode 100644 apps/desktop-tauri/src/surfaces/settings/providers/sections/MuseBrowserTeamSection.tsx create mode 100644 rust/src/providers/muse/web_quota/mod.rs create mode 100644 rust/src/providers/muse/web_quota/tests.rs 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 f9344fcfba..0c93ff623a 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -223,6 +223,7 @@ fn cookie_source_provider(provider_id: &str) -> Option ProviderId::Sakana, "notion" => ProviderId::Notion, "grok" => ProviderId::Grok, + "muse" => ProviderId::Muse, "replicate" => ProviderId::Replicate, "helmcode" => ProviderId::Helmcode, "typesafe" => ProviderId::TypeSafe, @@ -329,6 +330,7 @@ fn workspace_provider(provider_id: &str) -> Option { "xai" => ProviderId::Xai, "v0" => ProviderId::V0, "helmcode" => ProviderId::Helmcode, + "muse" => ProviderId::Muse, _ => return None, }) } @@ -386,6 +388,21 @@ mod tests { use super::{litellm_workspace_change_allowed, workspace_provider}; + #[test] + fn muse_exposes_cookie_source_and_browser_team_settings() { + assert_eq!(workspace_provider("muse"), Some(ProviderId::Muse)); + assert_eq!( + super::cookie_source_provider("muse"), + Some(ProviderId::Muse) + ); + let values: Vec = + super::cookie_source_options_for("muse", codexbar::settings::Language::English) + .into_iter() + .map(|option| option.value) + .collect(); + assert_eq!(values, ["auto", "manual", "off"]); + } + #[test] fn maps_opencode_go_workspace_provider() { assert_eq!( @@ -733,6 +750,29 @@ pub fn cookie_source_options_for(provider_id: &str, lang: Language) -> Vec vec![ + cookie_option( + lang, + "auto", + "Reads the selected team's quota with the signed-in dev.meta.ai browser session.", + "", + None, + ), + cookie_option( + lang, + "manual", + "", + "Paste the llama_dev_sess cookie from dev.meta.ai. Nothing is read until you paste one.", + None, + ), + cookie_option( + lang, + "off", + "", + "", + Some("Browser sessions are never read for Muse Code."), + ), + ], "replicate" => vec![ cookie_option( lang, diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..881ec90146 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -90,6 +90,12 @@ pub(crate) fn build_fetch_context( let has_opencodego_api_key = id == ProviderId::OpenCodeGo && api_key.as_deref().is_some_and(|key| !key.trim().is_empty()); + // Providers whose cookies only enrich an API result keep the configured + // usage source. Off and the default manual-without-cookie state never read + // a browser; only Automatic imports one (`browser_cookie_import`). + let cookies_only_enrich_usage = provider.cookies_only_enrich_usage(); + let browser_cookie_import = + cookies_only_enrich_usage && matches!(cookie_source, "auto" | "browser" | "web"); let (mut source_mode, mut cookie_header, fails_closed_without_cookie) = if id.cookie_domain().is_none() { let source_mode = if active_token_env.is_some() { @@ -98,6 +104,11 @@ pub(crate) fn build_fetch_context( usage_source }; (source_mode, None, false) + } else if cookies_only_enrich_usage { + let cookie_header = (cookie_source == "manual") + .then(|| active_token_cookie.or(stored_cookie)) + .flatten(); + (usage_source, cookie_header, false) } else { match cookie_source { // #433: an explicitly selected, non-empty Claude manual cookie is @@ -180,7 +191,11 @@ pub(crate) fn build_fetch_context( // historically mapped "manual + no cookie" to Cli, which surfaces as // "Source mode 'Cli' not supported". Remap to Web and try browser cookies // unless the user explicitly disabled cookies ("off"). - if source_mode == SourceMode::Cli && cookie_source != "off" && !provider.supports_cli() { + if source_mode == SourceMode::Cli + && cookie_source != "off" + && !provider.supports_cli() + && !cookies_only_enrich_usage + { if cookie_header .as_deref() .map(str::trim) @@ -229,6 +244,7 @@ pub(crate) fn build_fetch_context( api_region: (!api_region.is_empty()).then_some(api_region), gateway_url, auto_prefer_web, + browser_cookie_import, ..FetchContext::default() } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index d25a0a064a..26895526eb 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -1952,3 +1952,71 @@ fn bootstrap_payload_exposes_every_provider_variant() { assert!(encoded.contains("\"providers\"")); assert!(encoded.contains("\"settings\"")); } + +fn muse_fetch_context(cookie_source: Option<&str>, stored_cookie: Option<&str>) -> FetchContext { + let mut settings = Settings::default(); + if let Some(source) = cookie_source { + settings.set_cookie_source(ProviderId::Muse, source); + } + let mut cookies = ManualCookies::default(); + if let Some(header) = stored_cookie { + cookies.set(ProviderId::Muse.cli_name(), header); + } + super::build_fetch_context( + ProviderId::Muse, + &settings, + &cookies, + &ApiKeys::default(), + &HashMap::new(), + ) +} + +#[test] +fn muse_default_cookie_source_reads_no_browser_and_keeps_the_login_source() { + let ctx = muse_fetch_context(None, None); + + assert!(!ctx.browser_cookie_import); + assert!(ctx.manual_cookie_header.is_none()); + assert_eq!(ctx.source_mode, SourceMode::Auto); +} + +#[test] +fn muse_cookie_source_off_never_reads_or_forwards_a_cookie() { + let ctx = muse_fetch_context(Some("off"), Some("llama_dev_sess=abc")); + + assert!(!ctx.browser_cookie_import); + assert!(ctx.manual_cookie_header.is_none()); +} + +#[test] +fn muse_automatic_cookie_source_requests_the_browser_import() { + let ctx = muse_fetch_context(Some("auto"), Some("llama_dev_sess=abc")); + + assert!(ctx.browser_cookie_import); + assert!(ctx.manual_cookie_header.is_none()); + assert_eq!(ctx.source_mode, SourceMode::Auto); +} + +#[test] +fn muse_manual_cookie_source_forwards_only_the_pasted_header() { + let ctx = muse_fetch_context(Some("manual"), Some("llama_dev_sess=abc")); + + assert!(!ctx.browser_cookie_import); + assert_eq!( + ctx.manual_cookie_header.as_deref(), + Some("llama_dev_sess=abc") + ); +} + +#[test] +fn other_providers_never_request_the_muse_browser_import() { + let ctx = super::build_fetch_context( + ProviderId::Cursor, + &Settings::default(), + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + + assert!(!ctx.browser_cookie_import); +} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..66018d452b 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -754,6 +754,10 @@ export const ALL_LOCALE_KEYS = [ "OpenCodeGoWorkspaceTitle", "OpenCodeGoWorkspaceLabel", "OpenCodeGoWorkspaceHelp", + "MuseBrowserTeamTitle", + "MuseBrowserTeamChoose", + "MuseBrowserTeamHelp", + "MuseBrowserTeamNoTeams", "OpenRouterManagementKeyTitle", "OpenRouterManagementKeyLabel", "OpenRouterManagementKeyHelp", diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx index 0dadfe492d..b2b07c9a09 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProviderDetailPane.tsx @@ -48,6 +48,7 @@ import { CredentialStorageSection } from "./sections/CredentialStorageSection"; import { CredentialsDispatcher } from "./sections/CredentialsDispatcher"; import { WayfinderGatewaySection } from "./sections/WayfinderGatewaySection"; import { AzureApiVersionSection } from "./sections/AzureApiVersionSection"; +import { MuseBrowserTeamSection } from "./sections/MuseBrowserTeamSection"; interface Props { providerId: string | null; @@ -370,6 +371,15 @@ export function ProviderDetailPane({ onChanged={reload} /> )} + {detail.id === "muse" && ( + + )} ({ + getProviderWorkspaceId: vi.fn(), + setProviderWorkspaceId: vi.fn(), +})); + +function row(id: string, title: string, value: string): ProviderDisplayDetail { + return { id, title, value, secondaryValue: null, progress: null }; +} + +const DETAILS = [ + row("browser-teams-status", "Browser teams", "Choose a browser team ID in Muse Code settings"), + row("browser-team-11", "Alpha", "11"), + row("browser-team-22", "Beta", "22"), +]; + +function renderSection(details: ProviderDisplayDetail[] | undefined = DETAILS) { + const onChanged = vi.fn(); + render( + key} + onChanged={onChanged} + />, + ); + return onChanged; +} + +describe("MuseBrowserTeamSection", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getProviderWorkspaceId).mockResolvedValue(null); + vi.mocked(setProviderWorkspaceId).mockResolvedValue(undefined); + }); + + it("lists only team rows and never preselects the first team", async () => { + renderSection(); + const select = (await screen.findByRole("combobox")) as HTMLSelectElement; + + expect(select.value).toBe(""); + expect(Array.from(select.options).map((option) => option.value)).toEqual([ + "", + "11", + "22", + ]); + expect(select.options[2].textContent).toBe("Beta (22)"); + expect(screen.getByText("MuseBrowserTeamHelp")).toBeInTheDocument(); + }); + + it("persists the chosen team as the provider workspace value", async () => { + const onChanged = renderSection(); + const select = await screen.findByRole("combobox"); + + await act(async () => { + fireEvent.change(select, { target: { value: "22" } }); + }); + + await vi.waitFor(() => + expect(setProviderWorkspaceId).toHaveBeenCalledWith("muse", "22"), + ); + expect(onChanged).toHaveBeenCalledOnce(); + }); + + it("keeps a saved team visible when the session no longer lists it", async () => { + vi.mocked(getProviderWorkspaceId).mockResolvedValue("99"); + renderSection([]); + + const select = (await screen.findByRole("combobox")) as HTMLSelectElement; + await vi.waitFor(() => expect(select.value).toBe("99")); + expect(screen.getByText("MuseBrowserTeamNoTeams")).toBeInTheDocument(); + }); + + it("reports a failed save without changing the selection", async () => { + vi.mocked(setProviderWorkspaceId).mockRejectedValue("Muse browser team ID is invalid"); + const onChanged = renderSection(); + const select = (await screen.findByRole("combobox")) as HTMLSelectElement; + + await act(async () => { + fireEvent.change(select, { target: { value: "11" } }); + }); + + expect(await screen.findByText("Muse browser team ID is invalid")).toBeInTheDocument(); + expect(onChanged).not.toHaveBeenCalled(); + expect(select.value).toBe(""); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/MuseBrowserTeamSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MuseBrowserTeamSection.tsx new file mode 100644 index 0000000000..02953b4cde --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/MuseBrowserTeamSection.tsx @@ -0,0 +1,109 @@ +import { useEffect, useMemo, useState } from "react"; +import type { LocaleKey } from "../../../../i18n/keys"; +import { + getProviderWorkspaceId, + setProviderWorkspaceId, +} from "../../../../lib/tauri"; +import type { ProviderDisplayDetail } from "../../../../types/bridge"; + +/** Detail-row id prefix the Muse provider uses for each listed browser team. */ +export const MUSE_BROWSER_TEAM_ROW_PREFIX = "browser-team-"; + +interface Props { + providerId: string; + details: ProviderDisplayDetail[] | undefined; + disabled: boolean; + t: (key: LocaleKey) => string; + onChanged: () => void; +} + +interface TeamOption { + id: string; + label: string; +} + +function listedTeams(details: ProviderDisplayDetail[] | undefined): TeamOption[] { + return (details ?? []) + .filter((row) => row.id.startsWith(MUSE_BROWSER_TEAM_ROW_PREFIX)) + .map((row) => ({ id: row.value, label: `${row.title} (${row.value})` })); +} + +/** + * Picker for the dev.meta.ai team whose quota Muse Code reads when the + * device login reports none. The choice is stored as the provider workspace + * value; an empty choice never falls back to the first listed team. + */ +export function MuseBrowserTeamSection({ + providerId, + details, + disabled, + t, + onChanged, +}: Props) { + const [selected, setSelected] = useState(""); + const [busy, setBusy] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + let stale = false; + setError(null); + void getProviderWorkspaceId(providerId) + .then((value) => { + if (!stale) setSelected(value ?? ""); + }) + .catch((reason: unknown) => { + if (!stale) setError(String(reason)); + }); + return () => { + stale = true; + }; + }, [providerId]); + + const teams = useMemo(() => listedTeams(details), [details]); + const options = useMemo( + () => + !selected || teams.some((team) => team.id === selected) + ? teams + : [...teams, { id: selected, label: selected }], + [teams, selected], + ); + + const handleChange = async (next: string) => { + if (next === selected || busy || disabled) return; + setBusy(true); + setError(null); + try { + await setProviderWorkspaceId(providerId, next); + setSelected(next); + onChanged(); + } catch (reason: unknown) { + setError(String(reason)); + } finally { + setBusy(false); + } + }; + + return ( +
+

{t("MuseBrowserTeamTitle")}

+ +

+ {t(teams.length === 0 ? "MuseBrowserTeamNoTeams" : "MuseBrowserTeamHelp")} +

+ {error &&

{error}

} +
+ ); +} diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 237f7dab3d..7536f430b2 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -55,6 +55,29 @@ spend and, when the optional balance request succeeds, prepaid credit balance. It keeps those values in the cost/detail surfaces and does not invent a quota percentage or use a Replicate API token as a website credential. +### Muse Code browser team quota + +Muse Code reads its quota from the device-code login. When that response has no +`subs_usage`, CodexBar can optionally read the quota of one dev.meta.ai team you +select in Settings → Providers → Muse Code → **Browser team**. + +- The cookie source is **Off** by default. The Windows default ("Manual" with no + pasted cookie) also reads nothing, and never touches a browser. +- **Automatic** imports the `llama_dev_sess` cookie from the selected browser. + **Manual** accepts a pasted Cookie header or cURL capture. +- Requests, in order: `GET /api/auth/me` (its email must match the login email, + case-insensitively), `GET /api/portal/teams`, then + `GET /api/portal/teams/{id}/subscription-quota` for the selected team only. The + team's `tier` must equal the login plan. At most 5 requests, 8 seconds each. +- The team list appears as the picker options and starts at "Choose a team..."; + the first team is never selected for you. The team ID is stored as the provider + workspace value and accepts digits only. +- The source label becomes `oauth+web`. This port has no `estimated` data + confidence, so the pace shown for this reading is marked non-authoritative + instead. The device-code token is only sent to `api.meta.ai`; dev.meta.ai + requests carry only the browser cookie. +- The CLI does not read browser cookies for Muse Code. + ## Listing what is enabled ```powershell diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 64d0b54d59..a548059079 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -188,6 +188,7 @@ async fn collect_provider_diagnostic( .provider_config(provider_id) .and_then(|config| config.gateway_url.clone()), auto_prefer_web: false, + browser_cookie_import: false, // Diagnostics keep the short optional-join grace (upstream #2583 gate). requires_optional_usage_completeness: false, }; diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index b9ecf5dce8..e04a369c1f 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -323,6 +323,7 @@ async fn fetch_guard_outcome( api_region: None, gateway_url: None, auto_prefer_web: false, + browser_cookie_import: false, // Guard checks keep the short optional-join grace. requires_optional_usage_completeness: false, }; diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 2f26498cf1..7884d8c6f9 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -301,6 +301,7 @@ async fn hooks_watch_observation( api_region: (!region.is_empty()).then(|| region.to_string()), gateway_url: (!gateway.is_empty()).then(|| gateway.to_string()), auto_prefer_web: false, + browser_cookie_import: false, // Hook watches keep the short optional-join grace. requires_optional_usage_completeness: false, }; diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 8d28758df4..39ea263fb2 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -164,6 +164,7 @@ async fn fetch_provider_envelope( api_region: None, gateway_url: None, auto_prefer_web: false, + browser_cookie_import: false, requires_optional_usage_completeness: false, }; let fetch = bounded_fetch(provider_id, ctx, None, fetch_timeout).await; @@ -297,6 +298,7 @@ async fn collect_claude_accounts(claude_enabled: bool) -> Option) -> String { api_region: None, gateway_url: None, auto_prefer_web: false, + browser_cookie_import: false, // Serve `/usage` is a background poll read: keep the short optional- // join grace (upstream #2583), unlike `codexbar usage` which blocks // for the full completeness window. diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 2b1dfff0af..1160009861 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -254,6 +254,7 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch api_region: None, gateway_url: None, auto_prefer_web: false, + browser_cookie_import: false, // `codexbar usage` is a foreground read: optional enrichment (e.g. the // OpenCode Go Zen balance) is worth its full bounded wait (#2583). requires_optional_usage_completeness: true, diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..137ff8b913 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -413,7 +413,7 @@ impl ProviderId { ProviderId::DeepInfra => None, ProviderId::Fireworks => None, ProviderId::Meta => None, - ProviderId::Muse => None, + ProviderId::Muse => Some("dev.meta.ai"), ProviderId::Nous => None, ProviderId::AiAnd => None, ProviderId::Windsurf => None, @@ -771,6 +771,12 @@ pub struct FetchContext { /// manual cookie source, etc.). Workspace overrides are checked separately. pub auto_prefer_web: bool, + /// The user chose automatic browser cookie import for a provider whose + /// cookies only enrich its API usage (`Provider::cookies_only_enrich_usage`). + /// False for the default manual-without-cookie state and for Off, so those + /// never read a browser. + pub browser_cookie_import: bool, + /// Foreground usage reads (`codexbar usage`, `codexbar serve`) set this so /// providers join slow optional enrichment with the full optional-item /// timeout budget measured from task start; background/UI polls keep the @@ -794,6 +800,7 @@ impl Default for FetchContext { api_region: None, gateway_url: None, auto_prefer_web: false, + browser_cookie_import: false, requires_optional_usage_completeness: false, } } @@ -892,6 +899,15 @@ pub trait Provider: Send + Sync { true } + /// Whether cookies only enrich an API-backed result instead of being a + /// usage source of their own. The shell then keeps the configured usage + /// source, forwards a manual cookie as-is, and reads a browser only when + /// the cookie source is Automatic (`FetchContext::browser_cookie_import`); + /// the default manual state with no cookie never triggers a browser read. + fn cookies_only_enrich_usage(&self) -> bool { + false + } + /// Whether browser-cookie discovery/recovery is owned by the provider. fn owns_browser_cookie_resolution(&self) -> bool { false @@ -1450,7 +1466,7 @@ mod tests { fn test_provider_id_muse() { assert_eq!(ProviderId::Muse.cli_name(), "muse"); assert_eq!(ProviderId::Muse.display_name(), "Muse Code"); - assert_eq!(ProviderId::Muse.cookie_domain(), None); + assert_eq!(ProviderId::Muse.cookie_domain(), Some("dev.meta.ai")); assert_eq!(ProviderId::from_cli_name("muse"), Some(ProviderId::Muse)); assert_eq!( ProviderId::from_cli_name("muse-code"), diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..a48802199b 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -1033,6 +1033,10 @@ locale_keys! { OpenCodeGoWorkspaceTitle, OpenCodeGoWorkspaceLabel, OpenCodeGoWorkspaceHelp, + MuseBrowserTeamTitle, + MuseBrowserTeamChoose, + MuseBrowserTeamHelp, + MuseBrowserTeamNoTeams, OpenRouterManagementKeyTitle, OpenRouterManagementKeyLabel, OpenRouterManagementKeyHelp, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..ee7876219d 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -659,6 +659,10 @@ CredentialsSectionTitle = Credentials OpenCodeGoWorkspaceTitle = OpenCode Go OpenCodeGoWorkspaceLabel = Workspace ID OpenCodeGoWorkspaceHelp = Optional. Use a workspace ID when automatic detection fails. +MuseBrowserTeamTitle = Browser team +MuseBrowserTeamChoose = Choose a team… +MuseBrowserTeamHelp = Used only when the login response has no quota. Turn on the browser cookie source above, refresh Muse Code, then choose the dev.meta.ai team whose quota to show. +MuseBrowserTeamNoTeams = No teams listed yet. Refresh Muse Code after turning on the browser cookie source. CredsStatusAuthenticated = Authenticated CredsStatusNotSignedIn = Not signed in CredsStatusDetected = Detected diff --git a/rust/src/providers/muse/mod.rs b/rust/src/providers/muse/mod.rs index 940981888f..bda9bb4b62 100644 --- a/rust/src/providers/muse/mod.rs +++ b/rust/src/providers/muse/mod.rs @@ -6,6 +6,7 @@ //! not for inference keys (`LLM_` / `LLM|`). pub mod local_usage; +mod web_quota; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -78,7 +79,7 @@ impl MuseProvider { } } - async fn fetch_api(&self) -> Result { + async fn fetch_api(&self, ctx: &FetchContext) -> Result { let token = resolve_device_token()?; let response = timeout( REQUEST_TIMEOUT, @@ -102,7 +103,15 @@ impl MuseProvider { // read_bounded_body, which also covers unknown or lying // content-length headers; no separate precheck is needed. let body = read_bounded_body(response).await?; - parse_response(&body) + match parse_login_response(&body)? { + LoginResponse::Complete(result) => Ok(*result), + // The login omits quotas while the 5-hour window is idle. The + // selected dev.meta.ai team can fill them in when the user opted in. + LoginResponse::WithoutQuota(identity) => { + let reading = web_quota::fetch_reading(ctx, &identity).await; + Ok(web_quota::windowless_result(&identity, reading)) + } + } } } @@ -144,7 +153,7 @@ impl Provider for MuseProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => self.fetch_api().await, + SourceMode::Auto | SourceMode::OAuth => self.fetch_api(ctx).await, SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) } @@ -158,6 +167,12 @@ impl Provider for MuseProvider { fn supports_oauth(&self) -> bool { true } + + /// dev.meta.ai cookies only fill in quotas the login response omits; the + /// device-code login stays the usage source. + fn cookies_only_enrich_usage(&self) -> bool { + true + } } fn resolve_device_token() -> Result { @@ -247,7 +262,22 @@ fn status_error(status: StatusCode) -> ProviderError { } } +/// A parsed login response: a complete result, or an active subscription +/// whose response carried no quota windows. +enum LoginResponse { + Complete(Box), + WithoutQuota(web_quota::LoginIdentity), +} + +#[cfg(test)] fn parse_response(body: &[u8]) -> Result { + match parse_login_response(body)? { + LoginResponse::Complete(result) => Ok(*result), + LoginResponse::WithoutQuota(identity) => Ok(web_quota::windowless_result(&identity, None)), + } +} + +fn parse_login_response(body: &[u8]) -> Result { let decoded: Value = serde_json::from_slice(body).map_err(|_| { ProviderError::Parse( "Could not parse Muse Code subscription usage: expected JSON".to_string(), @@ -270,24 +300,10 @@ fn parse_response(body: &[u8]) -> Result { let plan = optional_text(root.get("subs_tier_name"), "subs_tier_name")?; let email = optional_text(root.get("user_email"), "user_email")?; let Some(raw_usage) = root.get("subs_usage").filter(|value| !value.is_null()) else { - let mut usage = UsageSnapshot::new(RateWindow::informational( - "Subscription active; quota was not included in this login response", - )) - .with_login_method("Muse login"); - if let Some(email) = email { - usage = usage.with_email(email); - } - let mut result = ProviderFetchResult::new(usage, "oauth") - .with_non_authoritative_pace() - .with_display_detail(ProviderDisplayDetail::new( - "quota", - "Quota", - "Not included in this login response", - )); - if let Some(plan) = plan { - result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", plan)); - } - return Ok(result); + return Ok(LoginResponse::WithoutQuota(web_quota::LoginIdentity { + email, + plan, + })); }; let usage = object(raw_usage, "subs_usage")?; let window = object( @@ -358,7 +374,7 @@ fn parse_response(body: &[u8]) -> Result { if let Some(plan) = plan { result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", plan)); } - Ok(result) + Ok(LoginResponse::Complete(Box::new(result))) } fn object<'a>( diff --git a/rust/src/providers/muse/web_quota/mod.rs b/rust/src/providers/muse/web_quota/mod.rs new file mode 100644 index 0000000000..388eeedd19 --- /dev/null +++ b/rust/src/providers/muse/web_quota/mod.rs @@ -0,0 +1,665 @@ +//! Selected-team quota read from the dev.meta.ai browser session. +//! +//! The Muse login response omits `subs_usage` while the 5-hour window is idle. +//! When the user opted in (automatic browser import or a pasted `llama_dev_sess` +//! cookie) and chose a browser team, the same subscription quota is read from +//! the dev.meta.ai portal. A session can see several teams and nothing in the +//! responses links a team to the CLI login, so the team is never guessed. +//! +//! Wire shapes come from upstream CodexBar v0.68.0 (`muse.ts`, #4011). The +//! device-code token is never sent to dev.meta.ai; those requests carry only +//! the session cookie. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use reqwest::Client; +use serde::Deserialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use tokio::time::{Duration, timeout}; + +use super::{MAX_RESET_SECONDS, WEEKLY_WINDOW_MINUTES}; +use crate::browser::cookies::{CookieError, get_cookie_headers_for_domain}; +use crate::core::{ + FetchContext, ProviderDisplayDetail, ProviderFetchResult, RateWindow, UsageSnapshot, +}; +use crate::providers::{cookie_values, normalize_cookie_header, read_bounded_response}; + +#[cfg(test)] +mod tests; + +const DEV_META_ORIGIN: &str = "https://dev.meta.ai"; +const DEV_META_DOMAIN: &str = "dev.meta.ai"; +const SESSION_COOKIE: &str = "llama_dev_sess"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(8); +/// Upstream budget: five requests across every candidate session. +const REQUEST_BUDGET: u8 = 5; +/// The 15 s login request plus this deadline stay inside the 60 s Muse fetch. +const WEB_DEADLINE: Duration = Duration::from_secs(45); +const MAX_BODY_BYTES: usize = 256 * 1024; +const MIN_WINDOW_SECONDS: f64 = 60.0; +const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991; + +/// Identity from the login response that the browser session must match. +#[derive(Debug, Clone, Default)] +pub(super) struct LoginIdentity { + pub(super) email: Option, + pub(super) plan: Option, +} + +/// How the fetch context allows dev.meta.ai to be contacted. +#[derive(Debug, Clone, PartialEq, Eq)] +enum CookieAccess { + /// No cookie was configured: never read a browser or contact dev.meta.ai. + Off, + /// A pasted Cookie header or cURL capture. + Manual(String), + /// Import the session from an installed browser. + Browser, +} + +impl CookieAccess { + fn from_context(ctx: &FetchContext) -> Self { + match ctx.manual_cookie_header.as_deref() { + Some(raw) if !raw.trim().is_empty() => Self::Manual(raw.to_string()), + _ if ctx.browser_cookie_import => Self::Browser, + _ => Self::Off, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct BrowserTeam { + id: String, + name: String, +} + +/// Why no quota is shown even though a browser session was available. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TeamNote { + ChooseTeam, + TeamNotVisible, + NoQuota, + PlanDiffers, + CookiesProtected, +} + +impl TeamNote { + fn message(self) -> &'static str { + match self { + Self::ChooseTeam => "Choose a browser team in Muse Code settings", + Self::TeamNotVisible => "The selected browser team is not visible to this session", + Self::NoQuota => "No subscription quota for the selected team", + Self::PlanDiffers => "The selected team's plan differs from the Muse login", + Self::CookiesProtected => { + "Browser cookies are protected by app-bound encryption; paste the Cookie header manually" + } + } + } +} + +#[derive(Debug, Clone)] +struct TeamQuota { + team: BrowserTeam, + primary: RateWindow, + secondary: RateWindow, +} + +/// What the browser session revealed. `None` from the readers means the +/// browser path contributed nothing and the login-only result stands. +#[derive(Debug, Clone, Default)] +pub(super) struct WebReading { + teams: Vec, + note: Option, + quota: Option, +} + +/// One dev.meta.ai GET. The transport is a trait so the request budget and +/// session handling are testable without a network. +#[async_trait] +trait DevMetaApi: Send + Sync { + async fn get(&self, path: &str, session_cookie: &str) -> Result; +} + +struct DevMetaResponse { + status: u16, + body: Vec, +} + +/// A transport or decoding failure that ends the whole browser path, like an +/// exception in the upstream plugin. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Abort; + +struct ReqwestDevMeta { + client: Client, +} + +impl ReqwestDevMeta { + fn new() -> Self { + Self { + client: crate::core::credentialed_http_client_builder() + .timeout(REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()), + } + } +} + +#[async_trait] +impl DevMetaApi for ReqwestDevMeta { + async fn get(&self, path: &str, session_cookie: &str) -> Result { + let response = self + .client + .get(format!("{DEV_META_ORIGIN}{path}")) + .header("Cookie", session_cookie) + .header("User-Agent", "CodexBar") + .send() + .await + .map_err(|_| Abort)?; + let status = response.status().as_u16(); + if status != 200 { + return Ok(DevMetaResponse { + status, + body: Vec::new(), + }); + } + let body = read_bounded_response(response, MAX_BODY_BYTES) + .await + .map_err(|_| Abort)?; + Ok(DevMetaResponse { status, body }) + } +} + +/// Read the selected team's quota, or `None` when the browser path is off or +/// yields nothing usable. Never returns an error: the login result stands. +pub(super) async fn fetch_reading( + ctx: &FetchContext, + identity: &LoginIdentity, +) -> Option { + let access = CookieAccess::from_context(ctx); + if access == CookieAccess::Off { + return None; + } + let team_id = ctx + .workspace_id + .as_deref() + .map_or("", str::trim) + .to_string(); + let identity = identity.clone(); + timeout(WEB_DEADLINE, async move { + let sessions = resolve_sessions(access).await; + if sessions.cookies.is_empty() { + return sessions.protected.then(|| WebReading { + note: Some(TeamNote::CookiesProtected), + ..WebReading::default() + }); + } + read_team_quota( + &ReqwestDevMeta::new(), + &sessions.cookies, + &identity, + &team_id, + Utc::now(), + ) + .await + }) + .await + .unwrap_or_else(|_| { + tracing::debug!("Muse browser team quota timed out"); + None + }) +} + +struct Sessions { + /// Cookie headers reduced to `llama_dev_sess=`. + cookies: Vec, + /// Chromium app-bound encryption hid every browser cookie. + protected: bool, +} + +impl Sessions { + fn none() -> Self { + Self { + cookies: Vec::new(), + protected: false, + } + } +} + +async fn resolve_sessions(access: CookieAccess) -> Sessions { + match access { + CookieAccess::Off => Sessions::none(), + CookieAccess::Manual(raw) => Sessions { + cookies: manual_session_cookie(&raw).into_iter().collect(), + protected: false, + }, + CookieAccess::Browser => { + let scan = + tokio::task::spawn_blocking(|| get_cookie_headers_for_domain(DEV_META_DOMAIN)) + .await; + match scan { + Ok(Ok(headers)) => { + let mut cookies: Vec = Vec::new(); + for cookie in headers + .iter() + .filter_map(|(_, header)| session_cookie(header)) + { + if !cookies.contains(&cookie) { + cookies.push(cookie); + } + } + Sessions { + cookies, + protected: false, + } + } + Ok(Err(CookieError::AppBoundEncryption)) => Sessions { + cookies: Vec::new(), + protected: true, + }, + _ => Sessions::none(), + } + } + } +} + +/// Accepts a bare Cookie header, a `Cookie:` line, or a DevTools cURL capture. +fn manual_session_cookie(raw: &str) -> Option { + let header = if crate::core::looks_like_curl_capture(raw) { + let fields = crate::core::header_fields(raw); + crate::core::header_value("cookie", &fields)? + } else { + raw.to_string() + }; + session_cookie(&normalize_cookie_header(&header)?) +} + +/// Reduce a Cookie header to the session cookie alone; anything else the +/// browser holds for dev.meta.ai is not forwarded. Duplicate values are +/// ambiguous, so they are rejected rather than picking one. +fn session_cookie(header: &str) -> Option { + let values = cookie_values(header, SESSION_COOKIE); + let [value] = values.as_slice() else { + return None; + }; + value + .chars() + .all(|c| c.is_ascii_graphic()) + .then(|| format!("{SESSION_COOKIE}={value}")) +} + +struct Budget(u8); + +/// Requests for one session. `rejected` is set on 401/403 so the caller moves +/// on to the next session. +struct SessionClient<'a> { + api: &'a dyn DevMetaApi, + cookie: &'a str, + budget: &'a mut Budget, + rejected: bool, +} + +impl SessionClient<'_> { + /// `Ok(None)` for a rejected session, a non-200 status, or a non-object + /// body; `Err` for an exhausted budget, transport failure, or invalid JSON. + async fn get(&mut self, path: &str) -> Result, Abort> { + if self.budget.0 == 0 { + return Err(Abort); + } + self.budget.0 -= 1; + let response = self.api.get(path, self.cookie).await?; + if matches!(response.status, 401 | 403) { + self.rejected = true; + return Ok(None); + } + if response.status != 200 { + return Ok(None); + } + let value: Value = serde_json::from_slice(&response.body).map_err(|_| Abort)?; + if !value.is_object() { + return Ok(None); + } + Ok(serde_json::from_value(value).ok()) + } +} + +/// Deserialize a field leniently: a value of the wrong type is treated as absent. +fn lenient<'de, D, T>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, + T: DeserializeOwned, +{ + let value = Value::deserialize(deserializer)?; + Ok(serde_json::from_value(value).ok()) +} + +#[derive(Deserialize)] +struct MeResponse { + #[serde(default, deserialize_with = "lenient")] + email: Option, +} + +#[derive(Deserialize)] +struct TeamsResponse { + #[serde(default, deserialize_with = "lenient")] + teams: Option>, +} + +#[derive(Deserialize, Default)] +struct TeamEntry { + #[serde(default, deserialize_with = "lenient")] + team_id: Option, + #[serde(default, deserialize_with = "lenient")] + team_name: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum TeamId { + Text(String), + Number(u64), +} + +impl TeamEntry { + /// Entries without a digits-only id are ignored, as upstream does. + fn into_team(self) -> Option { + let id = match self.team_id? { + TeamId::Text(text) => text, + TeamId::Number(number) if number <= MAX_SAFE_INTEGER => number.to_string(), + TeamId::Number(_) => return None, + }; + if id.is_empty() || !id.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + let name = self + .team_name + .map(|name| name.trim().to_string()) + .filter(|name| !name.is_empty()) + .unwrap_or_else(|| id.clone()); + Some(BrowserTeam { id, name }) + } +} + +#[derive(Deserialize)] +struct QuotaResponse { + #[serde(default)] + subscription_quota: Value, +} + +/// Limits and usage are weighted token counts encoded as decimal strings. +/// Every field stays a raw `Value` because null, absent, and malformed are +/// different cases for the reset times. +#[derive(Deserialize, Default)] +struct QuotaRecord { + #[serde(default)] + tier: Value, + #[serde(default)] + window_weighted_used: Value, + #[serde(default)] + window_weighted_limit: Value, + #[serde(default)] + window_duration_secs: Value, + #[serde(default)] + window_resets_at: Value, + #[serde(default)] + weekly_weighted_used: Value, + #[serde(default)] + weekly_weighted_limit: Value, + #[serde(default)] + weekly_resets_at: Value, +} + +/// Walk the candidate sessions until one belongs to the login's account. +async fn read_team_quota( + api: &dyn DevMetaApi, + sessions: &[String], + identity: &LoginIdentity, + team_id: &str, + now: DateTime, +) -> Option { + let mut budget = Budget(REQUEST_BUDGET); + for cookie in sessions { + if budget.0 == 0 { + break; + } + let mut client = SessionClient { + api, + cookie, + budget: &mut budget, + rejected: false, + }; + match read_session(&mut client, identity, team_id, now).await { + Ok(SessionOutcome::Next) => {} + Ok(SessionOutcome::Done(reading)) => return reading, + Err(Abort) => { + tracing::debug!("Muse browser team quota read stopped"); + return None; + } + } + } + None +} + +enum SessionOutcome { + /// This session was rejected or belongs to another account. + Next, + Done(Option), +} + +async fn read_session( + client: &mut SessionClient<'_>, + identity: &LoginIdentity, + team_id: &str, + now: DateTime, +) -> Result { + // The browser session must belong to the same Meta account as the CLI login. + let me: Option = client.get("/api/auth/me").await?; + let web_email = me + .and_then(|me| me.email) + .map(|email| email.trim().to_lowercase()); + let login_email = identity.email.as_deref().map(str::to_lowercase); + if login_email.is_none() || web_email != login_email { + return Ok(SessionOutcome::Next); + } + + let listed: Option = client.get("/api/portal/teams").await?; + if client.rejected { + return Ok(SessionOutcome::Next); + } + let Some(entries) = listed.and_then(|listed| listed.teams) else { + return Ok(SessionOutcome::Done(None)); + }; + let teams: Vec = entries + .into_iter() + .filter_map(|entry| { + serde_json::from_value::(entry) + .unwrap_or_default() + .into_team() + }) + .collect(); + + let done = |teams: Vec, note: Option, quota: Option| { + Ok(SessionOutcome::Done(Some(WebReading { + teams, + note, + quota, + }))) + }; + if team_id.is_empty() { + return done(teams, Some(TeamNote::ChooseTeam), None); + } + let Some(team) = teams.iter().find(|team| team.id == team_id).cloned() else { + return done(teams, Some(TeamNote::TeamNotVisible), None); + }; + + let response: Option = client + .get(&format!("/api/portal/teams/{}/subscription-quota", team.id)) + .await?; + if client.rejected { + return Ok(SessionOutcome::Next); + } + let record = response + .map(|response| response.subscription_quota) + .filter(Value::is_object) + .and_then(|value| serde_json::from_value::(value).ok()); + let Some(record) = record else { + return done(teams, Some(TeamNote::NoQuota), None); + }; + // The team's quota must be for the same plan the CLI login reports. + let tier = record.tier.as_str().map(str::trim); + if identity.plan.is_none() || tier != identity.plan.as_deref() { + return done(teams, Some(TeamNote::PlanDiffers), None); + } + let quota = parse_team_quota(&record, now).map(|(primary, secondary)| TeamQuota { + team, + primary, + secondary, + }); + done(teams, None, quota) +} + +fn amount(value: &Value) -> Option { + let parsed = match value { + Value::String(text) if !text.is_empty() && text.bytes().all(|b| b.is_ascii_digit()) => { + text.parse::().ok()? + } + Value::Number(number) => number.as_f64()?, + _ => return None, + }; + (parsed.is_finite() && parsed >= 0.0).then_some(parsed) +} + +fn percent(used: &Value, limit: &Value) -> Option { + let (used, limit) = (amount(used)?, amount(limit)?); + (limit > 0.0).then(|| (used / limit * 100.0).min(100.0)) +} + +fn reset_at(value: &Value) -> Option> { + let seconds = amount(value)?; + if seconds <= 0.0 || seconds > MAX_RESET_SECONDS { + return None; + } + #[expect( + clippy::cast_possible_truncation, + reason = "MAX_RESET_SECONDS bounds the value inside i64" + )] + let whole_seconds = seconds.trunc() as i64; + DateTime::::from_timestamp(whole_seconds, 0) +} + +/// `(primary, secondary)` windows, or `None` when the reading cannot be +/// trusted: a stale weekly quota, a missing limit, or a malformed reset shows +/// nothing rather than an old or invented value. +fn parse_team_quota(quota: &QuotaRecord, now: DateTime) -> Option<(RateWindow, RateWindow)> { + let weekly_percent = percent("a.weekly_weighted_used, "a.weekly_weighted_limit)?; + // A weekly quota without a future reset is stale. + let weekly_reset = reset_at("a.weekly_resets_at).filter(|reset| *reset > now)?; + + let seconds = amount("a.window_duration_secs)?; + if seconds.fract() != 0.0 || seconds < MIN_WINDOW_SECONDS || seconds > MAX_SAFE_INTEGER as f64 { + return None; + } + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "seconds is a validated non-negative integer; out-of-range minutes are rejected below" + )] + let window_minutes = u32::try_from((seconds / 60.0).round() as u64).ok()?; + + let window_reset = reset_at("a.window_resets_at); + let mut primary_percent = percent("a.window_weighted_used, "a.window_weighted_limit)?; + let no_window_reset = quota.window_resets_at.is_null(); + if (!no_window_reset && window_reset.is_none()) || (no_window_reset && primary_percent != 0.0) { + return None; + } + // An idle 5-hour window has no reset time; one whose reset has passed + // carries no usage into the next window. + let active_reset = window_reset.filter(|reset| *reset > now); + if active_reset.is_none() { + primary_percent = 0.0; + } + Some(( + RateWindow::with_details(primary_percent, Some(window_minutes), active_reset, None), + RateWindow::with_details( + weekly_percent, + Some(WEEKLY_WINDOW_MINUTES), + Some(weekly_reset), + None, + ), + )) +} + +/// Result for an active login whose response omitted `subs_usage`, optionally +/// filled in from the selected browser team. +pub(super) fn windowless_result( + identity: &LoginIdentity, + reading: Option, +) -> ProviderFetchResult { + let quota = reading.as_ref().and_then(|reading| reading.quota.as_ref()); + let (usage, source_label) = match quota { + Some(quota) => ( + UsageSnapshot::new(quota.primary.clone()).with_secondary(quota.secondary.clone()), + "oauth+web", + ), + None => ( + UsageSnapshot::new(RateWindow::informational( + "Subscription active; quota was not included in this login response", + )), + "oauth", + ), + }; + let mut usage = usage.with_login_method("Muse login"); + if let Some(email) = &identity.email { + usage = usage.with_email(email.clone()); + } + // Weighted usage of a user-selected team is not the login's own exact + // reading, so it never feeds pace advice. + let mut result = ProviderFetchResult::new(usage, source_label).with_non_authoritative_pace(); + if quota.is_none() { + result = result.with_display_detail(ProviderDisplayDetail::new( + "quota", + "Quota", + "Not included in this login response", + )); + } + if let Some(plan) = &identity.plan { + result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", plan)); + } + let Some(reading) = reading else { + return result; + }; + if let Some(note) = reading.note { + result = result.with_display_detail(ProviderDisplayDetail::new( + "browser-teams-status", + "Browser teams", + note.message(), + )); + } + for team in &reading.teams { + result = result.with_display_detail(ProviderDisplayDetail::new( + format!("browser-team-{}", team.id), + &team.name, + &team.id, + )); + } + if let Some(quota) = &reading.quota { + result = result + .with_display_detail(ProviderDisplayDetail::new( + "browser-team", + "Browser team quota (dev.meta.ai)", + "a.team.name, + )) + .with_display_detail(ProviderDisplayDetail::new( + "five-hour", + "5 hours", + super::format_percent(quota.primary.used_percent), + )) + .with_display_detail(ProviderDisplayDetail::new( + "weekly", + "Weekly", + super::format_percent(quota.secondary.used_percent), + )); + } + result +} diff --git a/rust/src/providers/muse/web_quota/tests.rs b/rust/src/providers/muse/web_quota/tests.rs new file mode 100644 index 0000000000..69a753a8a8 --- /dev/null +++ b/rust/src/providers/muse/web_quota/tests.rs @@ -0,0 +1,703 @@ +//! Fixtures follow the upstream Muse plugin test matrix (CodexBar v0.68.0, +//! `MusePluginTests.swift`): same paths, field names, and team/plan values. + +use super::*; +use std::sync::Mutex; + +const TEAM_ID: &str = "424242424242"; +const NOW: i64 = 1_790_341_873; +const FIXTURE_COOKIE: &str = "llama_dev_sess=fixture"; +const TEAMS: &str = r#"{"teams":[{"team_id":424242424242,"team_name":"My Team"}]}"#; +const ME: &str = r#"{"userId":"1","email":"Ada@Example.com","accountType":"META_ACCOUNT"}"#; + +fn now() -> DateTime { + DateTime::from_timestamp(NOW, 0).unwrap() +} + +fn identity() -> LoginIdentity { + LoginIdentity { + email: Some("ada@example.com".to_string()), + plan: Some("Muse Code Power Usage".to_string()), + } +} + +/// A dev.meta.ai team quota; the login fixture reports "Muse Code Power Usage". +fn quota( + tier: &str, + window_used: &str, + window_resets_at: Option, + weekly_used: &str, +) -> String { + quota_with_weekly_reset( + tier, + window_used, + window_resets_at, + weekly_used, + 1_790_553_600, + ) +} + +fn quota_with_weekly_reset( + tier: &str, + window_used: &str, + window_resets_at: Option, + weekly_used: &str, + weekly_resets_at: i64, +) -> String { + let window_reset = window_resets_at + .map(|value| format!(r#","window_resets_at":{value}"#)) + .unwrap_or_default(); + format!( + r#"{{"subscription_quota":{{"tier_id":"1","tier":"{tier}","as_of":{NOW},"window_weighted_limit":"20000000000","window_duration_secs":18000,"weekly_weighted_limit":"60000000000","weekly_resets_at":{weekly_resets_at},"window_weighted_used":"{window_used}","weekly_weighted_used":"{weekly_used}"{window_reset}}}}}"# + ) +} + +fn idle_quota() -> String { + quota("Muse Code Power Usage", "0", None, "9043782620") +} + +type Handler = Box (u16, String) + Send + Sync>; + +/// Serves canned dev.meta.ai responses and records `(cookie, path)` per request. +struct FakeApi { + handler: Handler, + log: Mutex>, +} + +impl FakeApi { + fn new(handler: impl Fn(&str, &str) -> (u16, String) + Send + Sync + 'static) -> Self { + Self { + handler: Box::new(handler), + log: Mutex::new(Vec::new()), + } + } + + /// The upstream `web(quota:)` fixture: me, teams, and the team quota. + fn with_quota(quota: String) -> Self { + Self::new(move |_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + path if path == format!("/api/portal/teams/{TEAM_ID}/subscription-quota") => { + (200, quota.clone()) + } + _ => (404, "{}".to_string()), + }) + } + + fn paths(&self) -> Vec { + self.log + .lock() + .unwrap() + .iter() + .map(|(_, path)| path.clone()) + .collect() + } + + fn cookies(&self) -> Vec { + self.log + .lock() + .unwrap() + .iter() + .map(|(cookie, _)| cookie.clone()) + .collect() + } + + fn quota_requested(&self) -> bool { + self.paths() + .iter() + .any(|path| path.ends_with("/subscription-quota")) + } +} + +#[async_trait] +impl DevMetaApi for FakeApi { + async fn get(&self, path: &str, session_cookie: &str) -> Result { + self.log + .lock() + .unwrap() + .push((session_cookie.to_string(), path.to_string())); + let (status, body) = (self.handler)(session_cookie, path); + Ok(DevMetaResponse { + status, + body: body.into_bytes(), + }) + } +} + +async fn read(api: &FakeApi, team_id: &str) -> Option { + read_team_quota( + api, + &[FIXTURE_COOKIE.to_string()], + &identity(), + team_id, + now(), + ) + .await +} + +fn detail_values(result: &ProviderFetchResult) -> Vec<(String, String, String)> { + result + .display_details() + .iter() + .map(|row| { + ( + row.id().to_string(), + row.title().to_string(), + row.value().to_string(), + ) + }) + .collect() +} + +fn has_detail(result: &ProviderFetchResult, id: &str) -> bool { + result.display_details().iter().any(|row| row.id() == id) +} + +fn detail<'a>(result: &'a ProviderFetchResult, id: &str) -> &'a ProviderDisplayDetail { + result + .display_details() + .iter() + .find(|row| row.id() == id) + .unwrap_or_else(|| panic!("missing detail {id}: {:?}", detail_values(result))) +} + +fn fallback_result(reading: Option) -> ProviderFetchResult { + windowless_result(&identity(), reading) +} + +#[tokio::test] +async fn selected_team_quota_fills_omitted_login_quotas() { + let api = FakeApi::with_quota(quota( + "Muse Code Power Usage", + "4000000000", + Some(NOW + 3600), + "9043782620", + )); + let reading = read(&api, TEAM_ID).await; + let result = fallback_result(reading); + + assert_eq!(result.source_label, "oauth+web"); + assert!(!result.pace_authoritative); + let primary = &result.usage.primary; + assert_eq!(primary.used_percent, 20.0); + assert_eq!(primary.window_minutes, Some(300)); + assert_eq!(primary.resets_at, DateTime::from_timestamp(NOW + 3600, 0)); + let weekly = result.usage.secondary.as_ref().unwrap(); + assert!((weekly.used_percent - 15.07297103).abs() < 0.0001); + assert_eq!(weekly.window_minutes, Some(10080)); + assert_eq!(weekly.resets_at, DateTime::from_timestamp(1_790_553_600, 0)); + assert_eq!(result.usage.login_method.as_deref(), Some("Muse login")); + assert_eq!( + result.usage.account_email.as_deref(), + Some("ada@example.com") + ); + + assert_eq!(detail(&result, "browser-team").value(), "My Team"); + assert_eq!( + detail(&result, "browser-team").title(), + "Browser team quota (dev.meta.ai)" + ); + assert_eq!(detail(&result, "five-hour").value(), "20%"); + assert_eq!(detail(&result, "weekly").value(), "15%"); + assert_eq!(detail(&result, "plan").value(), "Muse Code Power Usage"); + assert!(!has_detail(&result, "quota")); + assert_eq!( + detail(&result, &format!("browser-team-{TEAM_ID}")).value(), + TEAM_ID + ); + + // Only the session cookie is sent, and only to the three portal paths. + assert_eq!( + api.paths(), + vec![ + "/api/auth/me".to_string(), + "/api/portal/teams".to_string(), + format!("/api/portal/teams/{TEAM_ID}/subscription-quota"), + ] + ); + assert!(api.cookies().iter().all(|cookie| cookie == FIXTURE_COOKIE)); +} + +#[tokio::test] +async fn idle_and_expired_five_hour_windows_carry_no_usage_or_reset() { + for quota in [ + idle_quota(), + quota( + "Muse Code Power Usage", + "4000000000", + Some(NOW - 60), + "9043782620", + ), + ] { + let result = fallback_result(read(&FakeApi::with_quota(quota), TEAM_ID).await); + assert_eq!(result.source_label, "oauth+web"); + assert_eq!(result.usage.primary.used_percent, 0.0); + assert!(result.usage.primary.resets_at.is_none()); + assert_eq!(result.usage.primary.window_minutes, Some(300)); + assert!(result.usage.secondary.is_some()); + } +} + +#[tokio::test] +async fn weekly_quota_past_its_reset_withholds_the_whole_browser_reading() { + let quota = quota_with_weekly_reset( + "Muse Code Power Usage", + "4000000000", + Some(NOW + 3600), + "9043782620", + NOW - 60, + ); + let result = fallback_result(read(&FakeApi::with_quota(quota), TEAM_ID).await); + assert_eq!(result.source_label, "oauth"); + assert!(result.usage.primary.is_informational); + assert!(result.usage.secondary.is_none()); + assert!(has_detail(&result, "quota")); +} + +#[tokio::test] +async fn invalid_five_hour_resets_never_invent_an_idle_window() { + let valid = format!(r#""window_resets_at":{}"#, NOW + 3600); + for reset in ["\"invalid\"", "1e30", "-1", "true", "null", "0"] { + let body = quota( + "Muse Code Power Usage", + "4000000000", + Some(NOW + 3600), + "9043782620", + ) + .replace(&valid, &format!(r#""window_resets_at":{reset}"#)); + let result = fallback_result(read(&FakeApi::with_quota(body), TEAM_ID).await); + assert!(result.usage.secondary.is_none(), "reset {reset}"); + assert_eq!(result.source_label, "oauth", "reset {reset}"); + assert!(has_detail(&result, "quota"), "reset {reset}"); + } +} + +#[tokio::test] +async fn selected_quota_retains_the_team_list_for_settings() { + let result = fallback_result(read(&FakeApi::with_quota(idle_quota()), TEAM_ID).await); + let row = detail(&result, &format!("browser-team-{TEAM_ID}")); + assert_eq!(row.title(), "My Team"); + assert_eq!(row.value(), TEAM_ID); +} + +#[tokio::test] +async fn a_session_without_teams_reads_no_quota() { + let api = FakeApi::new(|_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, r#"{"teams":[]}"#.to_string()), + _ => (200, idle_quota()), + }); + let result = fallback_result(read(&api, TEAM_ID).await); + assert!(result.usage.secondary.is_none()); + assert!(!api.quota_requested()); + assert_eq!( + detail(&result, "browser-teams-status").value(), + "The selected browser team is not visible to this session" + ); +} + +#[tokio::test] +async fn without_a_selected_team_the_visible_teams_are_listed_and_no_quota_is_read() { + for team_id in ["", " "] { + let api = FakeApi::with_quota(idle_quota()); + let result = fallback_result( + read_team_quota( + &api, + &[FIXTURE_COOKIE.to_string()], + &identity(), + team_id.trim(), + now(), + ) + .await, + ); + assert_eq!(result.source_label, "oauth"); + assert!(result.usage.secondary.is_none()); + assert!(!api.quota_requested()); + assert_eq!( + detail(&result, "browser-team-424242424242").title(), + "My Team" + ); + assert_eq!( + detail(&result, "browser-teams-status").value(), + "Choose a browser team in Muse Code settings" + ); + } +} + +#[tokio::test] +async fn a_team_the_session_cannot_see_is_never_queried() { + let api = FakeApi::new(|_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (200, idle_quota()), + }); + let result = fallback_result(read(&api, "123").await); + assert!(result.usage.secondary.is_none()); + assert!(!api.quota_requested()); + assert!(has_detail(&result, "browser-team-424242424242")); +} + +#[tokio::test] +async fn the_selected_team_decides_the_quota_regardless_of_list_order() { + for teams in [ + r#"{"teams":[{"team_id":"11","team_name":"Alpha"},{"team_id":"22","team_name":"Beta"}]}"#, + r#"{"teams":[{"team_id":"22","team_name":"Beta"},{"team_id":"11","team_name":"Alpha"}]}"#, + ] { + let api = FakeApi::new(move |_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, teams.to_string()), + "/api/portal/teams/11/subscription-quota" => { + (200, quota("Muse Code Power Usage", "0", None, "6000000000")) + } + "/api/portal/teams/22/subscription-quota" => ( + 200, + quota("Muse Code Power Usage", "0", None, "48000000000"), + ), + _ => (404, "{}".to_string()), + }); + let result = fallback_result(read(&api, "22").await); + assert_eq!(result.usage.secondary.as_ref().unwrap().used_percent, 80.0); + assert_eq!(detail(&result, "browser-team").value(), "Beta"); + } +} + +#[tokio::test] +async fn a_team_quota_for_a_different_plan_is_not_shown() { + let api = FakeApi::with_quota(quota("Muse Code Everyday Usage", "0", None, "9043782620")); + let result = fallback_result(read(&api, TEAM_ID).await); + assert_eq!(result.source_label, "oauth"); + assert!(result.usage.secondary.is_none()); + assert_eq!( + detail(&result, "browser-teams-status").value(), + "The selected team's plan differs from the Muse login" + ); +} + +#[tokio::test] +async fn login_without_a_plan_never_matches_a_team_quota() { + let api = FakeApi::with_quota(idle_quota()); + let identity = LoginIdentity { + plan: None, + ..identity() + }; + let reading = read_team_quota( + &api, + &[FIXTURE_COOKIE.to_string()], + &identity, + TEAM_ID, + now(), + ) + .await + .unwrap(); + assert!(reading.quota.is_none()); + assert_eq!(reading.note, Some(TeamNote::PlanDiffers)); +} + +#[tokio::test] +async fn unusable_web_quotas_keep_the_login_response_result() { + let zero_limit = idle_quota().replace( + r#""window_weighted_limit":"20000000000""#, + r#""window_weighted_limit":"0""#, + ); + let huge_duration = idle_quota().replace("18000", "1e30"); + let short_duration = idle_quota().replace("18000", "59"); + let fractional_duration = idle_quota().replace("18000", "18000.5"); + let missing_weekly_limit = idle_quota().replace("weekly_weighted_limit", "weekly_other"); + for (status, body) in [ + (401_u16, r#"{"error":"Not authenticated"}"#.to_string()), + (200, r#"{"subscription_quota":null}"#.to_string()), + (200, zero_limit), + (200, "".to_string()), + (200, huge_duration), + (200, short_duration), + (200, fractional_duration), + (200, missing_weekly_limit), + (500, "{}".to_string()), + ] { + let api = FakeApi::new(move |_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (status, body.clone()), + }); + let result = fallback_result(read(&api, TEAM_ID).await); + assert_eq!(result.source_label, "oauth", "status {status}"); + assert!(result.usage.primary.is_informational, "status {status}"); + assert!(result.usage.secondary.is_none(), "status {status}"); + assert_eq!(result.usage.login_method.as_deref(), Some("Muse login")); + assert!(has_detail(&result, "quota"), "status {status}"); + } +} + +#[tokio::test] +async fn a_browser_session_for_another_account_never_supplies_quotas() { + for me in [ + r#"{"email":"bob@example.com"}"#, + r#"{"userId":"1"}"#, + r#"{"email":7}"#, + ] { + let api = FakeApi::new(move |_, path| match path { + "/api/auth/me" => (200, me.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (200, idle_quota()), + }); + let reading = read(&api, TEAM_ID).await; + assert!(reading.is_none(), "{me}"); + assert!( + api.paths() + .iter() + .all(|path| !path.starts_with("/api/portal")), + "{me}" + ); + let result = fallback_result(reading); + assert_eq!( + result.usage.account_email.as_deref(), + Some("ada@example.com") + ); + assert!(result.usage.secondary.is_none()); + } +} + +#[tokio::test] +async fn a_login_without_an_email_never_matches_a_session() { + let api = FakeApi::with_quota(idle_quota()); + let identity = LoginIdentity { + email: None, + ..identity() + }; + let reading = read_team_quota( + &api, + &[FIXTURE_COOKIE.to_string()], + &identity, + TEAM_ID, + now(), + ) + .await; + assert!(reading.is_none()); + assert!(!api.quota_requested()); +} + +#[tokio::test] +async fn rejected_and_wrong_account_sessions_advance_to_the_matching_account() { + for first_status in [200_u16, 401, 403, 500] { + let api = FakeApi::new(move |cookie, path| { + if cookie == "llama_dev_sess=first" { + assert_eq!(path, "/api/auth/me"); + return (first_status, r#"{"email":"other@example.com"}"#.to_string()); + } + match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (200, idle_quota()), + } + }); + let sessions = [ + "llama_dev_sess=first".to_string(), + "llama_dev_sess=matching".to_string(), + ]; + let reading = read_team_quota(&api, &sessions, &identity(), TEAM_ID, now()).await; + let result = fallback_result(reading); + assert_eq!(result.source_label, "oauth+web", "status {first_status}"); + assert!(result.usage.secondary.is_some()); + } +} + +#[tokio::test] +async fn a_session_rejected_after_the_email_check_advances() { + let api = FakeApi::new(|cookie, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" if cookie == "llama_dev_sess=first" => (401, "{}".to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (200, idle_quota()), + }); + let sessions = [ + "llama_dev_sess=first".to_string(), + "llama_dev_sess=second".to_string(), + ]; + let reading = read_team_quota(&api, &sessions, &identity(), TEAM_ID, now()).await; + assert_eq!(fallback_result(reading).source_label, "oauth+web"); +} + +#[tokio::test] +async fn browser_session_retries_stay_within_the_request_budget() { + let api = FakeApi::new(|_, _| (200, r#"{"email":"other@example.com"}"#.to_string())); + let sessions: Vec = (0..10) + .map(|index| format!("llama_dev_sess=s{index}")) + .collect(); + let reading = read_team_quota(&api, &sessions, &identity(), TEAM_ID, now()).await; + assert!(reading.is_none()); + assert_eq!(api.paths().len(), 5); +} + +#[tokio::test] +async fn an_exhausted_budget_stops_before_the_quota_request() { + // Four wrong-account sessions use four requests, so the matching session + // has one left for /api/auth/me and none for the team list. + let api = FakeApi::new(|cookie, path| match path { + "/api/auth/me" if cookie == "llama_dev_sess=match" => (200, ME.to_string()), + "/api/auth/me" => (200, r#"{"email":"other@example.com"}"#.to_string()), + "/api/portal/teams" => (200, TEAMS.to_string()), + _ => (200, idle_quota()), + }); + let mut sessions: Vec = (0..4) + .map(|index| format!("llama_dev_sess=s{index}")) + .collect(); + sessions.push("llama_dev_sess=match".to_string()); + let reading = read_team_quota(&api, &sessions, &identity(), TEAM_ID, now()).await; + assert!(reading.is_none()); + assert_eq!(api.paths().len(), 5); + assert!(!api.quota_requested()); +} + +#[tokio::test] +async fn a_non_array_team_list_ends_the_browser_path() { + let api = FakeApi::new(|_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, r#"{"teams":"none"}"#.to_string()), + _ => (200, idle_quota()), + }); + assert!(read(&api, TEAM_ID).await.is_none()); + assert!(!api.quota_requested()); +} + +#[tokio::test] +async fn team_entries_need_a_digits_only_id_and_fall_back_to_it_for_a_name() { + let teams = r#"{"teams":[ + {"team_id":"11","team_name":" Alpha "}, + {"team_id":22}, + {"team_id":"abc","team_name":"Bad"}, + {"team_id":-4,"team_name":"Negative"}, + {"team_id":1.5,"team_name":"Fractional"}, + {"team_name":"No id"}, + "not an object", + {"team_id":"33","team_name":7} + ]}"#; + let api = FakeApi::new(move |_, path| match path { + "/api/auth/me" => (200, ME.to_string()), + "/api/portal/teams" => (200, teams.to_string()), + _ => (404, "{}".to_string()), + }); + let reading = read(&api, "").await.unwrap(); + let listed: Vec<(String, String)> = reading + .teams + .iter() + .map(|team| (team.id.clone(), team.name.clone())) + .collect(); + assert_eq!( + listed, + vec![ + ("11".to_string(), "Alpha".to_string()), + ("22".to_string(), "22".to_string()), + ("33".to_string(), "33".to_string()), + ] + ); +} + +#[test] +fn cookie_access_is_off_until_a_cookie_or_browser_import_is_configured() { + let mut ctx = FetchContext::default(); + assert_eq!(CookieAccess::from_context(&ctx), CookieAccess::Off); + ctx.manual_cookie_header = Some(" ".to_string()); + assert_eq!(CookieAccess::from_context(&ctx), CookieAccess::Off); + ctx.manual_cookie_missing = true; + assert_eq!(CookieAccess::from_context(&ctx), CookieAccess::Off); + ctx.browser_cookie_import = true; + ctx.manual_cookie_header = None; + assert_eq!(CookieAccess::from_context(&ctx), CookieAccess::Browser); + ctx.manual_cookie_header = Some("llama_dev_sess=abc".to_string()); + assert_eq!( + CookieAccess::from_context(&ctx), + CookieAccess::Manual("llama_dev_sess=abc".to_string()) + ); +} + +#[tokio::test] +async fn disabled_cookies_never_contact_dev_meta_ai() { + // The default context is the Windows default: manual source, no cookie. + assert!( + fetch_reading(&FetchContext::default(), &identity()) + .await + .is_none() + ); +} + +#[test] +fn only_the_session_cookie_is_forwarded() { + assert_eq!( + session_cookie("theme=dark; llama_dev_sess=abc123; other=1").as_deref(), + Some("llama_dev_sess=abc123") + ); + assert_eq!(session_cookie("theme=dark"), None); + assert_eq!(session_cookie("llama_dev_sess="), None); + // Duplicates are ambiguous and rejected. + assert_eq!(session_cookie("llama_dev_sess=a; llama_dev_sess=b"), None); +} + +#[test] +fn manual_cookies_accept_headers_cookie_lines_and_curl_captures() { + let expected = Some("llama_dev_sess=abc123".to_string()); + assert_eq!(manual_session_cookie("llama_dev_sess=abc123"), expected); + assert_eq!( + manual_session_cookie("Cookie: a=1; llama_dev_sess=abc123"), + expected + ); + assert_eq!( + manual_session_cookie( + "curl 'https://dev.meta.ai/api/auth/me' -H 'accept: */*' -H 'Cookie: a=1; llama_dev_sess=abc123' -H 'user-agent: x'" + ), + expected + ); + assert_eq!( + manual_session_cookie("curl 'https://dev.meta.ai/api/auth/me' -H 'accept: */*'"), + None + ); + assert_eq!(manual_session_cookie("unrelated=1"), None); +} + +#[test] +fn weighted_amounts_accept_decimal_strings_and_reject_everything_else() { + use serde_json::json; + assert_eq!(amount(&json!("60000000000")), Some(60_000_000_000.0)); + assert_eq!(amount(&json!(5)), Some(5.0)); + for invalid in [ + json!("1e3"), + json!("-1"), + json!(""), + json!(-1), + json!(true), + json!(null), + json!("0x10"), + ] { + assert_eq!(amount(&invalid), None, "{invalid}"); + } + assert_eq!(percent(&json!("50"), &json!("0")), None); + assert_eq!(percent(&json!("500"), &json!("100")), Some(100.0)); +} + +#[test] +fn windowless_result_without_a_reading_matches_the_login_only_shape() { + let result = fallback_result(None); + assert_eq!(result.source_label, "oauth"); + assert!(!result.pace_authoritative); + assert!(result.usage.primary.is_informational); + assert_eq!( + detail(&result, "quota").value(), + "Not included in this login response" + ); + assert_eq!(detail(&result, "plan").value(), "Muse Code Power Usage"); + assert_eq!(result.display_details().len(), 2); +} + +#[test] +fn app_bound_encryption_is_reported_as_a_status_note() { + let reading = WebReading { + note: Some(TeamNote::CookiesProtected), + ..WebReading::default() + }; + let result = fallback_result(Some(reading)); + assert!( + detail(&result, "browser-teams-status") + .value() + .contains("app-bound encryption") + ); + assert!(has_detail(&result, "quota")); +} diff --git a/rust/src/settings/provider_workspace.rs b/rust/src/settings/provider_workspace.rs index 338747c8ef..d6c720e486 100644 --- a/rust/src/settings/provider_workspace.rs +++ b/rust/src/settings/provider_workspace.rs @@ -30,6 +30,9 @@ pub fn validate_provider_workspace_value( .chars() .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_') }), + ProviderId::Muse => validate_id(trimmed, "Muse browser team ID", |value| { + value.chars().all(|c| c.is_ascii_digit()) + }), ProviderId::Zed => validate_zed_url(trimmed), ProviderId::Xai => { if trimmed.contains('/') || trimmed == "." || trimmed == ".." { @@ -171,6 +174,24 @@ fn is_blocked_ipv6(ip: Ipv6Addr) -> bool { mod tests { use super::*; + #[test] + fn muse_browser_team_id_is_digits_only() { + assert_eq!( + validate_provider_workspace_value(ProviderId::Muse, " 424242424242 ").unwrap(), + "424242424242" + ); + assert_eq!( + validate_provider_workspace_value(ProviderId::Muse, " ").unwrap(), + "" + ); + for invalid in ["team-1", "12/34", "1 2", "-5", "1e3", "../1"] { + assert!( + validate_provider_workspace_value(ProviderId::Muse, invalid).is_err(), + "{invalid} must be rejected" + ); + } + } + #[test] fn validates_workspace_ids_by_provider() { assert_eq!( From 9458d879317e357b33405511442a63803ed36eed Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:24:23 +0700 Subject: [PATCH 123/407] Port upstream 0.67.0: show Grok product usage breakdown from the credits proxy --- rust/src/providers/grok/credits_proxy.rs | 19 +- rust/src/providers/grok/mod.rs | 12 + rust/src/providers/grok/product_usage.rs | 113 ++++++ .../src/providers/grok/product_usage_tests.rs | 338 ++++++++++++++++++ 4 files changed, 480 insertions(+), 2 deletions(-) create mode 100644 rust/src/providers/grok/product_usage.rs create mode 100644 rust/src/providers/grok/product_usage_tests.rs diff --git a/rust/src/providers/grok/credits_proxy.rs b/rust/src/providers/grok/credits_proxy.rs index dd89e57668..8172b8d501 100644 --- a/rust/src/providers/grok/credits_proxy.rs +++ b/rust/src/providers/grok/credits_proxy.rs @@ -15,6 +15,7 @@ use crate::core::ProviderError; use crate::providers::{BoundedBodyError, read_bounded_response}; use super::billing::GrokBillingSnapshot; +use super::product_usage::{GrokProductUsage, LossyProductUsage}; use super::{GrokCredentials, GrokProvider, grok_plan_display_name}; pub(super) const CREDITS_PROXY_ENDPOINT: &str = @@ -31,6 +32,9 @@ const MAX_RESPONSE_BYTES: usize = 256 * 1024; pub(super) struct BearerBilling { pub(super) billing: GrokBillingSnapshot, pub(super) subscription_tier: Option, + /// Product shares of the same payload's wire credit percent; empty for + /// every other percent source (on-demand ratio, period-only, gRPC-web). + pub(super) product_usage: Vec, } #[derive(Deserialize)] @@ -56,6 +60,8 @@ struct CreditsConfig { on_demand_used: Option, #[serde(rename = "subscriptionTier")] subscription_tier: Option, + #[serde(default, rename = "productUsage")] + product_usage: LossyProductUsage, } #[derive(Deserialize)] @@ -126,9 +132,14 @@ pub(super) fn parse_credits_response( }; let window_minutes = window_minutes(period_start, resets_at, now); - let used_percent = config + let wire_percent = config .credit_usage_percent - .filter(|percent| percent.is_finite()) + .filter(|percent| percent.is_finite()); + // Composition is checked against the raw percent, before display clamping. + let product_usage = wire_percent + .map(|percent| config.product_usage.composing(percent)) + .unwrap_or_default(); + let used_percent = wire_percent .or_else(|| { let cap = config.on_demand_cap?.val.filter(|cap| *cap > 0.0)?; Some(config.on_demand_used?.val? / cap * 100.0) @@ -149,6 +160,7 @@ pub(super) fn parse_credits_response( window_minutes, }, subscription_tier, + product_usage, }) } @@ -168,6 +180,8 @@ fn adopt_grpc_percent(proxy: BearerBilling, grpc: GrokBillingSnapshot) -> Bearer ..proxy.billing }, subscription_tier: proxy.subscription_tier, + // The grok.com percent is a different total than any proxy share list. + product_usage: Vec::new(), } } @@ -190,6 +204,7 @@ impl GrokProvider { return Ok(BearerBilling { billing, subscription_tier: None, + product_usage: Vec::new(), }); } }; diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 0db1fdcc7a..7c72b7421d 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -8,6 +8,7 @@ pub mod accounts; mod billing; mod credits_proxy; pub mod local_sessions; +mod product_usage; use async_trait::async_trait; use chrono::{DateTime, Utc}; @@ -137,6 +138,7 @@ impl GrokProvider { let BearerBilling { billing, subscription_tier, + product_usage, } = match self.fetch_bearer_billing(credentials).await { Ok(bearer) => bearer, Err(error) => { @@ -165,6 +167,12 @@ impl GrokProvider { credentials.team_id.clone(), plan, ); + // Reset-credit enrichment adds inventory only, so the breakdown survives it. + let result = product_usage::display_details(&product_usage) + .into_iter() + .fold(result, |result, detail| { + result.with_display_detail(Some(detail)) + }); Ok(match reset_credits { Some(credits) => result.with_inventory_item(credits), None => result, @@ -837,3 +845,7 @@ mod tests; #[cfg(test)] #[path = "credits_proxy_tests.rs"] mod credits_proxy_tests; + +#[cfg(test)] +#[path = "product_usage_tests.rs"] +mod product_usage_tests; diff --git a/rust/src/providers/grok/product_usage.rs b/rust/src/providers/grok/product_usage.rs new file mode 100644 index 0000000000..6483e398be --- /dev/null +++ b/rust/src/providers/grok/product_usage.rs @@ -0,0 +1,113 @@ +//! Per-product credit shares carried by the Grok credits proxy +//! (`config.productUsage`), shown as plain detail rows beside the total. +//! +//! Every share is a slice of the same credit pool as the primary window, so +//! it never becomes a rate window or progress bar. Shares are trusted only +//! when they compose the credit percentage published in the same payload; any +//! malformed entry drops the whole list so a partial list can never pass the +//! composition check as if it were complete. Follows upstream +//! `GrokCreditsProxyFetcher` (`composingProducts`, `LossyProductUsageArray`) +//! and `GrokProductUsageDetails` (v0.67.0). + +use serde::{Deserialize, Deserializer}; + +use crate::core::ProviderDisplayDetail; + +/// Rounding allowance between the summed shares and the raw credit percent. +const COMPOSITION_TOLERANCE_PERCENT: f64 = 1.0; + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct GrokProductUsage { + pub(super) product: String, + pub(super) used_percent: f64, +} + +#[derive(Deserialize)] +struct RawProductUsage { + product: String, + #[serde(rename = "usagePercent")] + usage_percent: f64, +} + +/// `config.productUsage` decoded all-or-nothing: a non-array value or a single +/// malformed entry leaves `None` without failing the surrounding payload. +#[derive(Debug, Default)] +pub(super) struct LossyProductUsage(Option>); + +impl<'de> Deserialize<'de> for LossyProductUsage { + fn deserialize>(deserializer: D) -> Result { + // Buffer the value first so a bad shape cannot desynchronize the parser. + let value = serde_json::Value::deserialize(deserializer)?; + let products = serde_json::from_value::>(value) + .ok() + .and_then(|raw| { + raw.into_iter() + .map(|entry| { + let product = entry.product.trim(); + (!product.is_empty() + && entry.usage_percent.is_finite() + && entry.usage_percent >= 0.0) + .then(|| GrokProductUsage { + product: product.to_string(), + used_percent: entry.usage_percent, + }) + }) + .collect::>>() + }); + Ok(Self(products)) + } +} + +impl LossyProductUsage { + /// Shares whose sum matches the raw (unclamped) credit percent from the + /// same payload; anything else is dropped. + pub(super) fn composing(self, credit_usage_percent: f64) -> Vec { + let products = self.0.unwrap_or_default(); + let sum: f64 = products.iter().map(|entry| entry.used_percent).sum(); + if products.is_empty() || (sum - credit_usage_percent).abs() > COMPOSITION_TOLERANCE_PERCENT + { + return Vec::new(); + } + products + } +} + +fn product_label(product: &str) -> &str { + match product { + "GrokBuild" => "Grok Build", + "GrokChat" => "Grok Chat", + "GrokImagine" => "Grok Imagine", + "GrokAppBuilder" => "Grok App Builder", + other => other, + } +} + +fn format_share(percent: f64) -> String { + let clamped = percent.clamp(0.0, 100.0); + if clamped > 0.0 && clamped < 1.0 { + "<1%".to_string() + } else { + format!("{clamped:.0}%") + } +} + +/// Plain detail rows sorted by share, largest first (ties keep wire order), +/// with zero-share products omitted. +pub(super) fn display_details(products: &[GrokProductUsage]) -> Vec { + let mut shares: Vec<&GrokProductUsage> = products + .iter() + .filter(|entry| entry.used_percent > 0.0) + .collect(); + // `sort_by` is stable, so equal shares keep their wire order. + shares.sort_by(|left, right| right.used_percent.total_cmp(&left.used_percent)); + shares + .into_iter() + .filter_map(|entry| { + ProviderDisplayDetail::new( + format!("grok.product.{}", entry.product), + product_label(&entry.product), + format_share(entry.used_percent), + ) + }) + .collect() +} diff --git a/rust/src/providers/grok/product_usage_tests.rs b/rust/src/providers/grok/product_usage_tests.rs new file mode 100644 index 0000000000..0cf53e2365 --- /dev/null +++ b/rust/src/providers/grok/product_usage_tests.rs @@ -0,0 +1,338 @@ +use super::credits_proxy::{BearerBilling, parse_credits_response}; +use super::product_usage::{GrokProductUsage, display_details}; +use super::tests::billing_response_with_percent; +use super::*; + +/// Upstream live capture (`LiveMultiProductCreditsPayload.p6`). +const LIVE_MULTI_PRODUCT_BODY: &str = r#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-09-20T18:42:45.537749+00:00","end":"2026-09-27T18:42:45.537749+00:00"},"creditUsagePercent":6.0,"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"productUsage":[{"product":"GrokChat","usagePercent":4.0},{"product":"GrokBuild","usagePercent":2.0}],"isUnifiedBillingUser":true,"prepaidBalance":{"val":0},"topUpMethod":"TOP_UP_METHOD_SAVED_PAYMENT_METHOD","billingPeriodStart":"2026-09-20T18:42:45.537749+00:00","billingPeriodEnd":"2026-09-27T18:42:45.537749+00:00"}}"#; +const BASE_PERIOD: &str = + r#""currentPeriod":{"start":"2026-09-20T00:00:00Z","end":"2026-09-27T00:00:00Z"}"#; + +fn product(name: &str, used_percent: f64) -> GrokProductUsage { + GrokProductUsage { + product: name.to_string(), + used_percent, + } +} + +fn parse(body: &str) -> BearerBilling { + let now = DateTime::parse_from_rfc3339("2026-09-23T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + parse_credits_response(body.as_bytes(), now).unwrap() +} + +fn rows(products: &[GrokProductUsage]) -> Vec<(String, String, String)> { + display_details(products) + .iter() + .map(|row| { + assert!(row.progress().is_none() && row.secondary_value().is_none()); + ( + row.id().to_string(), + row.title().to_string(), + row.value().to_string(), + ) + }) + .collect() +} + +#[test] +fn live_multi_product_capture_keeps_the_composed_breakdown() { + let parsed = parse(LIVE_MULTI_PRODUCT_BODY); + assert_eq!(parsed.billing.used_percent, Some(6.0)); + assert_eq!( + parsed.product_usage, + vec![product("GrokChat", 4.0), product("GrokBuild", 2.0)] + ); + assert_eq!( + rows(&parsed.product_usage), + vec![ + ( + "grok.product.GrokChat".to_string(), + "Grok Chat".to_string(), + "4%".to_string() + ), + ( + "grok.product.GrokBuild".to_string(), + "Grok Build".to_string(), + "2%".to_string() + ), + ] + ); +} + +#[test] +fn weekly_capture_with_one_product_keeps_it() { + let parsed = parse( + r#"{"config":{"creditUsagePercent":1.0,"productUsage":[{"product":"GrokBuild","usagePercent":1.0}]}}"#, + ); + assert_eq!(parsed.product_usage, vec![product("GrokBuild", 1.0)]); +} + +#[test] +fn malformed_products_drop_the_breakdown_but_not_total_or_period() { + let baseline = parse(&format!( + r#"{{"config":{{"creditUsagePercent":42,{BASE_PERIOD}}}}}"# + )); + assert!(baseline.product_usage.is_empty()); + let fragments = [ + r#""productUsage":null"#, + r#""productUsage":{}"#, + r#""productUsage":"wrong""#, + r#""productUsage":[]"#, + r#""productUsage":[{"product":"GrokBuild","usagePercent":"1"}]"#, + r#""productUsage":[{"usagePercent":1}]"#, + r#""productUsage":[{"product":42,"usagePercent":1}]"#, + r#""productUsage":[{"product":"GrokBuild"}]"#, + r#""productUsage":[{"product":"GrokBuild","usagePercent":-1}]"#, + r#""productUsage":[{"product":" ","usagePercent":1}]"#, + r#""productUsage":[{"product":"GrokChat","usagePercent":42},42]"#, + ]; + for fragment in fragments { + let parsed = parse(&format!( + r#"{{"config":{{"creditUsagePercent":42,{BASE_PERIOD},{fragment}}}}}"# + )); + assert!(parsed.product_usage.is_empty(), "{fragment}"); + assert_eq!(parsed.billing.used_percent, Some(42.0), "{fragment}"); + assert_eq!(parsed.billing.resets_at, baseline.billing.resets_at); + assert_eq!( + parsed.billing.window_minutes, + baseline.billing.window_minutes + ); + assert!(parsed.billing.used_percent_is_wire_published); + } +} + +#[test] +fn one_malformed_entry_cannot_hide_behind_the_tolerance() { + // 5 + a missing entry would otherwise compose 6 within 1 point. + let parsed = parse( + r#"{"config":{"creditUsagePercent":6,"productUsage":[{"product":"GrokChat","usagePercent":5},{"product":"GrokBuild","usagePercent":"1"}]}}"#, + ); + assert!(parsed.product_usage.is_empty()); + assert_eq!(parsed.billing.used_percent, Some(6.0)); + let exact_remainder = parse( + r#"{"config":{"creditUsagePercent":5,"productUsage":[{"product":"GrokChat","usagePercent":5},{"usagePercent":1}]}}"#, + ); + assert!(exact_remainder.product_usage.is_empty()); + assert_eq!(exact_remainder.billing.used_percent, Some(5.0)); +} + +#[test] +fn products_attach_only_to_the_wire_credit_percent() { + let cases = [ + ( + r#""onDemandCap":{"val":100},"onDemandUsed":{"val":3}"#, + Some(3.0), + ), + (r#""billingPeriodEnd":"2026-09-27T00:00:00Z""#, None), + ]; + for (fields, percent) in cases { + let parsed = parse(&format!( + r#"{{"config":{{{fields},"productUsage":[{{"product":"GrokBuild","usagePercent":3}}]}}}}"# + )); + assert_eq!(parsed.billing.used_percent, percent, "{fields}"); + assert!(parsed.product_usage.is_empty(), "{fields}"); + } +} + +#[test] +fn shares_that_do_not_compose_the_percent_are_dropped() { + for products in [ + r#"[{"product":"GrokBuild","usagePercent":60}]"#, + r#"[{"product":"GrokBuild","usagePercent":20},{"product":"GrokChat","usagePercent":5}]"#, + ] { + let parsed = parse(&format!( + r#"{{"config":{{"creditUsagePercent":30,{BASE_PERIOD},"productUsage":{products}}}}}"# + )); + assert!(parsed.product_usage.is_empty(), "{products}"); + assert_eq!(parsed.billing.used_percent, Some(30.0)); + assert_eq!(parsed.billing.window_minutes, Some(10080)); + } +} + +#[test] +fn shares_within_rounding_of_the_percent_are_kept() { + let rounded = parse( + r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":6.0},{"product":"GrokChat","usagePercent":3.6}]}}"#, + ); + assert_eq!( + rounded.product_usage, + vec![product("GrokBuild", 6.0), product("GrokChat", 3.6)] + ); + let exactly_one_point_off = parse( + r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":9.0}]}}"#, + ); + assert_eq!( + exactly_one_point_off.product_usage, + vec![product("GrokBuild", 9.0)] + ); + let just_outside = parse( + r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":8.9}]}}"#, + ); + assert!(just_outside.product_usage.is_empty()); +} + +#[test] +fn shares_compose_the_raw_unclamped_percent() { + let parsed = parse( + r#"{"config":{"creditUsagePercent":120,"productUsage":[{"product":"GrokBuild","usagePercent":120}]}}"#, + ); + assert_eq!(parsed.billing.used_percent, Some(100.0)); + assert_eq!(parsed.product_usage, vec![product("GrokBuild", 120.0)]); +} + +#[test] +fn rows_sort_by_share_keep_ties_in_wire_order_and_omit_zero() { + let products = [ + product("GrokBuild", 1.0), + product("GrokChat", 5.0), + product("GrokImagine", 0.0), + product("FutureGrok", 0.2), + product("GrokAppBuilder", 5.0), + ]; + let rows = rows(&products); + assert_eq!( + rows.iter().map(|row| row.0.as_str()).collect::>(), + [ + "grok.product.GrokChat", + "grok.product.GrokAppBuilder", + "grok.product.GrokBuild", + "grok.product.FutureGrok", + ] + ); + assert_eq!( + rows.iter().map(|row| row.1.as_str()).collect::>(), + ["Grok Chat", "Grok App Builder", "Grok Build", "FutureGrok"] + ); + assert_eq!( + rows.iter().map(|row| row.2.as_str()).collect::>(), + ["5%", "5%", "1%", "<1%"] + ); + assert!(display_details(&[]).is_empty()); + assert!(display_details(&[product("GrokChat", 0.0)]).is_empty()); +} + +#[test] +fn product_names_are_trimmed_and_imagine_is_labelled() { + let parsed = parse( + r#"{"config":{"creditUsagePercent":3,"productUsage":[{"product":" GrokImagine ","usagePercent":2.5},{"product":"Other","usagePercent":0.5}]}}"#, + ); + assert_eq!(parsed.product_usage[0].product, "GrokImagine"); + assert_eq!( + rows(&parsed.product_usage)[0], + ( + "grok.product.GrokImagine".to_string(), + "Grok Imagine".to_string(), + "2%".to_string() + ) + ); +} + +fn provider_for(server: &mockito::ServerGuard) -> GrokProvider { + GrokProvider::new() + .with_billing_endpoint_for_tests(format!("{}/billing", server.url())) + .with_credits_proxy_endpoint_for_tests(format!("{}/credits", server.url())) +} + +fn context_without_credits() -> FetchContext { + FetchContext { + include_credits: false, + ..FetchContext::default() + } +} + +#[tokio::test] +async fn bearer_result_shows_the_breakdown_as_plain_rows_under_one_bar() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body(LIVE_MULTI_PRODUCT_BODY) + .create_async() + .await; + + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context_without_credits(), + ) + .await + .unwrap(); + + assert_eq!(result.usage.primary.used_percent, 6.0); + assert!(result.usage.secondary.is_none()); + assert!(result.usage.tertiary.is_none()); + let shown: Vec<(&str, &str)> = result + .display_details() + .iter() + .map(|row| (row.title(), row.value())) + .collect(); + assert_eq!(shown, [("Grok Chat", "4%"), ("Grok Build", "2%")]); + assert!( + result + .display_details() + .iter() + .all(|row| row.progress().is_none()) + ); +} + +#[tokio::test] +async fn reset_credit_enrichment_preserves_the_breakdown() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body(LIVE_MULTI_PRODUCT_BODY) + .create_async() + .await; + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context_without_credits(), + ) + .await + .unwrap() + .with_inventory_item(ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count: 1, + next_expires_at: None, + }); + + assert_eq!(result.display_details().len(), 2); + assert_eq!(result.inventory.len(), 1); +} + +#[tokio::test] +async fn grpc_percent_never_borrows_the_proxy_products() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body( + r#"{"config":{"currentPeriod":{"start":"2026-08-06T00:00:00Z","end":"2026-08-13T00:00:00Z"},"productUsage":[{"product":"GrokBuild","usagePercent":12}]}}"#, + ) + .create_async() + .await; + server + .mock("POST", "/billing") + .with_status(200) + .with_body(billing_response_with_percent(12.0)) + .create_async() + .await; + + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context_without_credits(), + ) + .await + .unwrap(); + + assert_eq!(result.usage.primary.used_percent, 12.0); + assert!(result.display_details().is_empty()); +} From 1a798f645f02605d2eeb32f0c9b59f3ae186b848 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:25:38 +0700 Subject: [PATCH 124/407] Port upstream 0.68.0: show Mistral plan amounts as detail lines --- .../src-tauri/src/auto_resume_tests.rs | 1 + .../src-tauri/src/commands/bridge.rs | 6 ++ .../src-tauri/src/commands/bridge/status.rs | 45 +++++++++ .../src-tauri/src/commands/providers.rs | 1 + apps/desktop-tauri/src-tauri/src/powertoys.rs | 1 + .../src-tauri/src/tray_bridge.rs | 6 ++ .../src-tauri/src/usage_metric.rs | 1 + .../src/components/MenuCard.test.tsx | 42 ++++++++- .../src/components/MenuCardDetails.tsx | 7 +- .../src/floatbar/FloatBar.test.tsx | 22 +++++ apps/desktop-tauri/src/floatbar/FloatBar.tsx | 3 +- .../src/lib/usageWindows.test.ts | 35 ++++++- apps/desktop-tauri/src/lib/usageWindows.ts | 19 ++++ apps/desktop-tauri/src/styles.css | 18 +++- .../src/surfaces/menu-surface-layout.css | 3 +- .../providers/sections/UsageSection.test.tsx | 45 +++++++++ .../providers/sections/UsageSection.tsx | 5 +- apps/desktop-tauri/src/types/bridge.ts | 2 + rust/src/cli/usage/render.rs | 27 +++++- rust/src/cli/usage_tests.rs | 92 ++++++++++++++++++- rust/src/core/rate_window.rs | 15 +++ rust/src/providers/copilot/api.rs | 1 + rust/src/providers/mistral/mod.rs | 18 +++- 23 files changed, 401 insertions(+), 14 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/auto_resume_tests.rs b/apps/desktop-tauri/src-tauri/src/auto_resume_tests.rs index 07b289498f..bc9ba159ef 100644 --- a/apps/desktop-tauri/src-tauri/src/auto_resume_tests.rs +++ b/apps/desktop-tauri/src-tauri/src/auto_resume_tests.rs @@ -42,6 +42,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ab5efbef8a..9f94466a70 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -23,6 +23,9 @@ pub struct RateWindowSnapshot { pub is_exhausted: bool, #[serde(default)] pub is_informational: bool, + /// `reset_description` is a detail line (for example spend amounts), not reset wording. + #[serde(default)] + pub description_is_detail: bool, #[serde(default)] pub reserve_percent: Option, #[serde(default)] @@ -53,6 +56,7 @@ impl RateWindowSnapshot { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: rw.description_is_detail, } } @@ -518,6 +522,7 @@ impl ProviderUsageSnapshot { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }, primary_label: Some(metadata.session_label.to_string()), secondary: None, @@ -980,6 +985,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge/status.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge/status.rs index 6498eec07d..04bbb1ac48 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge/status.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge/status.rs @@ -32,6 +32,10 @@ fn compact_reset_description( return Some(format_compact_reset_countdown(dt, lang)); } + if window.description_is_detail { + return None; + } + window .reset_description .as_deref() @@ -128,3 +132,44 @@ pub(crate) fn friendly_provider_error(id: ProviderId, error: &str) -> String { trimmed.to_string() } + +#[cfg(test)] +mod tests { + use super::*; + use codexbar::core::RateWindow; + use codexbar::settings::Language; + + #[test] + fn detail_backed_description_is_never_a_tray_reset_label() { + let window = RateWindowSnapshot::from_rate_window( + &RateWindow::with_details( + 13.0, + None, + None, + Some("34.07 EUR / 255.00 EUR · 220.93 EUR remaining".to_string()), + ) + .with_description_as_detail(), + ); + + assert!(window.description_is_detail); + assert_eq!(compact_tray_status_label(&window, Language::English), "13%"); + } + + #[test] + fn detail_backed_window_keeps_its_countdown_when_reset_is_known() { + let window = RateWindowSnapshot::from_rate_window( + &RateWindow::with_details( + 13.0, + None, + Some(chrono::Utc::now() + chrono::Duration::minutes(125)), + Some("34.07 EUR / 255.00 EUR · 220.93 EUR remaining".to_string()), + ) + .with_description_as_detail(), + ); + + let label = compact_tray_status_label(&window, Language::English); + + assert!(label.starts_with("13% • "), "{label}"); + assert!(!label.contains("EUR"), "{label}"); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..57652d7f04 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -1320,6 +1320,7 @@ mod reset_backfill_tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, } } diff --git a/apps/desktop-tauri/src-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 99012ae7c3..ff257376c2 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -178,6 +178,7 @@ mod tests { reserve_percent: None, reserve_description: None, reserve_eta_seconds: None, + description_is_detail: false, reserve_will_last_to_reset: false, } } diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index aebb7c2115..882573b666 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -1034,6 +1034,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }, primary_label: None, secondary: secondary_percent.map(|pct| crate::commands::RateWindowSnapshot { @@ -1048,6 +1049,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }), secondary_label: None, model_specific: None, @@ -1063,6 +1065,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }), tertiary_label: None, extra_rate_windows: Vec::new(), @@ -1127,6 +1130,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }, } } @@ -1537,6 +1541,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }); let (primary, _) = selected_tray_percents(&snapshot, &settings); @@ -1581,6 +1586,7 @@ mod tests { reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, }); snapshot.extra_rate_windows.push(fake_extra_window(90.0)); diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index f6e8f5f532..095b181006 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -267,6 +267,7 @@ fn derived_window(used_percent: f64, resets_at: Option) -> RateWindowSna reserve_description: None, reserve_will_last_to_reset: false, reserve_eta_seconds: None, + description_is_detail: false, } } diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index c8ec795546..3b840688be 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -55,7 +55,7 @@ function rateWindow( function provider( error: string | null, usedPercent = 0, - opts: { exhausted?: boolean; resetDescription?: string | null } = {}, + opts: { exhausted?: boolean; resetDescription?: string | null; resetsAt?: string | null } = {}, ): ProviderUsageSnapshot { return { providerId: "claude", @@ -356,6 +356,46 @@ describe("MenuCard", () => { expect(await screen.findByText("Fable only")).toBeInTheDocument(); }); + it.each([false, true])( + "shows detail-backed amounts as their own line, never as reset text (resetsAt: %s)", + async (hasReset) => { + const resetsAt = hasReset + ? new Date(Date.now() + 3 * 60 * 60 * 1000 + 30_000).toISOString() + : null; + const snapshot = provider(null, 75, { + resetDescription: "19.17 EUR / 25.50 EUR · 6.33 EUR remaining", + resetsAt, + }); + snapshot.providerId = "mistral"; + snapshot.primary.descriptionIsDetail = true; + snapshot.selectedMetric.descriptionIsDetail = true; + snapshot.extraRateWindows = [ + { + id: "mistral-monthly-plan", + title: "Monthly Plan", + window: { + ...rateWindow(13, { + resetDescription: "34.07 EUR / 255.00 EUR · 220.93 EUR remaining", + resetsAt, + }), + descriptionIsDetail: true, + }, + }, + ]; + + renderCard(snapshot); + + const primaryDetail = await screen.findByText( + "19.17 EUR / 25.50 EUR · 6.33 EUR remaining", + ); + const planDetail = screen.getByText("34.07 EUR / 255.00 EUR · 220.93 EUR remaining"); + expect(primaryDetail).toHaveClass("menu-metric__detail"); + expect(planDetail).toHaveClass("menu-metric__detail"); + expect(screen.queryByText(/Resets .*EUR/)).not.toBeInTheDocument(); + expect(document.querySelectorAll(".menu-metric__reset")).toHaveLength(hasReset ? 2 : 0); + }, + ); + it("renders informational metrics without quota percentages", async () => { const snapshot = provider(null, 20); snapshot.extraRateWindows = [ diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 31634bd910..3826fc9823 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -24,6 +24,7 @@ import { SimpleBarChart, StackedBarChart } from "./MiniBarChart"; import { InventoryItemRow } from "./InventoryRows"; import { QuotaWindowHistory } from "./QuotaWindowHistory"; import { getPaceBudget, type PaceBudget } from "../lib/paceBudget"; +import { resetDescriptionFallback, windowDetailText } from "../lib/usageWindows"; import PaceDetailsChart from "./PaceDetailsChart"; /** Format a reserve description from raw pace data at render time. */ @@ -337,9 +338,10 @@ function MetricRow({ const barDisplayPct = showAsUsed ? barPct : Math.max(0, Math.min(100, remain)); const displayLabel = showAsUsed ? t("PanelUsedSuffix") : t("PanelLeftSuffix"); const level = levelOf(remain, snap.isExhausted); + const detailText = windowDetailText(snap); const resetText = useFormattedResetTime( snap.resetsAt, - isInformational ? null : snap.resetDescription, + isInformational ? null : resetDescriptionFallback(snap), resetTimeRelative, resetFormatMode ?? "reset", ); @@ -380,6 +382,9 @@ function MetricRow({ {resetText} )} + {!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 b61d8c69f37e966045e3f6333aac6d3fb21d9d69 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:31:08 +0700 Subject: [PATCH 125/407] Port upstream 0.66.0: DevPass plan credits and premium weekly provider --- README.md | 1 + .../providers/icons/ProviderIcon-devpass.svg | 3 + .../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 + docs/PROVIDERS.md | 13 + rust/src/cli/serve/dashboard/icons.rs | 4 + .../dashboard/icons/ProviderIcon-devpass.svg | 3 + rust/src/core/provider.rs | 10 +- rust/src/core/provider_factory.rs | 25 +- rust/src/core/token_accounts.rs | 1 + rust/src/providers/devpass/mod.rs | 248 ++++++++++ rust/src/providers/devpass/tests.rs | 467 ++++++++++++++++++ rust/src/providers/devpass/wire.rs | 139 ++++++ rust/src/providers/mod.rs | 2 + rust/src/settings/api_keys.rs | 11 + 17 files changed, 920 insertions(+), 14 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-devpass.svg create mode 100644 rust/src/cli/serve/dashboard/icons/ProviderIcon-devpass.svg create mode 100644 rust/src/providers/devpass/mod.rs create mode 100644 rust/src/providers/devpass/tests.rs create mode 100644 rust/src/providers/devpass/wire.rs diff --git a/README.md b/README.md index 2e052ad2f9..9148b86707 100755 --- a/README.md +++ b/README.md @@ -120,6 +120,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Deepgram | API Key | Project Usage | | Groq | API Key | Enterprise Metrics | | LLM Proxy | API Key | Quota Stats | +| DevPass | API Key | Plan credits, Premium weekly, API-key spend | diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-devpass.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-devpass.svg new file mode 100644 index 0000000000..00bcb52fe8 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-devpass.svg @@ -0,0 +1,3 @@ + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 2fdc4380d7..b2570f8d7e 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -19,6 +19,7 @@ import crossmodel from "./icons/ProviderIcon-crossmodel.svg?raw"; import cursor from "./icons/ProviderIcon-cursor.svg?raw"; import deepgram from "./icons/ProviderIcon-deepgram.svg?raw"; import deepinfra from "./icons/ProviderIcon-deepinfra.svg?raw"; +import devpass from "./icons/ProviderIcon-devpass.svg?raw"; import fireworks from "./icons/ProviderIcon-fireworks.svg?raw"; import aiand from "./icons/ProviderIcon-aiand.svg?raw"; import clinepass from "./icons/ProviderIcon-clinepass.svg?raw"; @@ -107,6 +108,7 @@ const RAW: Record = { cursor: tint(cursor), deepgram: tint(deepgram), deepinfra: tint(deepinfra), + devpass: tint(devpass), fireworks: tint(fireworks), aiand: tint(aiand), clinepass: tint(clinepass), @@ -174,6 +176,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { cursor: { id: "cursor", brandColor: "#00bfa5", fallbackLetter: "▸", svgPath: RAW.cursor }, deepgram: { id: "deepgram", brandColor: "#13ef93", fallbackLetter: "D", svgPath: RAW.deepgram }, deepinfra: { id: "deepinfra", brandColor: "#2a3275", fallbackLetter: "D", svgPath: RAW.deepinfra }, + devpass: { id: "devpass", brandColor: "#2563eb", fallbackLetter: "D", svgPath: RAW.devpass }, fireworks: { id: "fireworks", brandColor: "#f25b1c", fallbackLetter: "F", svgPath: RAW.fireworks }, aiand: { id: "aiand", brandColor: "#e25c2b", fallbackLetter: "&", svgPath: RAW.aiand }, clinepass: { id: "clinepass", brandColor: "#61a3fa", fallbackLetter: "C", svgPath: RAW.clinepass }, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..8cfe0e8636 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -36,7 +36,7 @@ const HAS_DASHBOARD = new Set([ "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", "opencode", "opencodego", "openrouter", "perplexity", "qoder", "codebuddy", "sakana", "stepfun", "t3chat", "venice", "vertexai", "warp", "windsurf", - "xai", "zai", "fireworks", "meta", "muse", "nous", + "xai", "zai", "fireworks", "meta", "muse", "nous", "devpass", ]); /** Provider IDs that have a status page URL in the backend */ const HAS_STATUS_PAGE = new Set([ diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx index 48c200779c..f7a960b71b 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx @@ -241,6 +241,7 @@ function providerSourceHintShort( case "zenmux": case "clinepass": case "neuralwatt": + case "devpass": case "doubao": case "crof": case "stepfun": diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index dda4c24441..5f517271d5 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -77,4 +77,5 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["meta", "Meta"], ["muse", "Muse Code"], ["nous", "Nous Portal"], + ["devpass", "DevPass"], ]; diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 237f7dab3d..f7a942f7c0 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -55,6 +55,19 @@ spend and, when the optional balance request succeeds, prepaid credit balance. It keeps those values in the cost/detail surfaces and does not invent a quota percentage or use a Replicate API token as a website credential. +### DevPass + +DevPass reads `GET https://api.llmgateway.io/v1/key` with a regular LLM Gateway +API key (`DEVPASS_API_KEY` or the Settings key field). Publishable keys and +end-user sessions get HTTP 403. Plan credits are the primary lane and the +premium weekly allowance is the secondary lane (seven-day window; the reset +comes only from the response, and an inactive window has no reset). A +pay-as-you-go key (`devPlan: none`) shows only key-scoped all-time spend. +Remaining plan credits are an allowance, not a wallet balance, so no balance is +reported. Amounts must be plain decimal strings; anything else fails the +refresh instead of showing as zero. The key is only sent to the fixed HTTPS +origin, redirects are not followed, and response bodies never appear in errors. + ## Listing what is enabled ```powershell diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 6ac230cc63..9c09ea2097 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -145,6 +145,10 @@ static ICONS: &[(&str, &[u8])] = &[ "ProviderIcon-devin", include_bytes!("icons/ProviderIcon-devin.svg"), ), + ( + "ProviderIcon-devpass", + include_bytes!("icons/ProviderIcon-devpass.svg"), + ), ( "ProviderIcon-doubao", include_bytes!("icons/ProviderIcon-doubao.svg"), diff --git a/rust/src/cli/serve/dashboard/icons/ProviderIcon-devpass.svg b/rust/src/cli/serve/dashboard/icons/ProviderIcon-devpass.svg new file mode 100644 index 0000000000..00bcb52fe8 --- /dev/null +++ b/rust/src/cli/serve/dashboard/icons/ProviderIcon-devpass.svg @@ -0,0 +1,3 @@ + + + diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..da22ad4740 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -93,6 +93,7 @@ pub enum ProviderId { Muse, Replicate, Nous, + DevPass, } impl ProviderId { @@ -179,6 +180,7 @@ impl ProviderId { ProviderId::Muse, ProviderId::Replicate, ProviderId::Nous, + ProviderId::DevPass, ] } @@ -226,6 +228,7 @@ impl ProviderId { ProviderId::Meta => "meta", ProviderId::Muse => "muse", ProviderId::Nous => "nous", + ProviderId::DevPass => "devpass", ProviderId::AiAnd => "aiand", ProviderId::Windsurf => "windsurf", ProviderId::Manus => "manus", @@ -313,6 +316,7 @@ impl ProviderId { ProviderId::Meta => "Meta", ProviderId::Muse => "Muse Code", ProviderId::Nous => "Nous Portal", + ProviderId::DevPass => "DevPass", ProviderId::AiAnd => "ai&", ProviderId::Windsurf => "Windsurf", ProviderId::Manus => "Manus", @@ -415,6 +419,7 @@ impl ProviderId { ProviderId::Meta => None, ProviderId::Muse => None, ProviderId::Nous => None, + ProviderId::DevPass => None, ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, @@ -495,6 +500,7 @@ impl ProviderId { "fireworks" | "fireworks-ai" | "fw" => Some(ProviderId::Fireworks), "muse" | "muse-code" | "muse code" => Some(ProviderId::Muse), "nous" | "nous-portal" | "nous portal" | "hermes" => Some(ProviderId::Nous), + "devpass" => Some(ProviderId::DevPass), "meta" | "metaspark" | "meta-spark" | "muse-spark" | "musespark" | "muse spark" | "meta muse spark" => Some(ProviderId::Meta), "aiand" | "ai&" | "ai-and" | "ai and" => Some(ProviderId::AiAnd), @@ -1106,6 +1112,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Muse => "#0668E1", ProviderId::Replicate => "#000000", ProviderId::Nous => "#D6A55C", + ProviderId::DevPass => "#2563EB", } } @@ -1120,7 +1127,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 80); + assert_eq!(all.len(), 81); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1181,6 +1188,7 @@ mod tests { assert!(all.contains(&ProviderId::Replicate)); assert!(all.contains(&ProviderId::Muse)); assert!(all.contains(&ProviderId::Nous)); + assert!(all.contains(&ProviderId::DevPass)); } #[test] diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 712c1d4fb9..c7e1ada519 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -11,18 +11,18 @@ use crate::providers::{ 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, HelmcodeProvider, 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, TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, - WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, - ZoomMateProvider, + CursorProvider, DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevPassProvider, + DevinProvider, DoubaoProvider, ElevenLabsProvider, FactoryProvider, FireworksProvider, + GeminiProvider, GrokProvider, GroqProvider, HelmcodeProvider, 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, TypeSafeProvider, V0Provider, VeniceProvider, + VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, + ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -111,6 +111,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Meta => Box::new(MetaProvider::new()), ProviderId::Muse => Box::new(MuseProvider::new()), ProviderId::Nous => Box::new(NousProvider::new()), + ProviderId::DevPass => Box::new(DevPassProvider::new()), } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 6dedc0fbde..f9f610192e 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -377,6 +377,7 @@ impl TokenAccountSupport { | ProviderId::Fireworks | ProviderId::Meta | ProviderId::Nous + | ProviderId::DevPass | ProviderId::Muse => None, } } diff --git a/rust/src/providers/devpass/mod.rs b/rust/src/providers/devpass/mod.rs new file mode 100644 index 0000000000..b8b0874c03 --- /dev/null +++ b/rust/src/providers/devpass/mod.rs @@ -0,0 +1,248 @@ +//! DevPass provider: plan credits and premium weekly usage through the LLM +//! Gateway API. +//! +//! Ported from upstream CodexBar 0.66.0 (`devpass` plugin). +//! `GET https://api.llmgateway.io/v1/key` with a regular gateway API key +//! returns the key's all-time spend plus the organization's DevPass plan +//! state. Only that fixed HTTPS origin receives the key, redirects are not +//! followed, and response bodies are never echoed into errors. +//! +//! Remaining plan credits are an allowance, not a wallet balance, so no typed +//! `CostSnapshot` balance is emitted. The premium window starts with the +//! first premium request; an inactive window has zero usage and no reset, and +//! no monthly reset is ever inferred. + +#[cfg(test)] +mod tests; +mod wire; + +use std::time::Duration; + +use async_trait::async_trait; +use reqwest::{Client, StatusCode}; + +use crate::core::{ + FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; +use crate::providers::{BoundedBodyError, read_bounded_response}; +use wire::{DevPlan, KeyData, Money, PlanFields}; + +const API_URL: &str = "https://api.llmgateway.io/v1/key"; +const CREDENTIAL_TARGET: &str = "codexbar-devpass"; +const ENV_KEYS: &[&str] = &["DEVPASS_API_KEY"]; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_BODY_BYTES: usize = 1024 * 1024; +/// The premium allowance window is seven days. +const PREMIUM_WINDOW_MINUTES: u32 = 7 * 24 * 60; + +pub struct DevPassProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl DevPassProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::DevPass, + display_name: "DevPass", + session_label: "Plan credits", + weekly_label: "Premium weekly", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://devpass.llmgateway.io/dashboard"), + status_page_url: None, + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()), + } + } +} + +impl Default for DevPassProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for DevPassProvider { + fn id(&self) -> ProviderId { + ProviderId::DevPass + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => { + let key = crate::providers::resolve_api_key( + ctx.api_key.as_deref(), + CREDENTIAL_TARGET, + ENV_KEYS, + )?; + fetch_key(&self.client, API_URL, &key).await + } + SourceMode::Web | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::OAuth] + } +} + +async fn fetch_key( + client: &Client, + url: &str, + key: &str, +) -> Result { + let response = client + .get(url) + .bearer_auth(key) + .header("Accept", "application/json") + .send() + .await?; + check_status(response.status())?; + let body = read_bounded_response(response, MAX_BODY_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => unrecognized(), + BoundedBodyError::Read(error) => ProviderError::Network(error), + })?; + parse_result(&body) +} + +fn parse_result(body: &[u8]) -> Result { + let data = wire::parse_key(body).ok_or_else(unrecognized)?; + let plan = match data.dev_plan { + DevPlan::None => None, + _ => Some(wire::parse_plan(body).ok_or_else(unrecognized)?), + }; + Ok(build_result(&data, plan.as_ref())) +} + +fn unrecognized() -> ProviderError { + ProviderError::Parse("DevPass returned an unrecognized usage response.".into()) +} + +fn check_status(status: StatusCode) -> Result<(), ProviderError> { + match status.as_u16() { + 200 => Ok(()), + 401 => Err(ProviderError::OAuthExpired( + "DevPass API key was rejected or is inactive.".into(), + )), + 403 => Err(ProviderError::Other( + "DevPass requires a regular gateway API key.".into(), + )), + 429 => Err(ProviderError::Other( + "DevPass usage requests are rate limited.".into(), + )), + 500.. => Err(ProviderError::Other( + "DevPass usage is temporarily unavailable.".into(), + )), + code => Err(ProviderError::Other(format!( + "DevPass returned HTTP {code}." + ))), + } +} + +fn usd(money: Money) -> String { + format!("${:.2}", money.value()) +} + +/// Percentage used, clamped to 0-100 so an over-limit amount fills the bar +/// while the row text keeps the actual figures. +fn percent(used: Money, limit: Money) -> f64 { + (used.value() / limit.value() * 100.0).clamp(0.0, 100.0) +} + +/// A used-of-limit row; the bar appears only for a positive limit. +fn allowance_row( + id: &str, + title: &str, + used: Money, + limit: Money, +) -> Option { + let row = ProviderDisplayDetail::new(id, title, format!("{} / {}", usd(used), usd(limit)))?; + if limit.value() > 0.0 { + row.with_progress(used.value().min(limit.value()), limit.value()) + } else { + Some(row) + } +} + +fn build_result(data: &KeyData, plan: Option<&PlanFields>) -> ProviderFetchResult { + let usage = match plan { + Some(plan) => plan_usage(plan), + None => UsageSnapshot::new(RateWindow::informational("Pay as you go, no DevPass plan")), + } + .with_login_method(data.dev_plan.login_method()); + + let mut result = ProviderFetchResult::new(usage, "api"); + if let Some(plan) = plan { + result = result + .with_display_detail(allowance_row( + "cycle-used", + "Cycle used", + plan.credits_used, + plan.credits_limit, + )) + .with_display_detail(ProviderDisplayDetail::new( + "cycle-remaining", + "Cycle remaining", + usd(plan.credits_remaining), + )) + .with_display_detail(allowance_row( + "premium-weekly", + "Premium weekly", + plan.premium_used, + plan.premium_limit, + )); + } + result = result.with_display_detail(ProviderDisplayDetail::new( + "key-usage", + "All-time key usage", + usd(data.usage), + )); + if let Some(limit) = data.limit { + result = result.with_display_detail(ProviderDisplayDetail::new( + "key-limit", + "Key spending limit", + usd(limit), + )); + } + result +} + +/// Plan credits are the primary lane and carry no reset or window length: +/// the endpoint exposes no monthly reset, so none is inferred. A zero +/// allowance has nothing to measure against and stays informational. +fn plan_usage(plan: &PlanFields) -> UsageSnapshot { + let primary = if plan.credits_limit.value() > 0.0 { + RateWindow::new(percent(plan.credits_used, plan.credits_limit)) + } else { + RateWindow::informational("No plan credit allowance reported") + }; + let mut usage = UsageSnapshot::new(primary); + if plan.premium_limit.value() > 0.0 { + usage = usage.with_secondary(RateWindow::with_details( + percent(plan.premium_used, plan.premium_limit), + Some(PREMIUM_WINDOW_MINUTES), + plan.premium_resets_at, + None, + )); + } + usage +} diff --git a/rust/src/providers/devpass/tests.rs b/rust/src/providers/devpass/tests.rs new file mode 100644 index 0000000000..3f7147eea0 --- /dev/null +++ b/rust/src/providers/devpass/tests.rs @@ -0,0 +1,467 @@ +use super::*; +use chrono::{TimeZone, Utc}; +use serde_json::{Value, json}; + +/// Public `/v1/key` example with synthetic values and a fixed reset +/// (upstream `DevPassPluginTests.fixture`). +fn fixture() -> Value { + json!({ + "label": "Fixture key", + "usage": "31.42", + "limit": null, + "devPlan": "pro", + "devPlanCreditsUsed": "25", + "devPlanCreditsLimit": "237", + "devPlanCreditsRemaining": "212.00", + "devPlanPremiumWeeklyLimit": "35.55", + "devPlanPremiumCreditsUsed": "5.00", + "devPlanPremiumWeekResetsAt": "2026-10-01T12:00:00.000Z", + }) +} + +fn body(data: &Value) -> Vec { + serde_json::to_vec(&json!({ "data": data })).unwrap() +} + +fn parse(data: &Value) -> Result { + parse_result(&body(data)) +} + +fn with(key: &str, value: Value) -> Value { + let mut data = fixture(); + data[key] = value; + data +} + +fn without(key: &str) -> Value { + let mut data = fixture(); + data.as_object_mut().unwrap().remove(key); + data +} + +fn row<'a>(result: &'a ProviderFetchResult, id: &str) -> &'a ProviderDisplayDetail { + result + .display_details() + .iter() + .find(|row| row.id() == id) + .unwrap_or_else(|| panic!("missing detail row {id}")) +} + +fn assert_close(actual: f64, expected: f64) { + assert!( + (actual - expected).abs() < 1e-4, + "{actual} is not close to {expected}" + ); +} + +#[test] +fn documented_cycle_and_weekly_credits_keep_their_scope() { + let result = parse(&fixture()).unwrap(); + let usage = &result.usage; + + assert_close(usage.primary.used_percent, 25.0 / 237.0 * 100.0); + assert!(!usage.primary.is_informational); + assert_eq!(usage.primary.resets_at, None); + assert_eq!(usage.primary.window_minutes, None); + + let weekly = usage.secondary.as_ref().expect("premium weekly window"); + assert_close(weekly.used_percent, 5.0 / 35.55 * 100.0); + assert_eq!(weekly.window_minutes, Some(10080)); + assert_eq!( + weekly.resets_at, + Some(Utc.with_ymd_and_hms(2026, 10, 1, 12, 0, 0).unwrap()) + ); + + assert!( + result.cost.is_none(), + "an allowance is not a wallet balance" + ); + assert_eq!(usage.login_method.as_deref(), Some("DevPass Pro")); + assert_eq!(row(&result, "cycle-used").value(), "$25.00 / $237.00"); + assert_eq!(row(&result, "cycle-remaining").value(), "$212.00"); + assert_eq!(row(&result, "premium-weekly").value(), "$5.00 / $35.55"); + assert_eq!(row(&result, "key-usage").value(), "$31.42"); + assert!( + result + .display_details() + .iter() + .all(|row| row.id() != "key-limit"), + "a null key limit has no row" + ); + let progress = row(&result, "cycle-used").progress().unwrap(); + assert_close(progress.used(), 25.0); + assert_close(progress.total(), 237.0); +} + +#[test] +fn key_spending_limit_row_appears_when_the_key_has_one() { + let result = parse(&with("limit", json!("50.5"))).unwrap(); + assert_eq!(row(&result, "key-limit").value(), "$50.50"); +} + +#[test] +fn every_plan_names_its_login_method() { + for (plan, expected) in [ + ("lite", "DevPass Lite"), + ("pro", "DevPass Pro"), + ("max", "DevPass Max"), + ] { + let result = parse(&with("devPlan", json!(plan))).unwrap(); + assert_eq!(result.usage.login_method.as_deref(), Some(expected)); + } +} + +#[test] +fn inactive_weekly_window_has_full_allowance_without_an_invented_reset() { + let mut data = with("devPlanPremiumCreditsUsed", json!("0.00")); + data["devPlanPremiumWeekResetsAt"] = Value::Null; + let weekly = parse(&data).unwrap().usage.secondary.unwrap(); + assert_eq!(weekly.used_percent, 0.0); + assert_eq!(weekly.resets_at, None); + assert_eq!(weekly.window_minutes, Some(10080)); +} + +#[test] +fn offset_reset_instants_normalize_to_utc() { + let data = with( + "devPlanPremiumWeekResetsAt", + json!("2026-10-01T14:00:00+02:00"), + ); + let weekly = parse(&data).unwrap().usage.secondary.unwrap(); + assert_eq!( + weekly.resets_at, + Some(Utc.with_ymd_and_hms(2026, 10, 1, 12, 0, 0).unwrap()) + ); +} + +#[test] +fn pay_as_you_go_shows_only_key_scoped_all_time_spend() { + let mut data = with("devPlan", json!("none")); + for key in [ + "devPlanCreditsUsed", + "devPlanCreditsLimit", + "devPlanCreditsRemaining", + "devPlanPremiumWeeklyLimit", + "devPlanPremiumCreditsUsed", + ] { + data[key] = json!("0"); + } + data["devPlanPremiumWeekResetsAt"] = Value::Null; + + let result = parse(&data).unwrap(); + assert!(result.usage.primary.is_informational); + assert!(result.usage.secondary.is_none()); + assert!(result.cost.is_none()); + assert_eq!(result.usage.login_method.as_deref(), Some("Pay as you go")); + let ids: Vec<_> = result.display_details().iter().map(|r| r.id()).collect(); + assert_eq!(ids, ["key-usage"]); + assert_eq!(row(&result, "key-usage").value(), "$31.42"); +} + +#[test] +fn pay_as_you_go_ignores_plan_fields_and_keeps_the_key_limit() { + let data = json!({ + "usage": "1.5", + "limit": "10", + "devPlan": "none", + "devPlanCreditsUsed": 7, + }); + let result = parse(&data).unwrap(); + assert_eq!(row(&result, "key-usage").value(), "$1.50"); + assert_eq!(row(&result, "key-limit").value(), "$10.00"); + assert_eq!(result.display_details().len(), 2); +} + +#[test] +fn zero_limits_omit_bars_and_overages_keep_actual_spend() { + let mut data = with("devPlanCreditsUsed", json!("250")); + data["devPlanCreditsRemaining"] = json!("0"); + data["devPlanPremiumWeeklyLimit"] = json!("0"); + let result = parse(&data).unwrap(); + + assert_eq!(result.usage.primary.used_percent, 100.0); + assert!(result.usage.secondary.is_none()); + assert_eq!(row(&result, "cycle-used").value(), "$250.00 / $237.00"); + let progress = row(&result, "cycle-used").progress().unwrap(); + assert_close(progress.used(), 237.0); + assert_eq!(row(&result, "premium-weekly").value(), "$5.00 / $0.00"); + assert!(row(&result, "premium-weekly").progress().is_none()); +} + +#[test] +fn a_zero_plan_allowance_is_informational() { + let mut data = with("devPlanCreditsLimit", json!("0")); + data["devPlanCreditsUsed"] = json!("0"); + let result = parse(&data).unwrap(); + assert!(result.usage.primary.is_informational); + assert!(row(&result, "cycle-used").progress().is_none()); + assert!(result.usage.secondary.is_some()); +} + +#[test] +fn invalid_monetary_strings_fail_closed() { + let huge = "9".repeat(400); + for value in [ + "", + " ", + "NaN", + "Infinity", + "1e999", + "1e2", + "-1", + "+1", + "0x10", + "1.", + ".5", + "1.2.3", + "1,5", + "25 ", + "２５", + huge.as_str(), + ] { + for key in ["usage", "devPlanCreditsUsed", "devPlanPremiumWeeklyLimit"] { + assert!( + matches!( + parse(&with(key, json!(value))), + Err(ProviderError::Parse(_)) + ), + "{key} = {value:?}" + ); + } + } +} + +#[test] +fn non_string_and_missing_amounts_fail_closed() { + for key in [ + "usage", + "limit", + "devPlanCreditsUsed", + "devPlanCreditsLimit", + "devPlanCreditsRemaining", + "devPlanPremiumWeeklyLimit", + "devPlanPremiumCreditsUsed", + "devPlanPremiumWeekResetsAt", + ] { + assert!( + matches!(parse(&without(key)), Err(ProviderError::Parse(_))), + "missing {key}" + ); + } + for key in ["usage", "limit", "devPlanCreditsUsed"] { + assert!( + matches!(parse(&with(key, json!(12))), Err(ProviderError::Parse(_))), + "numeric {key}" + ); + } + assert!(matches!( + parse(&with("usage", Value::Null)), + Err(ProviderError::Parse(_)) + )); +} + +#[test] +fn unknown_plans_and_invalid_reset_dates_fail_closed() { + for plan in ["unexpected", "", "Pro", "PRO", "team"] { + assert!( + matches!( + parse(&with("devPlan", json!(plan))), + Err(ProviderError::Parse(_)) + ), + "plan {plan:?}" + ); + } + assert!(matches!( + parse(&without("devPlan")), + Err(ProviderError::Parse(_)) + )); + for reset in [ + "unexpected", + "", + "2026-10-01", + "2026-10-01T12:00:00", + "2026-10-01 12:00:00Z", + "2026-10-01t12:00:00z", + "2026-10-01T12:00:00z", + "2026-13-01T12:00:00Z", + "2026-10-01T12:00:60Z", + "2026-10-01T12:00:00+0200", + ] { + assert!( + matches!( + parse(&with("devPlanPremiumWeekResetsAt", json!(reset))), + Err(ProviderError::Parse(_)) + ), + "reset {reset:?}" + ); + } + assert!(matches!( + parse(&with("devPlanPremiumWeekResetsAt", json!(1_790_000_000))), + Err(ProviderError::Parse(_)) + )); +} + +#[test] +fn malformed_responses_are_parse_failures_that_do_not_echo_the_body() { + for raw in [ + "not-json", + "private-response", + "{}", + r#"{"data":null}"#, + r#"{"data":[]}"#, + "[]", + "", + ] { + let error = parse_result(raw.as_bytes()).unwrap_err(); + assert!(matches!(error, ProviderError::Parse(_)), "{raw:?}"); + assert!(!error.to_string().contains("private-response")); + } +} + +#[test] +fn http_failures_are_classified_without_exposing_bodies() { + for (code, expected) in [ + (401, "DevPass API key was rejected or is inactive."), + (403, "DevPass requires a regular gateway API key."), + (429, "DevPass usage requests are rate limited."), + (500, "DevPass usage is temporarily unavailable."), + (503, "DevPass usage is temporarily unavailable."), + (400, "DevPass returned HTTP 400."), + (404, "DevPass returned HTTP 404."), + (204, "DevPass returned HTTP 204."), + ] { + let error = check_status(StatusCode::from_u16(code).unwrap()).unwrap_err(); + assert!(error.to_string().contains(expected), "{code}: {error}"); + } + assert!(matches!( + check_status(StatusCode::UNAUTHORIZED), + Err(ProviderError::OAuthExpired(_)) + )); + assert!(check_status(StatusCode::OK).is_ok()); +} + +fn key_url(server: &mockito::ServerGuard) -> String { + format!("{}/v1/key", server.url()) +} + +#[tokio::test] +async fn fetch_sends_the_bearer_key_and_maps_the_fixture() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/v1/key") + .match_header("authorization", "Bearer fixture-key") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(body(&fixture())) + .create_async() + .await; + + let result = fetch_key( + &DevPassProvider::new().client, + &key_url(&server), + "fixture-key", + ) + .await + .unwrap(); + + mock.assert_async().await; + assert_eq!(result.source_label, "api"); + assert_eq!(row(&result, "cycle-remaining").value(), "$212.00"); +} + +#[tokio::test] +async fn fetch_never_echoes_an_error_body() { + for (status, kind) in [ + (401, "OAuth"), + (403, "Other"), + (429, "Other"), + (503, "Other"), + ] { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/v1/key") + .with_status(status) + .with_body("private-response") + .create_async() + .await; + + let error = fetch_key(&DevPassProvider::new().client, &key_url(&server), "k") + .await + .unwrap_err(); + + assert!( + !error.to_string().contains("private-response"), + "{status}: {error}" + ); + assert!( + format!("{error:?}").starts_with(kind), + "{status}: {error:?}" + ); + } +} + +#[tokio::test] +async fn fetch_does_not_follow_redirects_with_the_key() { + let mut target = mockito::Server::new_async().await; + let leaked = target + .mock("GET", "/v1/key") + .with_status(200) + .with_body(body(&fixture())) + .expect(0) + .create_async() + .await; + let mut origin = mockito::Server::new_async().await; + origin + .mock("GET", "/v1/key") + .with_status(302) + .with_header("location", &key_url(&target)) + .create_async() + .await; + + let error = fetch_key(&DevPassProvider::new().client, &key_url(&origin), "secret") + .await + .unwrap_err(); + + leaked.assert_async().await; + assert!(error.to_string().contains("HTTP 302"), "{error}"); +} + +#[tokio::test] +async fn fetch_rejects_an_oversized_body() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/v1/key") + .with_status(200) + .with_body(vec![b' '; MAX_BODY_BYTES + 1]) + .create_async() + .await; + + let error = fetch_key(&DevPassProvider::new().client, &key_url(&server), "k") + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::Parse(_)), "{error}"); +} + +#[test] +fn provider_registration_matches_upstream_metadata() { + let provider = DevPassProvider::new(); + let metadata = provider.metadata(); + assert_eq!(metadata.display_name, "DevPass"); + assert_eq!(metadata.session_label, "Plan credits"); + assert_eq!(metadata.weekly_label, "Premium weekly"); + assert!(!metadata.default_enabled); + assert_eq!( + metadata.dashboard_url, + Some("https://devpass.llmgateway.io/dashboard") + ); + assert_eq!( + ProviderId::from_cli_name("devpass"), + Some(ProviderId::DevPass) + ); + assert_eq!(ProviderId::DevPass.cli_name(), "devpass"); + assert_eq!( + provider.available_sources(), + [SourceMode::Auto, SourceMode::OAuth] + ); +} diff --git a/rust/src/providers/devpass/wire.rs b/rust/src/providers/devpass/wire.rs new file mode 100644 index 0000000000..b487ea8188 --- /dev/null +++ b/rust/src/providers/devpass/wire.rs @@ -0,0 +1,139 @@ +//! DevPass `/v1/key` wire shapes. Every amount is a plain decimal string and +//! anything else (signs, `NaN`, exponents, hex, empty) fails the parse rather +//! than showing as zero usage. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Deserializer}; + +/// A non-negative, finite amount that arrived as `^\d+(\.\d+)?$`. +#[derive(Debug, Clone, Copy, PartialEq, Deserialize)] +#[serde(try_from = "String")] +pub(super) struct Money(f64); + +impl Money { + pub(super) fn value(self) -> f64 { + self.0 + } +} + +impl TryFrom for Money { + type Error = &'static str; + + fn try_from(raw: String) -> Result { + let (whole, fraction) = raw.split_once('.').unwrap_or((&raw, "0")); + let digits = |part: &str| !part.is_empty() && part.bytes().all(|b| b.is_ascii_digit()); + if !(digits(whole) && digits(fraction)) { + return Err("not a plain decimal amount"); + } + match raw.parse::() { + Ok(value) if value.is_finite() => Ok(Self(value)), + _ => Err("amount is out of range"), + } + } +} + +/// An ISO-8601 instant with an explicit `Z` or numeric offset. +#[derive(Debug, Clone, Copy, PartialEq, Deserialize)] +#[serde(try_from = "String")] +struct ResetInstant(DateTime); + +impl TryFrom for ResetInstant { + type Error = &'static str; + + fn try_from(raw: String) -> Result { + // chrono also accepts a lowercase `t`/`z` and a leap second, which + // the upstream strict shape and `Date` parsing reject. + if raw.as_bytes().get(10) != Some(&b'T') || raw.ends_with('z') { + return Err("not a strict ISO-8601 instant"); + } + let instant = DateTime::parse_from_rfc3339(&raw).map_err(|_| "invalid instant")?; + if instant.timestamp_subsec_nanos() >= 1_000_000_000 { + return Err("leap seconds are not valid instants"); + } + Ok(Self(instant.with_timezone(&Utc))) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(super) enum DevPlan { + None, + Lite, + Pro, + Max, +} + +impl DevPlan { + pub(super) fn login_method(self) -> &'static str { + match self { + Self::None => "Pay as you go", + Self::Lite => "DevPass Lite", + Self::Pro => "DevPass Pro", + Self::Max => "DevPass Max", + } + } +} + +/// A field that must be present but may be `null`; a missing field is a +/// parse failure, not `None`. +fn required_nullable<'de, D, T>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, + T: Deserialize<'de>, +{ + Option::::deserialize(deserializer) +} + +fn required_nullable_instant<'de, D>(deserializer: D) -> Result>, D::Error> +where + D: Deserializer<'de>, +{ + Ok(required_nullable::(deserializer)?.map(|instant| instant.0)) +} + +#[derive(Debug, Deserialize)] +struct Envelope { + data: T, +} + +/// Fields every key returns, including pay-as-you-go keys. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(super) struct KeyData { + pub(super) usage: Money, + #[serde(deserialize_with = "required_nullable")] + pub(super) limit: Option, + pub(super) dev_plan: DevPlan, +} + +/// Plan fields, required only when the key has a DevPass plan. +#[derive(Debug, Deserialize)] +pub(super) struct PlanFields { + #[serde(rename = "devPlanCreditsUsed")] + pub(super) credits_used: Money, + #[serde(rename = "devPlanCreditsLimit")] + pub(super) credits_limit: Money, + #[serde(rename = "devPlanCreditsRemaining")] + pub(super) credits_remaining: Money, + #[serde(rename = "devPlanPremiumWeeklyLimit")] + pub(super) premium_limit: Money, + #[serde(rename = "devPlanPremiumCreditsUsed")] + pub(super) premium_used: Money, + #[serde( + rename = "devPlanPremiumWeekResetsAt", + deserialize_with = "required_nullable_instant" + )] + pub(super) premium_resets_at: Option>, +} + +pub(super) fn parse_key(body: &[u8]) -> Option { + serde_json::from_slice::>(body) + .ok() + .map(|envelope| envelope.data) +} + +pub(super) fn parse_plan(body: &[u8]) -> Option { + serde_json::from_slice::>(body) + .ok() + .map(|envelope| envelope.data) +} diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index de9d265f6e..cd48fe326b 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -33,6 +33,7 @@ pub mod deepgram; pub mod deepinfra; pub mod deepseek; pub mod devin; +pub mod devpass; pub mod doubao; pub mod elevenlabs; pub mod factory; @@ -116,6 +117,7 @@ pub use deepgram::DeepgramProvider; pub use deepinfra::DeepInfraProvider; pub use deepseek::DeepSeekProvider; pub use devin::DevinProvider; +pub use devpass::DevPassProvider; pub use doubao::DoubaoProvider; pub use elevenlabs::ElevenLabsProvider; pub use factory::FactoryProvider; diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index c2adf3565c..4dd4f55ecc 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -426,6 +426,17 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: Some("https://portal.neuralwatt.com/dashboard"), }, + ProviderConfigInfo { + id: ProviderId::DevPass, + name: "DevPass", + requires_api_key: true, + api_key_env_var: Some("DEVPASS_API_KEY"), + api_key_help: Some( + "Use a regular LLM Gateway API key. Publishable keys and end-user sessions cannot read plan state.", + ), + config_file_path: None, + dashboard_url: Some("https://devpass.llmgateway.io/dashboard"), + }, ProviderConfigInfo { id: ProviderId::Doubao, name: "Doubao / Volcengine Ark", From 55ff7fb2059856850b83750f9660ce6bf9ddb0c6 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:37:43 +0700 Subject: [PATCH 126/407] Port upstream 0.67.0: add Stay Awake for live local agent sessions --- .../src-tauri/src/commands/bridge.rs | 2 + .../src-tauri/src/commands/settings.rs | 8 + apps/desktop-tauri/src-tauri/src/main.rs | 11 +- .../desktop-tauri/src-tauri/src/stay_awake.rs | 91 +++++ .../src-tauri/src/tray_bridge.rs | 7 +- apps/desktop-tauri/src-tauri/src/tray_menu.rs | 54 ++- apps/desktop-tauri/src/i18n/keys.ts | 3 + .../settings/tabs/AdvancedTab.test.tsx | 22 + .../surfaces/settings/tabs/AdvancedTab.tsx | 11 + apps/desktop-tauri/src/types/bridge.ts | 3 + rust/Cargo.toml | 1 + rust/src/agent_sessions.rs | 1 + rust/src/agent_sessions/native_processes.rs | 123 ++++++ rust/src/agent_sessions/parsers.rs | 57 ++- rust/src/lib.rs | 1 + rust/src/locale.rs | 3 + rust/src/locale/en-US.ftl | 3 + rust/src/power_assertion.rs | 377 ++++++++++++++++++ rust/src/settings.rs | 5 + rust/src/settings/raw.rs | 4 + 20 files changed, 756 insertions(+), 31 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/stay_awake.rs create mode 100644 rust/src/agent_sessions/native_processes.rs create mode 100644 rust/src/power_assertion.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..ed832d77ef 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -678,6 +678,7 @@ pub struct SettingsSnapshot { global_shortcut: String, codex_custom_sessions_dirs: Vec, agent_sessions_enabled: bool, + stay_awake_enabled: bool, agent_session_ssh_hosts: Vec, hooks_enabled: bool, http_proxy_enabled: bool, @@ -800,6 +801,7 @@ impl From for SettingsSnapshot { global_shortcut: settings.global_shortcut, codex_custom_sessions_dirs: settings.codex_custom_sessions_dirs, agent_sessions_enabled: settings.agent_sessions_enabled, + stay_awake_enabled: settings.stay_awake_enabled, agent_session_ssh_hosts: settings.agent_session_ssh_hosts, hooks_enabled: settings.hooks_enabled, http_proxy_enabled: settings.http_proxy_enabled, diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..e5cb97d4b2 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -43,6 +43,7 @@ pub struct SettingsUpdate { pub global_shortcut: Option, pub codex_custom_sessions_dirs: Option>, pub agent_sessions_enabled: Option, + pub stay_awake_enabled: Option, pub agent_session_ssh_hosts: Option>, pub hooks_enabled: Option, pub http_proxy_enabled: Option, @@ -327,6 +328,9 @@ impl SettingsUpdate { if let Some(v) = self.agent_sessions_enabled { settings.agent_sessions_enabled = v; } + if let Some(v) = self.stay_awake_enabled { + settings.stay_awake_enabled = v; + } if let Some(v) = self.agent_session_ssh_hosts.clone() { settings.agent_session_ssh_hosts = codexbar::agent_sessions::RemoteSessionFetcher::sanitized_hosts(&v); @@ -509,6 +513,7 @@ pub async fn update_settings( let clear_local_usage_cache = patch.codex_custom_sessions_dirs.is_some(); let rebuild_tray_menu = patch.rebuilds_tray_menu(); let refresh_tray_presentation = patch.refreshes_tray_presentation(); + let stay_awake_changed = patch.stay_awake_enabled.is_some(); let tray_promotion_changed = patch.changes_tray_promotion(); let tray_panel_always_on_top_changed = patch.tray_panel_always_on_top.is_some(); let previous_promoted = settings.promote_tray_icon; @@ -536,6 +541,9 @@ pub async fn update_settings( } crate::floatbar::after_settings_saved(&app, &float_bar_patch, &settings, notify_float_bar); + if stay_awake_changed { + crate::stay_awake::settings_changed(&app, settings.stay_awake_enabled); + } if rebuild_tray_menu { crate::tray_bridge::rebuild_tray_menu(&app); } diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index 02f48672b6..e12b4ad0cd 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -14,6 +14,7 @@ mod proof_harness; mod shell; mod shortcut_bridge; mod state; +mod stay_awake; mod surface; mod surface_target; mod tray_accounts; @@ -288,6 +289,7 @@ fn main() { shortcut_bridge::register(app.handle()); floatbar::install(app.handle()); auto_refresh::install(app.handle().clone()); + stay_awake::install(app.handle().clone()); if settings.powertoys_status_pipe_enabled { powertoys::install(app.handle().clone()); } @@ -415,8 +417,13 @@ fn main() { _ => {} } }) - .run(tauri::generate_context!()) - .expect("failed to run CodexBar desktop shell"); + .build(tauri::generate_context!()) + .expect("failed to build CodexBar desktop shell") + .run(|_app, event| { + if matches!(event, tauri::RunEvent::Exit) { + stay_awake::shutdown(); + } + }); } #[cfg(test)] diff --git a/apps/desktop-tauri/src-tauri/src/stay_awake.rs b/apps/desktop-tauri/src-tauri/src/stay_awake.rs new file mode 100644 index 0000000000..b9fd02ac6f --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/stay_awake.rs @@ -0,0 +1,91 @@ +//! Stay Awake scheduler: keeps Windows from idle-sleeping while a local agent +//! (Codex, Claude, Pi/OMP) process is running. +//! +//! While `stay_awake_enabled` is on, a local-only process scan runs every 30 +//! seconds, independent of `agent_sessions_enabled`. It never touches SSH or +//! Tailscale hosts and never reads transcript files, so file-only rollouts and +//! remote sessions cannot hold the machine awake. The policy itself lives in +//! `codexbar::power_assertion`. The setting reaches the controller through +//! [`settings_changed`] (every save goes through `update_settings`) and once at +//! startup; the scan loop never re-reads the file, so it cannot re-enable a +//! toggle the user just turned off. + +use std::sync::{Mutex, MutexGuard, OnceLock}; +use std::time::Duration; + +use codexbar::agent_sessions::LocalAgentSessionScanner; +use codexbar::power_assertion::{StayAwake, SystemPowerAssertion}; +use codexbar::settings::Settings; +use tokio::sync::Notify; + +const SCAN_INTERVAL: Duration = Duration::from_secs(30); + +type Controller = StayAwake; + +static CONTROLLER: OnceLock> = OnceLock::new(); +static WAKE: Notify = Notify::const_new(); + +fn controller() -> MutexGuard<'static, Controller> { + CONTROLLER + .get_or_init(|| Mutex::new(StayAwake::new(SystemPowerAssertion::new()))) + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// Whether an idle-sleep prevention is currently held (tray status line). +pub fn is_held() -> bool { + controller().is_held() +} + +/// Apply a saved settings change: turning Stay Awake off releases right away, +/// turning it on triggers a scan without waiting for the next interval. +pub fn settings_changed(app: &tauri::AppHandle, enabled: bool) { + let changed = controller().set_enabled(enabled); + if changed { + crate::tray_bridge::rebuild_tray_menu(app); + } + if enabled { + WAKE.notify_one(); + } +} + +/// Release for good on app exit. +pub fn shutdown() { + controller().shutdown(); +} + +pub fn install(app: tauri::AppHandle) { + controller().set_enabled(Settings::load().stay_awake_enabled); + tauri::async_runtime::spawn(async move { + loop { + scan_once(&app).await; + // Sleep one interval, or wake early when the setting is turned on. + tokio::time::timeout(SCAN_INTERVAL, WAKE.notified()) + .await + .ok(); + } + }); +} + +async fn scan_once(app: &tauri::AppHandle) { + if !controller().is_enabled() { + return; + } + let changed = + match tokio::task::spawn_blocking(LocalAgentSessionScanner::has_live_agent_process).await { + // The controller drops this result if the setting was turned off or + // the app shut down while the scan was running. + Ok(Ok(live)) => controller().observe(live), + Ok(Err(error)) => { + tracing::warn!("Stay Awake process scan failed: {error}"); + false + } + Err(error) => { + tracing::warn!("Stay Awake scan task failed: {error}"); + false + } + }; + if changed { + crate::tray_bridge::rebuild_tray_menu(app); + } +} diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index aebb7c2115..7ad4e455fa 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -20,7 +20,7 @@ use crate::surface::SurfaceMode; 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_menu::{MenuState, TrayMenuEntry, build_tray_menu_with}; #[derive(Debug, Clone, Copy)] struct MonitorScaleInfo { @@ -134,7 +134,10 @@ fn build_native_tray_menu( providers, status_labels, &enabled, - settings.float_bar_enabled, + MenuState { + float_bar_enabled: settings.float_bar_enabled, + stay_awake_held: crate::stay_awake::is_held(), + }, settings.ui_language, ); crate::tray_accounts::prepend_account_menus(&mut spec, &settings); diff --git a/apps/desktop-tauri/src-tauri/src/tray_menu.rs b/apps/desktop-tauri/src-tauri/src/tray_menu.rs index 688e56b42b..9f60c843b5 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_menu.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_menu.rs @@ -82,6 +82,14 @@ impl TrayMenuEntry { } } +/// Live app state the menu mirrors beyond the provider catalog. +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct MenuState { + pub(crate) float_bar_enabled: bool, + /// Stay Awake currently holds an idle-sleep prevention. + pub(crate) stay_awake_held: bool, +} + #[cfg(test)] pub(crate) fn build_tray_menu( providers: &[ProviderCatalogEntry], @@ -92,7 +100,7 @@ pub(crate) fn build_tray_menu( providers, status_labels, enabled_providers, - false, + MenuState::default(), Language::English, ) } @@ -101,7 +109,7 @@ pub(crate) fn build_tray_menu_with( providers: &[ProviderCatalogEntry], status_labels: &[(String, String)], enabled_providers: &HashSet, - float_bar_enabled: bool, + state: MenuState, lang: Language, ) -> Vec { let mut menu: Vec = Vec::new(); @@ -111,7 +119,13 @@ pub(crate) fn build_tray_menu_with( for (id, label) in status_labels { menu.push(TrayMenuEntry::status_row(format!("status_{id}"), label)); } - if !status_labels.is_empty() { + if state.stay_awake_held { + menu.push(TrayMenuEntry::status_row( + "stay_awake_active", + text(LocaleKey::TrayStayAwakeActive), + )); + } + if !status_labels.is_empty() || state.stay_awake_held { menu.push(TrayMenuEntry::separator()); } @@ -130,7 +144,7 @@ pub(crate) fn build_tray_menu_with( menu.push(TrayMenuEntry::check_item( "toggle_float_bar", text(LocaleKey::TrayShowFloatBar), - float_bar_enabled, + state.float_bar_enabled, )); menu.push(TrayMenuEntry::separator()); @@ -239,7 +253,10 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - /* float_bar_enabled = */ true, + MenuState { + float_bar_enabled: true, + ..MenuState::default() + }, Language::English, ); let toggle = menu_on @@ -253,7 +270,7 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - /* float_bar_enabled = */ false, + MenuState::default(), Language::English, ); let toggle = menu_off @@ -269,7 +286,7 @@ mod tests { &sample_provider_catalog(), &[], &both_enabled(), - false, + MenuState::default(), Language::Japanese, ); fn label_for<'a>(menu: &'a [TrayMenuEntry], id: &'a str) -> &'a str { @@ -311,4 +328,27 @@ mod tests { // Third item should be a separator. assert!(menu[2].is_separator); } + + #[test] + fn stay_awake_row_appears_only_while_held() { + let build = |stay_awake_held| { + build_tray_menu_with( + &sample_provider_catalog(), + &[], + &both_enabled(), + MenuState { + stay_awake_held, + ..MenuState::default() + }, + Language::English, + ) + }; + assert!(!menu_contains(&build(false), "stay_awake_active")); + + let held = build(true); + assert_eq!(held[0].id.as_deref(), Some("stay_awake_active")); + assert_eq!(held[0].label, "Stay Awake: local agent session is live"); + assert!(held[0].disabled); + assert!(held[1].is_separator); + } } diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..6eb0f17f1e 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -226,6 +226,7 @@ export const ALL_LOCALE_KEYS = [ "TrayPopOutDashboard", "TrayShowWindow", "TrayShowFloatBar", + "TrayStayAwakeActive", "TrayRefreshAll", "TrayProviders", "TraySettings", @@ -443,6 +444,8 @@ export const ALL_LOCALE_KEYS = [ "AgentSessionsTitle", "AgentSessionsEnableLabel", "AgentSessionsEnableHelper", + "AgentSessionsStayAwakeLabel", + "AgentSessionsStayAwakeHelper", "AgentSessionsSshHostsLabel", "AgentSessionsSshHostsHelper", "AgentSessionsLoading", diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx index 8d1d7f2148..18611b79d6 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx @@ -183,4 +183,26 @@ describe("AdvancedTab", () => { ).toHaveLength(0); expect(screen.getAllByText("HooksEnableLabel")).toHaveLength(1); }); + + it("toggles Stay Awake independently of agent session discovery", () => { + const set = vi.fn(); + render( + , + ); + + const toggle = screen + .getByText("AgentSessionsStayAwakeLabel") + .closest(".settings-field") + ?.querySelector('input[type="checkbox"]'); + expect(toggle).not.toBeNull(); + expect(toggle).not.toBeDisabled(); + expect(toggle).not.toBeChecked(); + + fireEvent.click(toggle as Element); + expect(set).toHaveBeenCalledWith({ stayAwakeEnabled: true }); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx index fcf98d09c3..ee278cc371 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.tsx @@ -192,6 +192,17 @@ export default function AdvancedTab({ settings, set, saving }: TabProps) { onChange={(v) => set({ agentSessionsEnabled: v })} /> + + set({ stayAwakeEnabled: v })} + /> + Result { + Ok(Self::holds_system_awake(&snapshot_processes()?)) + } + + /// Pure predicate behind [`Self::has_live_agent_process`]. + pub(super) fn holds_system_awake(processes: &[AgentProcessRecord]) -> bool { + processes + .iter() + .any(|process| process.pid > 0 && process.is_agent()) + } +} + +#[cfg(windows)] +fn snapshot_processes() -> Result, String> { + use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle}; + use windows::Win32::Foundation::HANDLE; + use windows::Win32::System::Diagnostics::ToolHelp::{ + CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW, + TH32CS_SNAPPROCESS, + }; + + // SAFETY: a successful call returns a unique snapshot handle owned below. + let snapshot = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) } + .map_err(|error| format!("process snapshot failed: {error}"))?; + // SAFETY: `snapshot` is a valid, unique handle returned just above; the + // `OwnedHandle` closes it when dropped. + let snapshot = unsafe { OwnedHandle::from_raw_handle(snapshot.0) }; + let handle = HANDLE(snapshot.as_raw_handle()); + + let mut entry = PROCESSENTRY32W { + dwSize: u32::try_from(std::mem::size_of::()) + .map_err(|error| format!("invalid process entry size: {error}"))?, + ..Default::default() + }; + let mut records = Vec::new(); + // SAFETY: the snapshot handle is live and `entry.dwSize` is correct. + if unsafe { Process32FirstW(handle, &mut entry) }.is_ok() { + loop { + let len = entry + .szExeFile + .iter() + .position(|&unit| unit == 0) + .unwrap_or(entry.szExeFile.len()); + let name = String::from_utf16_lossy(&entry.szExeFile[..len]); + records.push(WindowsProcessOutputParser::record( + entry.th32ProcessID, + entry.th32ParentProcessID, + None, + Some(name), + None, + )); + // SAFETY: `entry` is valid for the next snapshot entry. + if unsafe { Process32NextW(handle, &mut entry) }.is_err() { + break; + } + } + } + Ok(records) +} + +#[cfg(not(windows))] +fn snapshot_processes() -> Result, String> { + Err("native process snapshots are only available on Windows".to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn record(pid: u32, name: &str) -> AgentProcessRecord { + WindowsProcessOutputParser::record(pid, 1, None, Some(name.to_string()), None) + } + + #[test] + fn agent_processes_hold_the_system_awake() { + for name in ["claude.exe", "codex.exe"] { + let processes = [record(10, "explorer.exe"), record(11, name)]; + assert!( + LocalAgentSessionScanner::holds_system_awake(&processes), + "{name} should count" + ); + } + } + + #[test] + fn helpers_and_unrelated_processes_do_not_hold_the_system_awake() { + let processes = [ + record(10, "explorer.exe"), + record(11, "Codex (Renderer)"), + record(12, "codex app-server"), + record(13, "node.exe"), + ]; + assert!(!LocalAgentSessionScanner::holds_system_awake(&processes)); + assert!(!LocalAgentSessionScanner::holds_system_awake(&[])); + } + + #[test] + fn pid_zero_never_counts() { + let processes = [record(0, "claude.exe")]; + assert!(!LocalAgentSessionScanner::holds_system_awake(&processes)); + } + + #[cfg(windows)] + #[test] + fn native_snapshot_lists_the_current_process() { + let processes = snapshot_processes().expect("snapshot"); + assert!(processes.iter().any(|p| p.pid == std::process::id())); + } +} diff --git a/rust/src/agent_sessions/parsers.rs b/rust/src/agent_sessions/parsers.rs index 49e236e8e6..b29b4ea598 100644 --- a/rust/src/agent_sessions/parsers.rs +++ b/rust/src/agent_sessions/parsers.rs @@ -139,32 +139,49 @@ impl WindowsProcessOutputParser { .filter_map(|value| serde_json::from_value::(value).ok()) .filter(|process| process.process_id > 0 && seen.insert(process.process_id)) .map(|process| { - let display_name = process - .name - .as_deref() - .filter(|name| !name.trim().is_empty()) - .or(process.executable_path.as_deref()) - .unwrap_or_default(); - let classification = classify_process_command(display_name); - AgentProcessRecord { - pid: process.process_id, - ppid: process.parent_process_id, - started_at: process + Self::record( + process.process_id, + process.parent_process_id, + process .creation_date .as_deref() .and_then(parse_windows_creation_date), - provider: classification.provider, - source: classification.source, - executable: process.name.unwrap_or(classification.executable), - kind: classification.kind, - // Windows deliberately does not harvest other processes' - // command lines (privacy); ps-based remote scans carry - // them, and tests inject them. - command: None, - } + process.name, + process.executable_path.as_deref(), + ) }) .collect()) } + + /// Classify one Windows process. The image name is preferred over the + /// executable path, matching the CIM query and the native snapshot. + pub(super) fn record( + pid: u32, + ppid: u32, + started_at: Option>, + name: Option, + executable_path: Option<&str>, + ) -> AgentProcessRecord { + let display_name = name + .as_deref() + .filter(|name| !name.trim().is_empty()) + .or(executable_path) + .unwrap_or_default(); + let classification = classify_process_command(display_name); + AgentProcessRecord { + pid, + ppid, + started_at, + provider: classification.provider, + source: classification.source, + executable: name.unwrap_or(classification.executable), + kind: classification.kind, + // Windows deliberately does not harvest other processes' + // command lines (privacy); ps-based remote scans carry + // them, and tests inject them. + command: None, + } + } } fn parse_windows_creation_date(value: &str) -> Option> { diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 1dfa17cff5..b528d1fef7 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -19,6 +19,7 @@ pub mod login; #[cfg(windows)] pub mod managed_process; pub mod notifications; +pub mod power_assertion; pub mod providers; pub mod secure_file; pub mod settings; diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..658ecad57a 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -454,6 +454,7 @@ locale_keys! { TrayPopOutDashboard, TrayShowWindow, TrayShowFloatBar, + TrayStayAwakeActive, TrayRefreshAll, TrayProviders, TraySettings, @@ -687,6 +688,8 @@ locale_keys! { AgentSessionsTitle, AgentSessionsEnableLabel, AgentSessionsEnableHelper, + AgentSessionsStayAwakeLabel, + AgentSessionsStayAwakeHelper, AgentSessionsSshHostsLabel, AgentSessionsSshHostsHelper, AgentSessionsLoading, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..977a040d7c 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -211,6 +211,7 @@ TrayOpenCodexBar = Pop Out Dashboard TrayPopOutDashboard = Pop Out Dashboard TrayShowWindow = Show Window TrayShowFloatBar = Show Float Bar +TrayStayAwakeActive = Stay Awake: local agent session is live TrayRefreshAll = Refresh All TrayProviders = Providers TraySettings = Settings... @@ -380,6 +381,8 @@ AgentSessionsProviderPi = Pi AgentSessionsProviderOmp = OMP AgentSessionsEnableLabel = Show active Codex, Claude, and Pi/OMP sessions AgentSessionsEnableHelper = Discover local sessions and optional SSH hosts. Disabled by default. +AgentSessionsStayAwakeLabel = Stay awake while a local agent session is live +AgentSessionsStayAwakeHelper = Prevents idle sleep while a local Codex, Claude, or Pi/OMP process is running. The display can still turn off, and closing the lid still sleeps the PC. Disabled by default. AgentSessionsSshHostsLabel = SSH hosts AgentSessionsSshHostsHelper = Comma-separated SSH targets. Connections use non-interactive mode and short timeouts. AgentSessionsLoading = Looking for agent sessions... diff --git a/rust/src/power_assertion.rs b/rust/src/power_assertion.rs new file mode 100644 index 0000000000..b3dc426dd8 --- /dev/null +++ b/rust/src/power_assertion.rs @@ -0,0 +1,377 @@ +//! Stay Awake: hold one idle-sleep prevention while a local agent session is live. +//! +//! Port of upstream CodexBar 0.67.0 `AgentSessionPowerAssertion` and the +//! `updatePowerAssertion` policy in `AgentSessionsStore`. [`PowerAssertion`] is +//! the platform seam (injectable for tests); [`StayAwake`] is the pure policy +//! that decides when to acquire and release it. +//! +//! Only the system is kept awake, never the display: no display-required flag +//! is set, and lid close or an explicit sleep still sleeps the machine. + +/// Platform seam for one system idle-sleep prevention. +pub trait PowerAssertion: Send { + /// Start preventing idle system sleep. Returns `false` when the OS refused. + fn acquire(&mut self) -> bool; + /// Stop preventing idle system sleep. Safe to call when nothing is held. + fn release(&mut self); +} + +/// Stay Awake policy: at most one assertion, held only while enabled and a +/// live local agent session exists. +/// +/// Every method returns whether the held state changed so callers can refresh +/// UI that mirrors it. Scan results arrive asynchronously, so [`observe`] is +/// ignored once the feature is disabled or shut down: a stale scan completion +/// can never acquire after the toggle was turned off. +/// +/// [`observe`]: StayAwake::observe +pub struct StayAwake { + assertion: A, + enabled: bool, + shut_down: bool, + held: bool, +} + +impl StayAwake { + pub fn new(assertion: A) -> Self { + Self { + assertion, + enabled: false, + shut_down: false, + held: false, + } + } + + pub fn is_held(&self) -> bool { + self.held + } + + pub fn is_enabled(&self) -> bool { + self.enabled + } + + /// Apply the user setting. Turning it off releases immediately; turning it + /// on acquires nothing until a scan reports a live session. + pub fn set_enabled(&mut self, enabled: bool) -> bool { + self.enabled = enabled; + if enabled { false } else { self.release() } + } + + /// Record one scan outcome. A failed acquire leaves nothing held, so the + /// next scan that still sees a live session retries. + pub fn observe(&mut self, live_local_session: bool) -> bool { + if !self.enabled || self.shut_down { + return false; + } + match (live_local_session, self.held) { + (true, false) => { + self.held = self.assertion.acquire(); + self.held + } + (false, true) => self.release(), + _ => false, + } + } + + /// Release for good on app exit; later scans and toggles are ignored. + pub fn shutdown(&mut self) -> bool { + self.shut_down = true; + self.enabled = false; + self.release() + } + + fn release(&mut self) -> bool { + if !self.held { + return false; + } + self.assertion.release(); + self.held = false; + true + } +} + +#[cfg(windows)] +pub use windows_impl::SystemPowerAssertion; + +/// Holds `ES_SYSTEM_REQUIRED` for the process. +/// +/// `SetThreadExecutionState(ES_CONTINUOUS | ...)` is scoped to the calling +/// thread and lapses when that thread ends, so a dedicated long-lived thread +/// owns it. Calling it from a tokio worker would drop the request whenever the +/// runtime retires or reuses that thread. +#[cfg(windows)] +mod windows_impl { + use super::PowerAssertion; + use std::sync::mpsc::{self, Sender}; + use std::thread::JoinHandle; + use std::time::Duration; + use windows::Win32::System::Power::{ + ES_CONTINUOUS, ES_SYSTEM_REQUIRED, EXECUTION_STATE, SetThreadExecutionState, + }; + + const ACQUIRE_TIMEOUT: Duration = Duration::from_secs(2); + + enum Request { + Hold(Sender), + Clear, + } + + struct Worker { + requests: Sender, + thread: JoinHandle<()>, + } + + #[derive(Default)] + pub struct SystemPowerAssertion { + worker: Option, + } + + impl SystemPowerAssertion { + pub fn new() -> Self { + Self::default() + } + + fn worker(&mut self) -> Option<&Worker> { + if self.worker.is_none() { + let (requests, inbox) = mpsc::channel(); + let thread = std::thread::Builder::new() + .name("codexbar-stay-awake".to_string()) + .spawn(move || run(inbox)) + .map_err(|error| tracing::warn!("Stay Awake thread failed to start: {error}")) + .ok()?; + self.worker = Some(Worker { requests, thread }); + } + self.worker.as_ref() + } + } + + impl PowerAssertion for SystemPowerAssertion { + fn acquire(&mut self) -> bool { + let (reply, answer) = mpsc::channel(); + let Some(worker) = self.worker() else { + return false; + }; + if worker.requests.send(Request::Hold(reply)).is_err() { + // The thread is gone; the next attempt starts a fresh one. + self.worker = None; + return false; + } + answer.recv_timeout(ACQUIRE_TIMEOUT).unwrap_or(false) + } + + fn release(&mut self) { + if let Some(worker) = &self.worker + && worker.requests.send(Request::Clear).is_err() + { + self.worker = None; + } + } + } + + impl Drop for SystemPowerAssertion { + fn drop(&mut self) { + if let Some(Worker { requests, thread }) = self.worker.take() { + drop(requests); + thread.join().ok(); + } + } + } + + fn run(inbox: mpsc::Receiver) { + while let Ok(request) = inbox.recv() { + match request { + Request::Hold(reply) => { + reply + .send(set_state(ES_CONTINUOUS | ES_SYSTEM_REQUIRED)) + .ok(); + } + Request::Clear => { + set_state(ES_CONTINUOUS); + } + } + } + // Channel closed: clear explicitly instead of relying on thread exit. + set_state(ES_CONTINUOUS); + } + + /// `SetThreadExecutionState` returns the previous state, or 0 on failure. + fn set_state(flags: EXECUTION_STATE) -> bool { + // SAFETY: plain Win32 call with valid flags; it only touches the + // calling thread's execution state. + unsafe { SetThreadExecutionState(flags) }.0 != 0 + } +} + +#[cfg(not(windows))] +pub use fallback::SystemPowerAssertion; + +/// Other platforms have no implementation; nothing is ever held. +#[cfg(not(windows))] +mod fallback { + use super::PowerAssertion; + + #[derive(Default)] + pub struct SystemPowerAssertion; + + impl SystemPowerAssertion { + pub fn new() -> Self { + Self + } + } + + impl PowerAssertion for SystemPowerAssertion { + fn acquire(&mut self) -> bool { + false + } + + fn release(&mut self) {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Default)] + struct Counts { + acquires: u32, + releases: u32, + fail_next_acquires: u32, + } + + #[derive(Clone, Default)] + struct FakeAssertion(Arc>); + + impl FakeAssertion { + fn counts(&self) -> (u32, u32) { + let counts = self.0.lock().unwrap(); + (counts.acquires, counts.releases) + } + + fn fail_next(&self, times: u32) { + self.0.lock().unwrap().fail_next_acquires = times; + } + } + + impl PowerAssertion for FakeAssertion { + fn acquire(&mut self) -> bool { + let mut counts = self.0.lock().unwrap(); + counts.acquires += 1; + if counts.fail_next_acquires > 0 { + counts.fail_next_acquires -= 1; + return false; + } + true + } + + fn release(&mut self) { + self.0.lock().unwrap().releases += 1; + } + } + + fn enabled() -> (StayAwake, FakeAssertion) { + let fake = FakeAssertion::default(); + let mut stay_awake = StayAwake::new(fake.clone()); + stay_awake.set_enabled(true); + (stay_awake, fake) + } + + #[test] + fn disabled_never_acquires() { + let fake = FakeAssertion::default(); + let mut stay_awake = StayAwake::new(fake.clone()); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn enabling_alone_acquires_nothing() { + let (stay_awake, fake) = enabled(); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn live_session_holds_exactly_one_assertion_until_it_disappears() { + let (mut stay_awake, fake) = enabled(); + assert!(stay_awake.observe(true), "appearing session acquires"); + assert!(!stay_awake.observe(true), "idle session keeps the hold"); + assert!(!stay_awake.observe(true)); + assert!(stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 0)); + + assert!(stay_awake.observe(false), "disappearing session releases"); + assert!(!stay_awake.observe(false)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + + assert!(stay_awake.observe(true), "a later session holds again"); + assert_eq!(fake.counts(), (2, 1)); + } + + #[test] + fn toggling_off_releases_immediately() { + let (mut stay_awake, fake) = enabled(); + stay_awake.observe(true); + assert!(stay_awake.set_enabled(false)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + assert!(!stay_awake.set_enabled(false), "already released"); + assert_eq!(fake.counts(), (1, 1)); + } + + #[test] + fn stale_scan_after_disable_does_not_acquire() { + let (mut stay_awake, fake) = enabled(); + stay_awake.set_enabled(false); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (0, 0)); + } + + #[test] + fn shutdown_releases_and_blocks_later_acquires() { + let (mut stay_awake, fake) = enabled(); + stay_awake.observe(true); + assert!(stay_awake.shutdown()); + assert_eq!(fake.counts(), (1, 1)); + + assert!(!stay_awake.observe(true)); + stay_awake.set_enabled(true); + assert!(!stay_awake.observe(true)); + assert!(!stay_awake.is_held()); + assert_eq!(fake.counts(), (1, 1)); + } + + #[test] + fn failed_acquire_retries_on_the_next_scan() { + let (mut stay_awake, fake) = enabled(); + fake.fail_next(1); + assert!(!stay_awake.observe(true), "refused acquire is not held"); + assert!(!stay_awake.is_held()); + assert!(stay_awake.observe(true), "next scan retries"); + assert!(stay_awake.is_held()); + assert_eq!(fake.counts(), (2, 0)); + } + + #[test] + fn failed_acquire_then_no_session_never_releases() { + let (mut stay_awake, fake) = enabled(); + fake.fail_next(1); + stay_awake.observe(true); + assert!(!stay_awake.observe(false)); + assert_eq!(fake.counts(), (1, 0)); + } + + #[cfg(windows)] + #[test] + fn system_assertion_acquires_and_releases_on_its_own_thread() { + let mut assertion = SystemPowerAssertion::new(); + assert!(assertion.acquire()); + assertion.release(); + assert!(assertion.acquire(), "can be re-acquired after release"); + drop(assertion); + } +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..03764a68a5 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -220,6 +220,10 @@ pub struct Settings { #[serde(default)] pub agent_sessions_enabled: bool, + /// Keep the system awake (no idle sleep) while a local agent session is live. + #[serde(default)] + pub stay_awake_enabled: bool, + /// SSH targets queried for remote agent sessions. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub agent_session_ssh_hosts: Vec, @@ -562,6 +566,7 @@ impl Default for Settings { global_shortcut: default_global_shortcut(), // Ctrl+Shift+U by default codex_custom_sessions_dirs: Vec::new(), agent_sessions_enabled: false, + stay_awake_enabled: false, agent_session_ssh_hosts: Vec::new(), hooks_enabled: false, http_proxy_enabled: false, diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..b5a5222354 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -121,6 +121,8 @@ pub(super) struct RawSettings { global_shortcut: String, codex_custom_sessions_dirs: Vec, agent_sessions_enabled: bool, + #[serde(default)] + stay_awake_enabled: bool, agent_session_ssh_hosts: Vec, #[serde(default)] hooks_enabled: bool, @@ -255,6 +257,7 @@ impl Default for RawSettings { global_shortcut: s.global_shortcut, codex_custom_sessions_dirs: s.codex_custom_sessions_dirs, agent_sessions_enabled: s.agent_sessions_enabled, + stay_awake_enabled: s.stay_awake_enabled, agent_session_ssh_hosts: s.agent_session_ssh_hosts, hooks_enabled: s.hooks_enabled, http_proxy_enabled: s.http_proxy_enabled, @@ -558,6 +561,7 @@ impl From for Settings { global_shortcut: raw.global_shortcut, codex_custom_sessions_dirs: raw.codex_custom_sessions_dirs, agent_sessions_enabled: raw.agent_sessions_enabled, + stay_awake_enabled: raw.stay_awake_enabled, agent_session_ssh_hosts: raw.agent_session_ssh_hosts, hooks_enabled: raw.hooks_enabled, http_proxy_enabled: raw.http_proxy_enabled, From 7cddcacdfce5e4fa084639a00b9722488fa58b98 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:41:31 +0700 Subject: [PATCH 127/407] Port upstream 0.67.0: credential expiry alerts --- .../src-tauri/src/commands/bridge.rs | 3 + .../src/commands/credential_alerts.rs | 124 +++++ .../src-tauri/src/commands/mod.rs | 1 + .../src-tauri/src/commands/providers.rs | 23 +- .../src-tauri/src/commands/settings.rs | 4 + apps/desktop-tauri/src/App.test.tsx | 1 + .../src/floatbar/FloatBar.test.tsx | 1 + apps/desktop-tauri/src/i18n/keys.ts | 4 + .../src/surfaces/PopOutPanel.test.tsx | 1 + .../src/surfaces/TrayPanel.test.tsx | 1 + .../surfaces/settings/tabs/AboutTab.test.tsx | 1 + .../settings/tabs/AdvancedTab.test.tsx | 1 + .../settings/tabs/GeneralTab.test.tsx | 23 + .../src/surfaces/settings/tabs/GeneralTab.tsx | 12 + apps/desktop-tauri/src/types/bridge.test.ts | 1 + apps/desktop-tauri/src/types/bridge.ts | 2 + rust/src/core/provider_state.rs | 19 + rust/src/locale.rs | 4 + rust/src/locale/en-US.ftl | 4 + rust/src/notifications.rs | 27 +- rust/src/notifications/credential.rs | 432 ++++++++++++++++++ rust/src/settings.rs | 5 + rust/src/settings/raw.rs | 3 + rust/src/settings/tests.rs | 2 + 24 files changed, 693 insertions(+), 6 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/credential_alerts.rs create mode 100644 rust/src/notifications/credential.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..9c56a92eba 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -659,6 +659,7 @@ pub struct SettingsSnapshot { provider_usage_thresholds: std::collections::HashMap, predictive_pace_warning_enabled: bool, + credential_expiry_notifications_enabled: bool, show_pace: bool, tray_icon_mode: &'static str, switcher_shows_icons: bool, @@ -781,6 +782,8 @@ impl From for SettingsSnapshot { critical_usage_threshold: settings.critical_usage_threshold, provider_usage_thresholds: settings.provider_usage_thresholds, predictive_pace_warning_enabled: settings.predictive_pace_warning_enabled, + credential_expiry_notifications_enabled: settings + .credential_expiry_notifications_enabled, show_pace: settings.show_pace, tray_icon_mode: tray_icon_mode_label(settings.tray_icon_mode), switcher_shows_icons: settings.switcher_shows_icons, diff --git a/apps/desktop-tauri/src-tauri/src/commands/credential_alerts.rs b/apps/desktop-tauri/src-tauri/src/commands/credential_alerts.rs new file mode 100644 index 0000000000..5675b96ef7 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/credential_alerts.rs @@ -0,0 +1,124 @@ +//! Bridges provider refresh outcomes to the core credential-expiry episode +//! tracker (`codexbar::notifications`). Classification stays typed: the +//! failure kind is captured from the fresh fetch before any last-good +//! fallback replaces the published snapshot, so cached data can neither end +//! nor restart an episode. + +use codexbar::core::{ProviderId, ProviderStateKind}; +use codexbar::notifications::{CredentialAlertPolicy, NotificationManager}; + +use super::bridge::ProviderUsageSnapshot; +use super::providers::quota_notification_account_identity; + +/// Result of the fresh fetch, before last-good preservation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum FetchAttempt { + Succeeded, + Failed(ProviderStateKind), +} + +impl FetchAttempt { + pub(super) fn of(fresh: &ProviderUsageSnapshot) -> Self { + if fresh.error.is_none() { + Self::Succeeded + } else { + Self::Failed(fresh.error_state) + } + } +} + +/// Report one accepted refresh outcome. `published` is the snapshot that +/// replaced the cache; `last_good` is the previous error-free snapshot, the +/// only source of account identity when the fetch itself failed. +pub(super) fn observe_attempt( + manager: &mut NotificationManager, + policy: CredentialAlertPolicy, + id: ProviderId, + token_account_id: Option, + attempt: FetchAttempt, + published: &ProviderUsageSnapshot, + last_good: Option<&ProviderUsageSnapshot>, +) { + match attempt { + FetchAttempt::Succeeded => { + let account = quota_notification_account_identity(published, token_account_id); + manager.observe_credential_recovery(id, &account); + } + FetchAttempt::Failed(kind) => { + let account = quota_notification_account_identity( + last_good.unwrap_or(published), + token_account_id, + ); + manager.observe_credential_failure(policy, id, &account, kind); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use codexbar::core::{ProviderError, instantiate_provider}; + use codexbar::settings::Settings; + + fn failed(kind: ProviderStateKind) -> ProviderUsageSnapshot { + ProviderUsageSnapshot::from_error( + ProviderId::Codex, + instantiate_provider(ProviderId::Codex).metadata(), + "Authentication required".to_string(), + kind, + ) + } + + fn succeeded() -> ProviderUsageSnapshot { + let mut snapshot = failed(ProviderStateKind::Ready); + snapshot.error = None; + snapshot.account_email = Some("person@example.test".to_string()); + snapshot + } + + #[test] + fn attempt_reads_the_fresh_error_state() { + assert_eq!( + FetchAttempt::of(&failed(ProviderStateKind::ExpiredSession)), + FetchAttempt::Failed(ProviderStateKind::ExpiredSession) + ); + assert_eq!(FetchAttempt::of(&succeeded()), FetchAttempt::Succeeded); + } + + #[test] + fn typed_state_matches_provider_error_classification() { + let kind = ProviderError::AuthRequired.state_kind(); + assert_eq!( + FetchAttempt::of(&failed(kind)), + FetchAttempt::Failed(ProviderStateKind::NeedsAuthentication) + ); + } + + #[test] + fn observing_outcomes_with_alerts_off_never_delivers() { + // Default settings leave credential alerts off, so no toast can fire. + let policy = CredentialAlertPolicy::from_settings(&Settings::default()); + let mut manager = NotificationManager::new(); + let good = succeeded(); + let failure = failed(ProviderStateKind::NeedsAuthentication); + + observe_attempt( + &mut manager, + policy, + ProviderId::Codex, + None, + FetchAttempt::of(&failure), + &failure, + Some(&good), + ); + observe_attempt( + &mut manager, + policy, + ProviderId::Codex, + None, + FetchAttempt::Succeeded, + &good, + 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..dec5b8ec24 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -36,6 +36,7 @@ mod claude_accounts; pub(crate) mod claude_reconciliation; mod codex_accounts; mod codex_workspaces; +mod credential_alerts; mod credential_detection; mod credentials; mod diagnostics; diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 103d47b869..e377081cd1 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -1,3 +1,4 @@ +use super::credential_alerts::{self, FetchAttempt}; use super::provider_refresh::{ ProviderRefreshCompletion, ProviderRefreshReservation, complete_provider_refresh, reserve_provider_refresh, @@ -6,6 +7,7 @@ use super::warning_identity::WarningIdentity; use super::*; use chrono::{Local, Utc}; use codexbar::core::HookUsageWindow; +use codexbar::notifications::CredentialAlertPolicy; use serde::Serialize; use std::sync::Arc; @@ -519,6 +521,7 @@ fn spawn_provider_refreshes( .and_then(ProviderAccountData::active_account) .map(|account| account.id); let hooks_enabled = inputs.settings.hooks_enabled; + let credential_alerts = CredentialAlertPolicy::from_settings(&inputs.settings); handles.push(tokio::spawn(async move { let Ok(_permit) = fetch_permits.acquire_owned().await else { @@ -531,6 +534,7 @@ fn spawn_provider_refreshes( generation, token_account_id, hooks_enabled, + credential_alerts, ) .await; })); @@ -564,10 +568,13 @@ async fn refresh_provider( generation: u64, token_account_id: Option, hooks_enabled: bool, + credential_alerts: CredentialAlertPolicy, ) { let (snapshot, account_identity, failure_policy) = fetch_provider_snapshot(id, ctx, token_account_id).await; let fresh_snapshot = snapshot.error.is_none(); + // Captured before last-good preservation can swap in a cached snapshot. + let fetch_attempt = FetchAttempt::of(&snapshot); let state = app.state::>(); let published = if let Ok(mut guard) = state.lock() { @@ -601,6 +608,15 @@ async fn refresh_provider( .provider_cache_updated_at_by_provider .insert(id, std::time::Instant::now()); } + credential_alerts::observe_attempt( + &mut guard.notification_manager, + credential_alerts, + id, + token_account_id, + fetch_attempt, + &snapshot, + cached.as_ref(), + ); Some(snapshot) } } else { @@ -950,6 +966,11 @@ fn update_tray_and_notifications( crate::tray_bridge::update_tray_status_items(app, &cached); crate::tray_bridge::update_tray_icon_and_tooltip(app, &cached); notify_usage_thresholds(state, settings, token_accounts, &cached); + if let Ok(mut guard) = state.lock() { + guard + .notification_manager + .retire_credential_episodes_except(&settings.get_enabled_provider_ids()); + } Ok(()) } @@ -1053,7 +1074,7 @@ fn dispatch_quota_hooks( /// Stable account discriminator for threshold/session toast dedupe. /// Prefer token-account id, then email, org, plan; empty for single-account lanes. -fn quota_notification_account_identity( +pub(super) fn quota_notification_account_identity( snapshot: &ProviderUsageSnapshot, token_account_id: Option, ) -> String { diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..ca18cfcfa6 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -24,6 +24,7 @@ pub struct SettingsUpdate { pub provider_usage_thresholds: Option>, pub predictive_pace_warning_enabled: Option, + pub credential_expiry_notifications_enabled: Option, pub show_pace: Option, pub tray_icon_mode: Option, pub switcher_shows_icons: Option, @@ -298,6 +299,9 @@ impl SettingsUpdate { if let Some(v) = self.predictive_pace_warning_enabled { settings.predictive_pace_warning_enabled = v; } + if let Some(v) = self.credential_expiry_notifications_enabled { + settings.credential_expiry_notifications_enabled = v; + } if let Some(v) = self.show_pace { settings.show_pace = v; } diff --git a/apps/desktop-tauri/src/App.test.tsx b/apps/desktop-tauri/src/App.test.tsx index 42d55a7c27..4059e9a056 100644 --- a/apps/desktop-tauri/src/App.test.tsx +++ b/apps/desktop-tauri/src/App.test.tsx @@ -84,6 +84,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index 88860e087d..d5a2608710 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -149,6 +149,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..360ae9da0a 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -58,6 +58,10 @@ export const ALL_LOCALE_KEYS = [ "PredictivePaceWarningsHelper", "PredictivePaceWarningTitle", "PredictivePaceWarningBody", + "CredentialExpiryNotifications", + "CredentialExpiryNotificationsHelper", + "CredentialExpiryTitle", + "CredentialExpiryBody", "UsageDisplay", "ShowUsageAsUsed", "ShowUsageAsUsedHelper", diff --git a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx index 0e7f3053a9..b24605805e 100644 --- a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx @@ -154,6 +154,7 @@ function settings(): SettingsSnapshot { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 8e156a34bf..a5985a3820 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -141,6 +141,7 @@ function settings(overrides: Partial = {}): SettingsSnapshot { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: false, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx index 1a818ed97f..29e1bf7f5b 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx @@ -60,6 +60,7 @@ const settings: SettingsSnapshot = { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx index 8d1d7f2148..7ba050d564 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx @@ -29,6 +29,7 @@ const settings: SettingsSnapshot = { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx index 33a78a00e9..0b849644df 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx @@ -51,6 +51,7 @@ const settings: SettingsSnapshot = { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, @@ -157,6 +158,28 @@ describe("GeneralTab language picker", () => { expect(set).toHaveBeenCalledWith({ predictivePaceWarningEnabled: true }); }); + it("updates the credential expiry alert preference", () => { + const set = vi.fn(); + render( + , + ); + + const toggle = screen.getByRole("checkbox", { + name: "CredentialExpiryNotifications", + }); + expect(toggle).not.toBeChecked(); + fireEvent.click(toggle); + + expect(set).toHaveBeenCalledWith({ + credentialExpiryNotificationsEnabled: true, + }); + }); + it("updates the low power mode preference", () => { const set = vi.fn(); render(); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx index f2e700a5e3..99d40e5ab8 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx @@ -338,6 +338,18 @@ export default function GeneralTab({ onChange={(v) => set({ predictivePaceWarningEnabled: v })} /> + + set({ credentialExpiryNotificationsEnabled: v })} + /> + { highUsageThreshold: 70, criticalUsageThreshold: 90, predictivePaceWarningEnabled: false, + credentialExpiryNotificationsEnabled: false, trayIconMode: "single", switcherShowsIcons: true, menuBarShowsHighestUsage: true, diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 88754baa7d..d75458c4db 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -188,6 +188,7 @@ export interface SettingsSnapshot { criticalUsageThreshold: number; providerUsageThresholds?: Record; predictivePaceWarningEnabled: boolean; + credentialExpiryNotificationsEnabled: boolean; showPace?: boolean; trayIconMode: TrayIconMode; switcherShowsIcons: boolean; @@ -297,6 +298,7 @@ export interface SettingsUpdate { criticalUsageThreshold?: number; providerUsageThresholds?: Record; predictivePaceWarningEnabled?: boolean; + credentialExpiryNotificationsEnabled?: boolean; showPace?: boolean; trayIconMode?: TrayIconMode; switcherShowsIcons?: boolean; diff --git a/rust/src/core/provider_state.rs b/rust/src/core/provider_state.rs index a4087e1ca5..ec2eff1ab2 100644 --- a/rust/src/core/provider_state.rs +++ b/rust/src/core/provider_state.rs @@ -44,6 +44,16 @@ impl ProviderStateKind { pub const fn is_problem(self) -> bool { !matches!(self, ProviderStateKind::Ready) } + + /// Whether this state means the account must sign in again. The only + /// states that qualify for credential-expiry alerts; quota, permission, + /// rate-limit and transport failures map elsewhere and never do. + pub const fn needs_sign_in(self) -> bool { + matches!( + self, + ProviderStateKind::NeedsAuthentication | ProviderStateKind::ExpiredSession + ) + } } impl ProviderError { @@ -137,6 +147,15 @@ mod tests { } } + #[test] + fn only_authentication_states_need_sign_in() { + assert!(ProviderStateKind::NeedsAuthentication.needs_sign_in()); + assert!(ProviderStateKind::ExpiredSession.needs_sign_in()); + assert!(!ProviderStateKind::Ready.needs_sign_in()); + assert!(!ProviderStateKind::LocalRuntimeOffline.needs_sign_in()); + assert!(!ProviderStateKind::Unknown.needs_sign_in()); + } + #[test] fn serializes_as_camel_case_for_the_bridge() { assert_eq!( diff --git a/rust/src/locale.rs b/rust/src/locale.rs index e6ec918ac4..d7cd2868b8 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -244,6 +244,10 @@ locale_keys! { PredictivePaceWarningsHelper, PredictivePaceWarningTitle, PredictivePaceWarningBody, + CredentialExpiryNotifications, + CredentialExpiryNotificationsHelper, + CredentialExpiryTitle, + CredentialExpiryBody, // Display settings (Preferences) UsageDisplay, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..034974a2e3 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -51,6 +51,10 @@ PredictivePaceWarnings = Predictive Pace Warnings PredictivePaceWarningsHelper = Alert when Codex or Claude usage is likely to run out before reset PredictivePaceWarningTitle = { "{}" } { "{}" } pace warning PredictivePaceWarningBody = Quota may run out in { "{}" } +CredentialExpiryNotifications = Credential Expiry Alerts +CredentialExpiryNotificationsHelper = Alert once when a provider account needs to sign in again +CredentialExpiryTitle = { "{}" } needs sign-in +CredentialExpiryBody = Open CodexBar to review the account error and sign in again. UsageDisplay = Usage Display ShowUsageAsUsed = Show Usage as Used ShowUsageAsUsedHelper = Display as used percentage instead of remaining diff --git a/rust/src/notifications.rs b/rust/src/notifications.rs index 94a109bf40..bfde706afb 100755 --- a/rust/src/notifications.rs +++ b/rust/src/notifications.rs @@ -14,6 +14,11 @@ use crate::settings::Settings; use crate::sound::{NotificationSoundEvent, play_alert}; use chrono::{DateTime, Utc}; +mod credential; + +pub use credential::CredentialAlertPolicy; +use credential::CredentialEpisodes; + /// Notification types #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum NotificationType { @@ -132,6 +137,8 @@ pub struct NotificationManager { previous_session_percent: std::collections::HashMap, predictive_warning_keys: std::collections::HashSet, deepseek_pricing_period: Option, + /// Open credential-expiry episodes (in memory; reset on restart). + credential_episodes: CredentialEpisodes, } impl NotificationManager { @@ -141,6 +148,7 @@ impl NotificationManager { previous_session_percent: std::collections::HashMap::new(), predictive_warning_keys: std::collections::HashSet::new(), deepseek_pricing_period: None, + credential_episodes: CredentialEpisodes::default(), } } @@ -514,8 +522,10 @@ impl NotificationManager { Self::play_notification_sound(NotificationSoundEvent::StatusIssue, settings); } + /// Returns whether the toast was handed to the OS (a spawn failure is the + /// only failure the fire-and-forget PowerShell dispatch can observe). #[cfg(target_os = "windows")] - fn show_toast(&self, title: &str, body: &str) { + fn show_toast(&self, title: &str, body: &str) -> bool { use std::os::windows::process::CommandExt; use std::process::Command; use std::sync::Once; @@ -573,13 +583,19 @@ impl NotificationManager { .creation_flags(0x08000000) // CREATE_NO_WINDOW .spawn() { - Ok(_) => tracing::debug!("Toast notification dispatched: {}", title), - Err(e) => tracing::warn!("Failed to dispatch toast notification '{}': {}", title, e), + Ok(_) => { + tracing::debug!("Toast notification dispatched: {}", title); + true + } + Err(e) => { + tracing::warn!("Failed to dispatch toast notification '{}': {}", title, e); + false + } } } #[cfg(not(target_os = "windows"))] - fn show_toast(&self, title: &str, body: &str) { + fn show_toast(&self, title: &str, body: &str) -> bool { use std::process::Command; // Try notify-send first (works on most Linux distros including WSL with WSLg) @@ -594,10 +610,11 @@ impl NotificationManager { && output.status.success() { tracing::debug!("Sent notification via notify-send: {}", title); - return; + return true; } tracing::info!("Notification: {} - {}", title, body); + false } } diff --git a/rust/src/notifications/credential.rs b/rust/src/notifications/credential.rs new file mode 100644 index 0000000000..d0566f8069 --- /dev/null +++ b/rust/src/notifications/credential.rs @@ -0,0 +1,432 @@ +//! Opt-in credential-expiry alerts (upstream 0.67.0 `handleCredentialOutcome`). +//! +//! One toast per (provider, account scope) failure episode. Only a typed +//! sign-in state ([`ProviderStateKind::needs_sign_in`]) starts an episode; +//! only a fresh successful fetch ends it. Cached or last-good fallbacks, +//! repeated failures and unrelated failures (network, timeout, parse, quota) +//! leave it untouched, so a flapping provider never re-alerts. Episodes live +//! in memory and start fresh when the app restarts. +//! +//! The toast names the provider only; it never carries the raw error, email +//! or account id. + +use std::collections::HashSet; + +use super::NotificationManager; +use crate::core::{ProviderId, ProviderStateKind}; +use crate::locale::{self, LocaleKey}; +use crate::settings::{Language, Settings}; + +/// Delivery settings captured at refresh time, so a refresh task does not +/// need the whole [`Settings`] value. +#[derive(Debug, Clone, Copy)] +pub struct CredentialAlertPolicy { + enabled: bool, + language: Language, +} + +impl CredentialAlertPolicy { + /// The master notification switch also silences credential alerts, like + /// every other toast. + pub fn from_settings(settings: &Settings) -> Self { + Self { + enabled: settings.show_notifications + && settings.credential_expiry_notifications_enabled, + language: settings.ui_language, + } + } +} + +/// Open episodes keyed by provider and account scope. An empty scope means +/// no account evidence was available, so it acts as a wildcard for that +/// provider: an identity gap must neither hide a real recovery nor re-alert +/// on every refresh. +#[derive(Debug, Default)] +pub(super) struct CredentialEpisodes { + open: HashSet<(ProviderId, String)>, +} + +impl CredentialEpisodes { + fn covers(&self, provider: ProviderId, account: &str) -> bool { + if account.is_empty() { + return self.open.iter().any(|(open, _)| *open == provider); + } + self.open.contains(&(provider, account.to_string())) + || self.open.contains(&(provider, String::new())) + } + + /// Reserve an episode; `false` when one already covers this account. + fn begin(&mut self, provider: ProviderId, account: &str) -> bool { + if self.covers(provider, account) { + return false; + } + self.open.insert((provider, account.to_string())); + true + } + + fn release(&mut self, provider: ProviderId, account: &str) { + self.open.remove(&(provider, account.to_string())); + } + + fn end(&mut self, provider: ProviderId, account: &str) { + if account.is_empty() { + self.open.retain(|(open, _)| *open != provider); + } else { + self.open.remove(&(provider, account.to_string())); + self.open.remove(&(provider, String::new())); + } + } +} + +impl NotificationManager { + /// Record a failed refresh for `account`. Posts one toast when `kind` is a + /// sign-in state, alerts are enabled and no episode is open yet. Returns + /// whether a toast was delivered. + pub fn observe_credential_failure( + &mut self, + policy: CredentialAlertPolicy, + provider: ProviderId, + account: &str, + kind: ProviderStateKind, + ) -> bool { + self.observe_credential_failure_with(policy, provider, account, kind, Self::show_toast) + } + + fn observe_credential_failure_with( + &mut self, + policy: CredentialAlertPolicy, + provider: ProviderId, + account: &str, + kind: ProviderStateKind, + deliver: impl FnOnce(&Self, &str, &str) -> bool, + ) -> bool { + if !policy.enabled || !kind.needs_sign_in() { + return false; + } + if !self.credential_episodes.begin(provider, account) { + return false; + } + let title = locale::format_locale( + policy.language, + LocaleKey::CredentialExpiryTitle, + &[provider.display_name()], + ); + let body = locale::get_text(policy.language, LocaleKey::CredentialExpiryBody); + let delivered = deliver(self, &title, &body); + if !delivered { + // A failed toast must not swallow the alert: the next refresh retries. + self.credential_episodes.release(provider, account); + } + delivered + } + + /// Record a fresh successful fetch (no error, not a last-good replay). + /// Ends the episode even while alerts are disabled, so re-enabling them + /// never replays a stale failure. + pub fn observe_credential_recovery(&mut self, provider: ProviderId, account: &str) { + self.credential_episodes.end(provider, account); + } + + /// Forget episodes of providers that are no longer enabled. + pub fn retire_credential_episodes_except(&mut self, enabled: &[ProviderId]) { + self.credential_episodes + .open + .retain(|(provider, _)| enabled.contains(provider)); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::ProviderError; + use std::cell::RefCell; + + const ACCOUNT: &str = "person@example.test"; + + fn policy(enabled: bool) -> CredentialAlertPolicy { + CredentialAlertPolicy { + enabled, + language: Language::English, + } + } + + /// Records delivered toasts; `succeed` scripts the delivery result. + struct Toasts { + sent: RefCell>, + } + + impl Toasts { + fn new() -> Self { + Self { + sent: RefCell::new(Vec::new()), + } + } + + fn fail( + &self, + manager: &mut NotificationManager, + provider: ProviderId, + account: &str, + kind: ProviderStateKind, + succeed: bool, + ) -> bool { + manager.observe_credential_failure_with( + policy(true), + provider, + account, + kind, + |_, title, body| { + self.sent + .borrow_mut() + .push((title.to_string(), body.to_string())); + succeed + }, + ) + } + + fn count(&self) -> usize { + self.sent.borrow().len() + } + } + + fn network_error() -> ProviderError { + ProviderError::Network( + reqwest::Client::new() + .get("file:///private.example.test/unreachable") + .build() + .unwrap_err(), + ) + } + + #[test] + fn each_provider_error_variant_alerts_only_for_sign_in_states() { + use ProviderError as E; + let cases = [ + (E::NotInstalled("key missing".into()), true), + (E::AuthRequired, true), + (E::NoCookies, true), + (E::OAuth("credentials not found".into()), true), + (E::OAuthRevoked("revoked".into()), true), + (E::OAuthExpired("expired".into()), true), + (E::OAuthTransient("try again".into()), false), + (E::Parse("bad body".into()), false), + (network_error(), false), + (E::Timeout, false), + (E::UnsupportedSource(crate::core::SourceMode::Auto), false), + (E::Other("API error 500".into()), false), + (E::Other("quota exhausted".into()), false), + (E::Other("403 permission denied".into()), false), + (E::Other("429 rate limited".into()), false), + ]; + for (error, expect_alert) in cases { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let alerted = toasts.fail( + &mut manager, + ProviderId::Claude, + ACCOUNT, + error.state_kind(), + true, + ); + assert_eq!(alerted, expect_alert, "{error}"); + assert_eq!(toasts.count(), usize::from(expect_alert), "{error}"); + } + } + + #[test] + fn provider_specific_offline_state_does_not_alert() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + assert!(!toasts.fail( + &mut manager, + ProviderId::Antigravity, + "", + ProviderStateKind::LocalRuntimeOffline, + true, + )); + assert_eq!(toasts.count(), 0); + } + + #[test] + fn toast_names_the_provider_and_nothing_else() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + assert!(toasts.fail( + &mut manager, + ProviderId::Claude, + ACCOUNT, + ProviderStateKind::ExpiredSession, + true, + )); + let sent = toasts.sent.borrow(); + assert_eq!(sent[0].0, "Claude needs sign-in"); + assert_eq!( + sent[0].1, + "Open CodexBar to review the account error and sign in again." + ); + assert!(!sent[0].0.contains(ACCOUNT) && !sent[0].1.contains(ACCOUNT)); + } + + #[test] + fn episode_ends_only_on_fresh_success() { + // fail, fail, timeout, cached success (no recovery call), fresh + // success, fail => two toasts in total. + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + let p = ProviderId::Codex; + + assert!(toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + assert!(!toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + assert!(!toasts.fail(&mut manager, p, ACCOUNT, ProviderStateKind::Unknown, true)); + // A cached / last-good replay is never reported as a recovery, so the + // next sign-in failure is still the same episode. + assert!(!toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + + manager.observe_credential_recovery(p, ACCOUNT); + assert!(toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + assert_eq!(toasts.count(), 2); + } + + #[test] + fn failed_toast_releases_the_reservation_for_retry() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + + assert!(!toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, false)); + assert!(toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + assert!(!toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + assert_eq!(toasts.count(), 2); + } + + #[test] + fn accounts_and_providers_have_independent_episodes() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + + assert!(toasts.fail( + &mut manager, + ProviderId::Codex, + "a@example.test", + auth, + true + )); + assert!(toasts.fail( + &mut manager, + ProviderId::Codex, + "b@example.test", + auth, + true + )); + assert!(toasts.fail( + &mut manager, + ProviderId::Claude, + "a@example.test", + auth, + true + )); + + manager.observe_credential_recovery(ProviderId::Codex, "a@example.test"); + assert!(toasts.fail( + &mut manager, + ProviderId::Codex, + "a@example.test", + auth, + true + )); + assert!(!toasts.fail( + &mut manager, + ProviderId::Codex, + "b@example.test", + auth, + true + )); + } + + #[test] + fn identity_gap_neither_realerts_nor_hides_recovery() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + let p = ProviderId::Codex; + + // No identity yet, then the same account surfaces without evidence. + assert!(toasts.fail(&mut manager, p, "", auth, true)); + assert!(!toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + // Failing with a known identity is covered by that episode too. + manager.observe_credential_recovery(p, ACCOUNT); + assert!(toasts.fail(&mut manager, p, "", auth, true)); + // A success with no identity ends every episode of the provider. + manager.observe_credential_recovery(p, ""); + assert!(toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + assert_eq!(toasts.count(), 3); + } + + #[test] + fn disabled_alerts_deliver_nothing_but_keep_open_episodes() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + let p = ProviderId::Codex; + + assert!(toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + + let disabled = manager.observe_credential_failure_with( + policy(false), + p, + "other@example.test", + auth, + |_, _, _| panic!("no delivery while disabled"), + ); + assert!(!disabled); + + // Off then on: the unresolved episode is remembered, so no replay... + assert!(!toasts.fail(&mut manager, p, ACCOUNT, auth, true)); + // ...and the account that failed while disabled alerts once enabled. + assert!(toasts.fail(&mut manager, p, "other@example.test", auth, true)); + assert_eq!(toasts.count(), 2); + } + + #[test] + fn recovery_while_disabled_still_ends_the_episode() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::ExpiredSession; + + assert!(toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + manager.observe_credential_recovery(ProviderId::Codex, ACCOUNT); + assert!(toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + } + + #[test] + fn disabled_providers_retire_their_episodes() { + let mut manager = NotificationManager::new(); + let toasts = Toasts::new(); + let auth = ProviderStateKind::NeedsAuthentication; + + assert!(toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + assert!(toasts.fail(&mut manager, ProviderId::Claude, ACCOUNT, auth, true)); + + manager.retire_credential_episodes_except(&[ProviderId::Claude]); + + assert!(toasts.fail(&mut manager, ProviderId::Codex, ACCOUNT, auth, true)); + assert!(!toasts.fail(&mut manager, ProviderId::Claude, ACCOUNT, auth, true)); + } + + #[test] + fn policy_requires_master_switch_and_credential_toggle() { + let mut settings = Settings { + show_notifications: true, + credential_expiry_notifications_enabled: true, + ..Settings::default() + }; + assert!(CredentialAlertPolicy::from_settings(&settings).enabled); + settings.show_notifications = false; + assert!(!CredentialAlertPolicy::from_settings(&settings).enabled); + settings.show_notifications = true; + settings.credential_expiry_notifications_enabled = false; + assert!(!CredentialAlertPolicy::from_settings(&settings).enabled); + } +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..da636db21f 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -165,6 +165,10 @@ pub struct Settings { #[serde(default)] pub predictive_pace_warning_enabled: bool, + /// Notify once per failure episode when a provider account needs sign-in again. + #[serde(default)] + pub credential_expiry_notifications_enabled: bool, + /// Show pace visualizations and forecast text in provider menu cards. #[serde(default = "default_true")] pub show_pace: bool, @@ -549,6 +553,7 @@ impl Default for Settings { reset_time_relative: true, // Show relative times by default show_reset_when_exhausted: false, predictive_pace_warning_enabled: false, + credential_expiry_notifications_enabled: false, show_pace: true, menu_bar_display_mode: "detailed".to_string(), // Detailed mode by default overview_layout: default_overview_layout(), diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..8826eb336b 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -44,6 +44,7 @@ pub(super) struct RawSettings { reset_time_relative: bool, show_reset_when_exhausted: bool, predictive_pace_warning_enabled: bool, + credential_expiry_notifications_enabled: bool, #[serde(default = "default_true")] show_pace: bool, menu_bar_display_mode: String, @@ -214,6 +215,7 @@ impl Default for RawSettings { reset_time_relative: s.reset_time_relative, show_reset_when_exhausted: s.show_reset_when_exhausted, predictive_pace_warning_enabled: s.predictive_pace_warning_enabled, + credential_expiry_notifications_enabled: s.credential_expiry_notifications_enabled, show_pace: s.show_pace, menu_bar_display_mode: s.menu_bar_display_mode, overview_layout: s.overview_layout, @@ -541,6 +543,7 @@ impl From for Settings { reset_time_relative: raw.reset_time_relative, show_reset_when_exhausted: raw.show_reset_when_exhausted, predictive_pace_warning_enabled: raw.predictive_pace_warning_enabled, + credential_expiry_notifications_enabled: raw.credential_expiry_notifications_enabled, show_pace: raw.show_pace, menu_bar_display_mode: raw.menu_bar_display_mode, overview_layout: normalize_overview_layout(&raw.overview_layout), diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..d56ea02783 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -19,6 +19,7 @@ fn test_settings_default() { assert_eq!(settings.critical_usage_threshold, 90.0); assert!(!settings.show_reset_when_exhausted); assert!(!settings.predictive_pace_warning_enabled); + assert!(!settings.credential_expiry_notifications_enabled); assert!(!settings.float_bar_show_cost); assert!(!settings.tray_panel_always_on_top); assert_eq!(settings.overview_layout, "compact"); @@ -183,6 +184,7 @@ fn new_warning_and_reset_settings_are_backward_compatible() { assert!(!loaded.show_reset_when_exhausted); assert!(!loaded.predictive_pace_warning_enabled); + assert!(!loaded.credential_expiry_notifications_enabled); } #[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 128/407] 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}

}
); } diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 88754baa7d..edee983d32 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -628,6 +628,8 @@ export interface ProviderDisplayDetail { value: string; secondaryValue: string | null; progress: ProviderDisplayProgress | null; + /** Heading shared by consecutive rows; absent for ungrouped rows. */ + section?: string | null; } /** One metric or provider-emitted extra row available to visibility controls. */ @@ -945,6 +947,9 @@ export interface ProviderDetail { autoResumeAfterQuotaReset: boolean; /** Whether the active credential lane can be correlated to a local CLI session. */ autoResumeSupported: boolean; + /** LiteLLM and Claude expose one opt-in extra breakdown; other providers do not. */ + optionalDetailsSupported: boolean; + optionalDetailsEnabled: boolean; // Identity email: string | null; diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 64d0b54d59..e426dccb4d 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -190,6 +190,7 @@ async fn collect_provider_diagnostic( auto_prefer_web: false, // Diagnostics keep the short optional-join grace (upstream #2583 gate). requires_optional_usage_completeness: false, + optional_details_enabled: settings.optional_details_enabled(provider_id), }; let fetch_result = provider.fetch_usage(&ctx).await; diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index b9ecf5dce8..9ba99834b1 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -325,6 +325,7 @@ async fn fetch_guard_outcome( auto_prefer_web: false, // Guard checks keep the short optional-join grace. requires_optional_usage_completeness: false, + optional_details_enabled: false, }; match provider.fetch_usage(&ctx).await { diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 2f26498cf1..3dfa6377c1 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -303,6 +303,7 @@ async fn hooks_watch_observation( auto_prefer_web: false, // Hook watches keep the short optional-join grace. requires_optional_usage_completeness: false, + optional_details_enabled: settings.optional_details_enabled(provider_id), }; if ctx.api_key.is_none() { diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 8d28758df4..95ed7793b5 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -165,6 +165,7 @@ async fn fetch_provider_envelope( gateway_url: None, auto_prefer_web: false, requires_optional_usage_completeness: false, + optional_details_enabled: false, }; let fetch = bounded_fetch(provider_id, ctx, None, fetch_timeout).await; ProviderFetchEnvelope { @@ -298,6 +299,7 @@ async fn collect_claude_accounts(claude_enabled: bool) -> Option) -> String { // join grace (upstream #2583), unlike `codexbar usage` which blocks // for the full completeness window. requires_optional_usage_completeness: false, + optional_details_enabled: false, }; let mut results = Vec::new(); diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 2b1dfff0af..a18fdb2e2b 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -257,6 +257,9 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch // `codexbar usage` is a foreground read: optional enrichment (e.g. the // OpenCode Go Zen balance) is worth its full bounded wait (#2583). requires_optional_usage_completeness: true, + // Per-provider opt-ins live in settings, which this shared context + // does not load; the optional breakdowns stay off here. + optional_details_enabled: false, } } diff --git a/rust/src/cli/usage/render.rs b/rust/src/cli/usage/render.rs index 62fd50310d..fbe119f7c0 100644 --- a/rust/src/cli/usage/render.rs +++ b/rust/src/cli/usage/render.rs @@ -81,6 +81,7 @@ pub fn render_json_result( "title": detail.title(), "value": detail.value(), "secondaryValue": detail.secondary_value(), + "section": detail.section(), "progress": detail.progress().map(|progress| { serde_json::json!({ "used": progress.used(), @@ -298,7 +299,14 @@ fn append_inventory_lines(lines: &mut Vec, inventory: &[ProviderInventor } fn append_display_detail_lines(lines: &mut Vec, details: &[ProviderDisplayDetail]) { + let mut current_section: Option<&str> = None; for detail in details { + if detail.section() != current_section { + current_section = detail.section(); + if let Some(section) = current_section { + lines.push(format!(" {section}:")); + } + } let secondary = detail .secondary_value() .map(|value| format!(" ({value})")) diff --git a/rust/src/cli/usage_tests.rs b/rust/src/cli/usage_tests.rs index d332131b81..27315e6be3 100644 --- a/rust/src/cli/usage_tests.rs +++ b/rust/src/cli/usage_tests.rs @@ -300,6 +300,30 @@ fn display_details_are_rendered_in_full_text_and_json() { assert_eq!(json["details"][0]["progress"]["total"], 100.0); } +#[test] +fn sectioned_display_details_print_one_heading_per_group() { + let row = |id: &str, title: &str, section: Option<&str>| { + let row = ProviderDisplayDetail::new(id, title, "1"); + match section { + Some(section) => row.and_then(|row| row.with_section(section)), + None => row, + } + }; + let result = fetch_result(UsageSnapshot::new(RateWindow::new(10.0))) + .with_display_detail(row("plain", "Plain", None)) + .with_display_detail(row("a", "Alpha", Some("Model activity"))) + .with_display_detail(row("b", "Beta", Some("Model activity"))); + + let text = render_text_with_status(ProviderId::LiteLLM, &result, None, false); + let json = render_json_result(ProviderId::LiteLLM, result, None); + + assert_eq!(text.matches("Model activity:").count(), 1); + assert!(text.find("Plain: 1").unwrap() < text.find("Model activity:").unwrap()); + assert!(text.find("Model activity:").unwrap() < text.find("Alpha: 1").unwrap()); + assert_eq!(json["details"][0]["section"], serde_json::Value::Null); + assert_eq!(json["details"][2]["section"], "Model activity"); +} + #[test] fn json_inventory_is_additive_and_contains_no_redemption_token() { let result = fetch_result(UsageSnapshot::new(RateWindow::new(10.0))).with_inventory_item( diff --git a/rust/src/core/display_detail.rs b/rust/src/core/display_detail.rs index 1c80831aa3..6008f33162 100644 --- a/rust/src/core/display_detail.rs +++ b/rust/src/core/display_detail.rs @@ -22,6 +22,7 @@ pub struct ProviderDisplayDetail { value: String, secondary_value: Option, progress: Option, + section: Option, } #[derive(Debug, Clone, Copy, PartialEq)] @@ -44,9 +45,10 @@ impl ProviderDisplayDetail { value: value.into(), secondary_value: None, progress: None, + section: None, }; let valid = is_display_shape(&row.id, 64) - && is_display_shape(&row.title, 128) + && Self::is_valid_title(&row.title) && is_display_shape(&row.value, 512); valid.then_some(row) } @@ -61,6 +63,23 @@ impl ProviderDisplayDetail { Some(self) } + /// Group the row under a heading shared by consecutive rows; rejects + /// invalid text. + pub fn with_section(mut self, title: impl Into) -> Option { + let title = title.into(); + if !is_display_shape(&title, 128) { + return None; + } + self.section = Some(title); + Some(self) + } + + /// Whether `title` would be accepted as a row title, for providers that + /// must reject a whole payload when any wire-supplied label is unusable. + pub fn is_valid_title(title: &str) -> bool { + is_display_shape(title, 128) + } + /// Attach a progress bar; rejects non-finite or non-positive ranges. pub fn with_progress(mut self, used: f64, total: f64) -> Option { if !(used.is_finite() && total.is_finite() && used >= 0.0 && total > 0.0) { @@ -89,6 +108,10 @@ impl ProviderDisplayDetail { pub fn progress(&self) -> Option { self.progress } + + pub fn section(&self) -> Option<&str> { + self.section.as_deref() + } } impl ProviderDisplayProgress { diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..28d1de8285 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -777,6 +777,11 @@ pub struct FetchContext { /// short join grace instead (upstream 0.48.0 /// `requiresOptionalUsageCompleteness`, #2583). pub requires_optional_usage_completeness: bool, + + /// The user opted in to the requested provider's optional detail breakdown + /// (LiteLLM model activity, Claude workspace spend). Scoped to the provider + /// being fetched; false everywhere the setting is not consulted. + pub optional_details_enabled: bool, } impl Default for FetchContext { @@ -795,6 +800,7 @@ impl Default for FetchContext { gateway_url: None, auto_prefer_web: false, requires_optional_usage_completeness: false, + optional_details_enabled: false, } } } diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index ebbdcb8c20..e499d9bdba 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -758,6 +758,28 @@ mod tests { assert_eq!(details[0].value(), "12"); } + #[test] + fn display_detail_sections_are_validated() { + let row = ProviderDisplayDetail::new("a", "Alpha", "1").unwrap(); + assert_eq!(row.section(), None); + assert!(row.clone().with_section("").is_none()); + assert!( + row.clone() + .with_section( + "bad +section" + ) + .is_none() + ); + assert!(row.clone().with_section("x".repeat(129)).is_none()); + assert_eq!( + row.with_section("Model activity").unwrap().section(), + Some("Model activity") + ); + assert!(ProviderDisplayDetail::is_valid_title("fixture-alpha")); + assert!(!ProviderDisplayDetail::is_valid_title("")); + } + #[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.rs b/rust/src/locale.rs index e6ec918ac4..7df4a47610 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -1194,6 +1194,11 @@ locale_keys! { ProviderAutoResumeTitle, ProviderAutoResumeAfterQuotaReset, ProviderAutoResumeAfterQuotaResetHelper, + ProviderOptionalDetailsTitle, + ProviderLiteLLMModelActivity, + ProviderLiteLLMModelActivityHelper, + ProviderClaudeWorkspaceSpend, + ProviderClaudeWorkspaceSpendHelper, } #[cfg(test)] diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..11a93a465c 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -839,6 +839,11 @@ ProviderAccentColorInvalid = Invalid hex color. Use #RRGGBB format, e.g. #FF5733 ProviderAutoResumeTitle = Automatic session resume ProviderAutoResumeAfterQuotaReset = Reopen the CLI session when quota resets ProviderAutoResumeAfterQuotaResetHelper = Reopen the exact local Codex or Claude session after its quota becomes available. No prompt is sent automatically. +ProviderOptionalDetailsTitle = Optional details +ProviderLiteLLMModelActivity = Show model activity +ProviderLiteLLMModelActivityHelper = Adds tokens and requests per model for your user over the last 30 UTC days (top 20). Makes up to 3 extra requests to your LiteLLM proxy per refresh; if it is unavailable the section is hidden and budgets are unaffected. Off by default. +ProviderClaudeWorkspaceSpend = Show workspace spend +ProviderClaudeWorkspaceSpendHelper = Adds Admin API spend per workspace over the last 30 days (top 20, shown only when more than one workspace has spend). Organization totals are unchanged and no extra request is made. Off by default. OpenRouterManagementKeyTitle = OpenRouter spend history OpenRouterManagementKeyLabel = Management API key diff --git a/rust/src/providers/claude/admin_api.rs b/rust/src/providers/claude/admin_api.rs index e56f0cf907..24c841aa22 100644 --- a/rust/src/providers/claude/admin_api.rs +++ b/rust/src/providers/claude/admin_api.rs @@ -10,12 +10,15 @@ use serde::Deserialize; use std::collections::HashMap; use crate::core::{ - CostSnapshot, FetchContext, ProviderError, ProviderFetchResult, RateWindow, UsageSnapshot, + CostSnapshot, FetchContext, ProviderDisplayDetail, ProviderError, ProviderFetchResult, + RateWindow, UsageSnapshot, }; const COST_REPORT_URL: &str = "https://api.anthropic.com/v1/organizations/cost_report"; const MESSAGES_USAGE_URL: &str = "https://api.anthropic.com/v1/organizations/usage_report/messages"; const ANTHROPIC_VERSION: &str = "2023-06-01"; +const WORKSPACE_SECTION_TITLE: &str = "Workspace spend \u{b7} 30d"; +const MAX_WORKSPACE_ROWS: usize = 20; pub struct ClaudeAdminApiFetcher { client: Client, @@ -50,13 +53,7 @@ impl ClaudeAdminApiFetcher { let costs: CostReportResponse = self .fetch_json( COST_REPORT_URL, - &[ - ("starting_at", start.to_rfc3339()), - ("ending_at", end.to_rfc3339()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by[]", "description".to_string()), - ], + &cost_report_query(&start, &end, ctx.optional_details_enabled), &api_key, "cost_report", ) @@ -75,7 +72,12 @@ impl ClaudeAdminApiFetcher { "messages", ) .await?; - Ok(result_from_admin_usage(&costs, &messages, now)) + Ok(result_from_admin_usage( + &costs, + &messages, + now, + ctx.optional_details_enabled, + )) } fn api_key(ctx: &FetchContext) -> Option { @@ -126,6 +128,26 @@ impl ClaudeAdminApiFetcher { } } +/// Cost report query. Workspace spend only adds a second `group_by[]`; it never +/// costs another request. +fn cost_report_query( + start: &DateTime, + end: &DateTime, + workspace_spend: bool, +) -> Vec<(&'static str, String)> { + let mut query = vec![ + ("starting_at", start.to_rfc3339()), + ("ending_at", end.to_rfc3339()), + ("bucket_width", "1d".to_string()), + ("limit", "31".to_string()), + ("group_by[]", "description".to_string()), + ]; + if workspace_spend { + query.push(("group_by[]", "workspace_id".to_string())); + } + query +} + fn clean_key(raw: &str) -> Option { let trimmed = raw.trim().trim_matches('"').trim_matches('\'').trim(); (!trimmed.is_empty()).then(|| trimmed.to_string()) @@ -152,6 +174,9 @@ struct CostResult { amount: String, description: Option, cost_type: Option, + /// Present only when the cost report is grouped by `workspace_id`; null + /// for the default workspace. + workspace_id: Option, } #[derive(Debug, Deserialize)] @@ -188,6 +213,7 @@ fn result_from_admin_usage( costs: &CostReportResponse, messages: &MessagesUsageResponse, now: DateTime, + workspace_spend: bool, ) -> ProviderFetchResult { let cost_total: f64 = costs .data @@ -301,11 +327,51 @@ fn result_from_admin_usage( ); } - ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( + let mut result = ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( cost_total, "USD", "Last 30 days", - )) + )); + if workspace_spend { + for row in workspace_spend_rows(costs) { + result = result.with_display_detail(Some(row)); + } + } + result +} + +/// Per-workspace spend over the same buckets as the organization total. Shown +/// only when the report spans more than one workspace; a single workspace adds +/// nothing over the organization view. +fn workspace_spend_rows(costs: &CostReportResponse) -> Vec { + let mut workspaces: HashMap = HashMap::new(); + for result in costs.data.iter().flat_map(|bucket| &bucket.results) { + let name = result + .workspace_id + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or("Default"); + *workspaces.entry(name.to_string()).or_default() += usd_from_lowest_unit(&result.amount); + } + if workspaces.len() < 2 { + return Vec::new(); + } + let mut ranked: Vec<_> = workspaces.into_iter().collect(); + ranked.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + ranked + .into_iter() + .take(MAX_WORKSPACE_ROWS) + .enumerate() + .filter_map(|(index, (name, spend))| { + ProviderDisplayDetail::new( + format!("claude-workspace-{index}"), + name, + format!("${spend:.2}"), + )? + .with_section(WORKSPACE_SECTION_TITLE) + }) + .collect() } fn usd_from_lowest_unit(raw: &str) -> f64 { @@ -313,14 +379,5 @@ fn usd_from_lowest_unit(raw: &str) -> f64 { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn cleans_quoted_admin_key() { - assert_eq!( - clean_key(" 'sk-ant-admin-123' "), - Some("sk-ant-admin-123".to_string()) - ); - } -} +#[path = "admin_api_tests.rs"] +mod tests; diff --git a/rust/src/providers/claude/admin_api_tests.rs b/rust/src/providers/claude/admin_api_tests.rs new file mode 100644 index 0000000000..00c47b7c21 --- /dev/null +++ b/rust/src/providers/claude/admin_api_tests.rs @@ -0,0 +1,154 @@ +use super::*; + +fn now() -> DateTime { + DateTime::parse_from_rfc3339("2026-09-24T12:00:00Z") + .unwrap() + .with_timezone(&Utc) +} + +fn no_messages() -> MessagesUsageResponse { + serde_json::from_str(r#"{"data":[]}"#).unwrap() +} + +fn costs(json: &str) -> CostReportResponse { + serde_json::from_str(json).unwrap() +} + +fn workspace_rows(result: &ProviderFetchResult) -> Vec<(&str, &str)> { + result + .display_details() + .iter() + .map(|row| (row.title(), row.value())) + .collect() +} + +/// Wire shape from upstream 0.67.0 `ClaudeAdminAPIUsageTests`. +const WORKSPACE_FIXTURE: &str = r#"{"data":[ + {"starting_at":"2026-09-23T00:00:00Z","ending_at":"2026-09-24T00:00:00Z","results":[ + {"currency":"USD","amount":"1250","description":"Input tokens","workspace_id":"wrk_fixture"}, + {"currency":"USD","amount":"250","description":"Output tokens","workspace_id":"wrk_fixture"}, + {"currency":"USD","amount":"50","description":"Input tokens","workspace_id":null}]}, + {"starting_at":"2026-09-24T00:00:00Z","ending_at":"2026-09-25T00:00:00Z","results":[ + {"currency":"USD","amount":"75","description":"Input tokens","workspace_id":null}]}]}"#; + +#[test] +fn cleans_quoted_admin_key() { + assert_eq!( + clean_key(" 'sk-ant-admin-123' "), + Some("sk-ant-admin-123".to_string()) + ); +} + +#[test] +fn workspace_details_preserve_organization_totals_and_default_workspace_cents() { + let report = costs(WORKSPACE_FIXTURE); + for enabled in [false, true] { + let result = result_from_admin_usage(&report, &no_messages(), now(), enabled); + let cost = result.cost.as_ref().expect("organization cost"); + assert_eq!(cost.used, 16.25); + assert!( + result + .usage + .extra_rate_windows + .iter() + .any(|window| window.id == "cost-0" + && window.window.reset_description.as_deref() == Some("$13.75")) + ); + if enabled { + assert_eq!( + workspace_rows(&result), + [("wrk_fixture", "$15.00"), ("Default", "$1.25")] + ); + assert!( + result + .display_details() + .iter() + .all(|row| row.section() == Some("Workspace spend \u{b7} 30d")) + ); + } else { + assert!(result.display_details().is_empty()); + } + } +} + +#[test] +fn blank_workspace_ids_group_with_the_default_workspace() { + let report = costs( + r#"{"data":[{"starting_at":"2026-09-24T00:00:00Z","ending_at":"2026-09-25T00:00:00Z","results":[ + {"amount":"100","workspace_id":" "}, + {"amount":"200"}, + {"amount":"400","workspace_id":" wrk_a "}]}]}"#, + ); + let result = result_from_admin_usage(&report, &no_messages(), now(), true); + assert_eq!( + workspace_rows(&result), + [("wrk_a", "$4.00"), ("Default", "$3.00")] + ); +} + +#[test] +fn equal_workspace_spend_is_ordered_by_name() { + let report = costs( + r#"{"data":[{"starting_at":"2026-09-24T00:00:00Z","ending_at":"2026-09-25T00:00:00Z","results":[ + {"amount":"100","workspace_id":"wrk_b"}, + {"amount":"100","workspace_id":"wrk_a"}]}]}"#, + ); + let result = result_from_admin_usage(&report, &no_messages(), now(), true); + assert_eq!( + workspace_rows(&result), + [("wrk_a", "$1.00"), ("wrk_b", "$1.00")] + ); +} + +#[test] +fn workspace_rows_are_bounded_and_a_single_workspace_keeps_the_organization_view() { + for count in [1_u32, 25] { + let rows = (0..count) + .map(|index| { + format!( + r#"{{"amount":"{}","workspace_id":"wrk_fixture_{index}"}}"#, + index * 100 + ) + }) + .collect::>() + .join(","); + let report = costs(&format!( + r#"{{"data":[{{"starting_at":"2026-09-24T00:00:00Z","ending_at":"2026-09-25T00:00:00Z","results":[{rows}]}}]}}"# + )); + let result = result_from_admin_usage(&report, &no_messages(), now(), true); + assert_eq!( + result.cost.as_ref().map(|cost| cost.used), + Some(f64::from(count * (count - 1) / 2)) + ); + let details = result.display_details(); + if count == 1 { + assert!(details.is_empty()); + } else { + assert_eq!(details.len(), 20); + assert_eq!( + details.first().map(|row| row.title()), + Some("wrk_fixture_24") + ); + assert_eq!(details.last().map(|row| row.title()), Some("wrk_fixture_5")); + } + } +} + +#[test] +fn workspace_option_groups_the_existing_cost_request_only() { + let start = now(); + let end = now() + Duration::days(1); + let groups = |enabled| { + cost_report_query(&start, &end, enabled) + .into_iter() + .filter(|(name, _)| *name == "group_by[]") + .map(|(_, value)| value) + .collect::>() + }; + assert_eq!(groups(false), ["description"]); + assert_eq!(groups(true), ["description", "workspace_id"]); + assert_eq!( + cost_report_query(&start, &end, true).len(), + cost_report_query(&start, &end, false).len() + 1 + ); +} diff --git a/rust/src/providers/litellm/mod.rs b/rust/src/providers/litellm/mod.rs index d27a573668..fc66141696 100644 --- a/rust/src/providers/litellm/mod.rs +++ b/rust/src/providers/litellm/mod.rs @@ -2,6 +2,8 @@ use async_trait::async_trait; use reqwest::{Client, Url}; use serde_json::Value; +mod model_activity; + use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, @@ -61,7 +63,7 @@ impl Provider for LiteLLMProvider { let response = self .client .get(management_url(&base, "key/info")?) - .bearer_auth(key) + .bearer_auth(&key) .header("Accept", "application/json") .send() .await?; @@ -79,7 +81,20 @@ impl Provider for LiteLLMProvider { 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)) + let mut result = result_from_key_info(&value); + if ctx.optional_details_enabled + && let Some(user_id) = user_id(&value) + && let Ok(endpoint) = management_url(&base, "user/daily/activity") + { + // Optional history must never fail the budget fetch. + let today = chrono::Utc::now().date_naive(); + for row in + model_activity::fetch(&self.client, &endpoint, &key, &user_id, today).await + { + result = result.with_display_detail(Some(row)); + } + } + Ok(result) } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -132,11 +147,20 @@ fn management_url(base: &str, path: &str) -> Result { .map_err(|e| ProviderError::Other(format!("Invalid LiteLLM URL: {e}"))) } -fn result_from_key_info(value: &Value) -> ProviderFetchResult { - let root = value +fn key_info_root(value: &Value) -> &Value { + value .get("info") .or_else(|| value.get("key")) - .unwrap_or(value); + .unwrap_or(value) +} + +/// The key owner. Team-only keys carry no `user_id`, so no activity is requested. +fn user_id(value: &Value) -> Option { + string(key_info_root(value), &["user_id"]) +} + +fn result_from_key_info(value: &Value) -> ProviderFetchResult { + let root = key_info_root(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 @@ -228,6 +252,20 @@ mod tests { ); } + #[test] + fn user_id_is_read_from_the_key_info_root_only() { + assert_eq!( + user_id(&serde_json::json!({"info":{"user_id":" fixture+user "}})).as_deref(), + Some("fixture+user") + ); + assert_eq!( + user_id(&serde_json::json!({"key":{"user_id":"u"}})).as_deref(), + Some("u") + ); + assert_eq!(user_id(&serde_json::json!({"info":{"team_id":"t"}})), None); + assert_eq!(user_id(&serde_json::json!({"info":{"user_id":" "}})), None); + } + #[test] fn saved_base_url_uses_only_app_saved_key() { let mut ctx = FetchContext { diff --git a/rust/src/providers/litellm/model_activity.rs b/rust/src/providers/litellm/model_activity.rs new file mode 100644 index 0000000000..5b4ae8824f --- /dev/null +++ b/rust/src/providers/litellm/model_activity.rs @@ -0,0 +1,296 @@ +//! Optional LiteLLM per-model activity for the last 30 UTC days. +//! +//! Mirrors upstream 0.67.0 `litellm.ts`: up to three pages of +//! `GET /user/daily/activity`, every page strictly validated. Any failure +//! (HTTP status, timeout, malformed or incomplete data) omits the section +//! instead of failing the budget fetch, and never surfaces the response body. + +use std::collections::{BTreeMap, HashMap}; +use std::time::Duration; + +use chrono::{Days, NaiveDate}; +use reqwest::{Client, StatusCode, Url}; +use serde::de::{DeserializeOwned, Error as _}; +use serde::{Deserialize, Deserializer}; +use serde_json::{Number, Value}; + +use crate::core::ProviderDisplayDetail; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +pub(super) const SECTION_TITLE: &str = "Model activity \u{b7} 30d UTC"; +const WINDOW_DAYS: u64 = 29; +const MAX_PAGES: u64 = 3; +const PAGE_SIZE: &str = "1000"; +const MAX_DAYS_PER_PAGE: usize = 31; +const MAX_MODELS: usize = 1000; +const MAX_ROWS: usize = 20; +const MAX_PAGE_BYTES: usize = 8 * 1024 * 1024; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(2); +/// `Number.MAX_SAFE_INTEGER`, the bound upstream applies to every counter. +const MAX_SAFE_INTEGER: u64 = (1 << 53) - 1; + +/// Fetch and summarise model activity ending on `today` (UTC). Returns no rows +/// when the feature is unavailable, malformed, or the window is empty. +pub(super) async fn fetch( + client: &Client, + endpoint: &Url, + api_key: &str, + user_id: &str, + today: NaiveDate, +) -> Vec { + match collect(client, endpoint, api_key, user_id, today).await { + Ok(totals) => detail_rows(totals), + Err(reason) => { + tracing::debug!(reason, "LiteLLM model activity omitted"); + Vec::new() + } + } +} + +async fn collect( + client: &Client, + endpoint: &Url, + api_key: &str, + user_id: &str, + today: NaiveDate, +) -> Result, &'static str> { + let start = today + .checked_sub_days(Days::new(WINDOW_DAYS)) + .ok_or("invalid activity window")? + .format("%Y-%m-%d") + .to_string(); + let end = today.format("%Y-%m-%d").to_string(); + let mut totals = BTreeMap::new(); + + for page in 1..=MAX_PAGES { + let mut url = endpoint.clone(); + url.query_pairs_mut() + .clear() + .append_pair("user_id", user_id) + .append_pair("start_date", &start) + .append_pair("end_date", &end) + .append_pair("page", &page.to_string()) + .append_pair("page_size", PAGE_SIZE); + let response = client + .get(url) + .bearer_auth(api_key) + .header("Accept", "application/json") + .timeout(REQUEST_TIMEOUT) + .send() + .await + .map_err(|_| "model activity request failed")?; + if response.status() != StatusCode::OK { + return Err("model activity unavailable"); + } + let body = read_bounded_response(response, MAX_PAGE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => "model activity page too large", + BoundedBodyError::Read(_) => "model activity read failed", + })?; + let parsed: Object = + serde_json::from_slice(&body).map_err(|_| "invalid activity page")?; + let Page { results, metadata } = parsed.0; + if results.len() > MAX_DAYS_PER_PAGE { + return Err("invalid activity page"); + } + for row in results { + accumulate_day(&mut totals, row.0, &start, &end)?; + } + + let metadata = metadata.map_or_else(Metadata::default, |wrapped| wrapped.0); + if metadata.page.is_some_and(|reported| reported.0 != page) { + return Err("repeated activity page"); + } + let total_pages = metadata.total_pages.map_or(1, |pages| pages.0); + if !metadata.has_more.unwrap_or(false) && page >= total_pages { + return Ok(totals); + } + } + Err("incomplete model activity") +} + +fn accumulate_day( + totals: &mut BTreeMap, + row: Day, + start: &str, + end: &str, +) -> Result<(), &'static str> { + if !is_iso_day(&row.date) || row.date.as_str() < start || row.date.as_str() > end { + return Err("invalid activity day"); + } + for (name, entry) in row.breakdown.0.models.0 { + if !ProviderDisplayDetail::is_valid_title(&name) { + return Err("invalid activity model name"); + } + let counters = entry_counters(&entry)?; + totals + .entry(name) + .or_default() + .add(&counters) + .ok_or("invalid activity count")?; + if totals.len() > MAX_MODELS { + return Err("too many activity models"); + } + } + Ok(()) +} + +/// Counters live under `metrics` when present, otherwise on the entry itself. +fn entry_counters(entry: &Value) -> Result { + let object = entry.as_object().ok_or("invalid activity count")?; + let source = match object.get("metrics") { + None | Some(Value::Null) => entry, + Some(metrics) => metrics, + }; + Object::::deserialize(source) + .map(|counters| counters.0) + .map_err(|_| "invalid activity count") +} + +fn is_iso_day(value: &str) -> bool { + let bytes = value.as_bytes(); + bytes.len() == 10 + && bytes.iter().enumerate().all(|(index, byte)| match index { + 4 | 7 => *byte == b'-', + _ => byte.is_ascii_digit(), + }) +} + +fn detail_rows(totals: BTreeMap) -> Vec { + let mut ranked: Vec<(String, Totals)> = totals.into_iter().collect(); + // `BTreeMap` iterates by name, and the sort is stable, so ties keep name order. + ranked.sort_by_key(|(_, totals)| std::cmp::Reverse(totals.total)); + ranked + .into_iter() + .take(MAX_ROWS) + .enumerate() + .filter_map(|(index, (name, totals))| { + ProviderDisplayDetail::new( + format!("litellm-model-{index}"), + name, + format!( + "{} tokens \u{b7} {} requests", + totals.total, totals.requests + ), + )? + .with_secondary_value(format!( + "Input {} \u{b7} Output {}", + totals.input, totals.output + ))? + .with_section(SECTION_TITLE) + }) + .collect() +} + +#[derive(Default)] +struct Totals { + input: u64, + output: u64, + total: u64, + requests: u64, +} + +impl Totals { + /// Adds one day's counters; `None` if any running sum leaves the safe + /// integer range. + fn add(&mut self, counters: &Counters) -> Option<()> { + let sum = |current: u64, next: u64| { + current + .checked_add(next) + .filter(|value| *value <= MAX_SAFE_INTEGER) + }; + let input = sum(self.input, counters.prompt_tokens.0)?; + let output = sum(self.output, counters.completion_tokens.0)?; + let total = sum(self.total, counters.total_tokens.0)?; + let requests = sum(self.requests, counters.api_requests.0)?; + *self = Self { + input, + output, + total, + requests, + }; + Some(()) + } +} + +/// Deserializes `T` only from a JSON object, so a struct is never filled +/// positionally from an array. +struct Object(T); + +impl<'de, T: DeserializeOwned> Deserialize<'de> for Object { + fn deserialize>(deserializer: D) -> Result { + let value = Value::deserialize(deserializer)?; + if !value.is_object() { + return Err(D::Error::custom("expected a JSON object")); + } + T::deserialize(value).map(Self).map_err(D::Error::custom) + } +} + +/// A non-negative integer no larger than `Number.MAX_SAFE_INTEGER`. Integral +/// floats such as `30.0` are accepted, like a JavaScript number. +#[derive(Clone, Copy)] +struct Count(u64); + +impl<'de> Deserialize<'de> for Count { + fn deserialize>(deserializer: D) -> Result { + let number = Number::deserialize(deserializer)?; + safe_count(&number) + .map(Self) + .ok_or_else(|| D::Error::custom("invalid count")) + } +} + +fn safe_count(number: &Number) -> Option { + if let Some(value) = number.as_u64() { + return (value <= MAX_SAFE_INTEGER).then_some(value); + } + let value = number.as_f64()?; + if !(value >= 0.0 && value.fract() == 0.0 && value <= MAX_SAFE_INTEGER as f64) { + return None; + } + #[expect( + clippy::cast_possible_truncation, + reason = "value is a non-negative integer no larger than 2^53 - 1" + )] + let count = value as u64; + Some(count) +} + +#[derive(Deserialize)] +struct Counters { + prompt_tokens: Count, + completion_tokens: Count, + total_tokens: Count, + api_requests: Count, +} + +#[derive(Deserialize)] +struct Page { + results: Vec>, + #[serde(default)] + metadata: Option>, +} + +#[derive(Deserialize)] +struct Day { + date: String, + breakdown: Object, +} + +#[derive(Deserialize)] +struct Breakdown { + models: Object>, +} + +#[derive(Default, Deserialize)] +struct Metadata { + page: Option, + total_pages: Option, + has_more: Option, +} + +#[cfg(test)] +#[path = "model_activity_tests.rs"] +mod tests; diff --git a/rust/src/providers/litellm/model_activity_tests.rs b/rust/src/providers/litellm/model_activity_tests.rs new file mode 100644 index 0000000000..3ab2ab857c --- /dev/null +++ b/rust/src/providers/litellm/model_activity_tests.rs @@ -0,0 +1,293 @@ +//! Scenarios ported from upstream 0.67.0 `LiteLLMModelUsageTests.swift`, run +//! against a mockito server. The window is fixed at 2026-08-26..=2026-09-24. + +use chrono::NaiveDate; +use mockito::{Matcher, Server}; +use reqwest::{Client, Url}; + +use super::*; + +const COUNTERS: &str = + r#"{"prompt_tokens":20,"completion_tokens":10,"total_tokens":30,"api_requests":1}"#; +const NESTED: &str = + r#"{"metrics":{"prompt_tokens":20,"completion_tokens":10,"total_tokens":30,"api_requests":2}}"#; + +fn today() -> NaiveDate { + NaiveDate::from_ymd_opt(2026, 9, 24).unwrap() +} + +fn page(date: &str, models: &str, total_pages: u32, page: u32) -> String { + format!( + r#"{{"results":[{{"date":"{date}","breakdown":{{"models":{models}}}}}],"metadata":{{"total_pages":{total_pages},"page":{page}}}}}"# + ) +} + +fn normal_pages() -> Vec<(u64, u16, String)> { + let first = format!(r#"{{"fixture-alpha":{COUNTERS},"fixture-beta":{COUNTERS}}}"#); + let second = format!(r#"{{"fixture-alpha":{NESTED}}}"#); + vec![ + (1, 200, page("2026-09-24", &first, 2, 1)), + (2, 200, page("2026-09-23", &second, 2, 2)), + ] +} + +async fn run(pages: Vec<(u64, u16, String)>) -> Vec { + let mut server = Server::new_async().await; + let mut mocks = Vec::new(); + for (number, status, body) in pages { + let mock = server + .mock("GET", "/user/daily/activity") + .match_header("authorization", "Bearer fixture-key") + .match_query(Matcher::AllOf(vec![ + Matcher::UrlEncoded("user_id".into(), "fixture+user".into()), + Matcher::UrlEncoded("start_date".into(), "2026-08-26".into()), + Matcher::UrlEncoded("end_date".into(), "2026-09-24".into()), + Matcher::UrlEncoded("page_size".into(), "1000".into()), + Matcher::UrlEncoded("page".into(), number.to_string()), + ])) + .with_status(status as usize) + .with_body(body) + .expect(1) + .create_async() + .await; + mocks.push(mock); + } + let endpoint = Url::parse(&server.url()) + .unwrap() + .join("user/daily/activity") + .unwrap(); + let rows = fetch( + &Client::new(), + &endpoint, + "fixture-key", + "fixture+user", + today(), + ) + .await; + for mock in mocks { + mock.assert_async().await; + } + rows +} + +#[tokio::test] +async fn combines_flat_and_nested_metrics_across_pages() { + let rows = run(normal_pages()).await; + let titles: Vec<_> = rows.iter().map(|row| row.title()).collect(); + assert_eq!(titles, ["fixture-alpha", "fixture-beta"]); + assert_eq!(rows[0].value(), "60 tokens \u{b7} 3 requests"); + assert_eq!(rows[0].secondary_value(), Some("Input 40 \u{b7} Output 20")); + assert_eq!(rows[1].value(), "30 tokens \u{b7} 1 requests"); + assert!( + rows.iter() + .all(|row| row.section() == Some("Model activity \u{b7} 30d UTC")) + ); +} + +#[tokio::test] +async fn one_day_can_be_split_across_pages() { + let first = format!(r#"{{"fixture-alpha":{COUNTERS}}}"#); + let second = format!(r#"{{"fixture-alpha":{NESTED}}}"#); + let rows = run(vec![ + (1, 200, page("2026-09-24", &first, 2, 1)), + (2, 200, page("2026-09-24", &second, 2, 2)), + ]) + .await; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].value(), "60 tokens \u{b7} 3 requests"); +} + +#[tokio::test] +async fn forbidden_history_is_omitted() { + assert!(run(vec![(1, 403, "{}".into())]).await.is_empty()); +} + +#[tokio::test] +async fn non_200_success_status_is_omitted() { + let body = page("2026-09-24", &format!(r#"{{"a":{COUNTERS}}}"#), 1, 1); + assert!(run(vec![(1, 201, body)]).await.is_empty()); +} + +#[tokio::test] +async fn unreachable_endpoint_is_omitted() { + let endpoint = Url::parse("http://127.0.0.1:1/user/daily/activity").unwrap(); + let rows = fetch( + &Client::new(), + &endpoint, + "fixture-key", + "fixture+user", + today(), + ) + .await; + assert!(rows.is_empty()); +} + +#[tokio::test] +async fn malformed_pages_are_omitted() { + for body in [ + r#"{"results":[{"date":"2026-09-24"}]}"#, + r#"{"results":{}}"#, + r#"[]"#, + r#"not json"#, + r#"{"results":[["2026-09-24",{}]]}"#, + r#"{"results":[{"date":"2026-09-24","breakdown":{"models":[]}}]}"#, + r#"{"results":[{"date":"2026-09-24","breakdown":{"models":{"m":[20,10,30,1]}}}]}"#, + r#"{"results":[{"date":"2026-09-24","breakdown":{"models":{"m":{"prompt_tokens":1}}}}]}"#, + r#"{"results":[],"metadata":[]}"#, + r#"{"results":[],"metadata":{"has_more":"yes"}}"#, + ] { + assert!(run(vec![(1, 200, body.into())]).await.is_empty(), "{body}"); + } +} + +#[tokio::test] +async fn too_many_days_on_one_page_is_omitted() { + let day = format!(r#"{{"date":"2026-09-24","breakdown":{{"models":{{"m":{COUNTERS}}}}}}}"#); + let days = vec![day; 32].join(","); + let body = format!(r#"{{"results":[{days}]}}"#); + assert!(run(vec![(1, 200, body)]).await.is_empty()); +} + +#[tokio::test] +async fn unsafe_and_invalid_counters_are_omitted() { + for counters in [ + r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":1e100,"api_requests":0}"#, + r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":9007199254740992,"api_requests":0}"#, + r#"{"prompt_tokens":-1,"completion_tokens":0,"total_tokens":0,"api_requests":0}"#, + r#"{"prompt_tokens":1.5,"completion_tokens":0,"total_tokens":0,"api_requests":0}"#, + r#"{"prompt_tokens":"1","completion_tokens":0,"total_tokens":0,"api_requests":0}"#, + ] { + let body = page( + "2026-09-24", + &format!(r#"{{"fixture-alpha":{counters}}}"#), + 1, + 1, + ); + assert!(run(vec![(1, 200, body)]).await.is_empty(), "{counters}"); + } +} + +#[tokio::test] +async fn running_sum_past_the_safe_integer_range_is_omitted() { + let near_max = r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":9007199254740991,"api_requests":0}"#; + let one = r#"{"prompt_tokens":0,"completion_tokens":0,"total_tokens":1,"api_requests":0}"#; + let first = page("2026-09-24", &format!(r#"{{"m":{near_max}}}"#), 2, 1); + let second = page("2026-09-23", &format!(r#"{{"m":{one}}}"#), 2, 2); + assert!( + run(vec![(1, 200, first), (2, 200, second)]) + .await + .is_empty() + ); +} + +#[tokio::test] +async fn integral_float_counters_are_accepted() { + let counters = + r#"{"prompt_tokens":20.0,"completion_tokens":10.0,"total_tokens":30.0,"api_requests":1.0}"#; + let body = page( + "2026-09-24", + &format!(r#"{{"fixture-alpha":{counters}}}"#), + 1, + 1, + ); + let rows = run(vec![(1, 200, body)]).await; + assert_eq!(rows[0].value(), "30 tokens \u{b7} 1 requests"); +} + +#[tokio::test] +async fn history_still_incomplete_after_three_pages_is_omitted() { + let models = format!(r#"{{"fixture-alpha":{COUNTERS}}}"#); + let pages = (1..=3) + .map(|number| { + let date = format!("2026-09-{}", 25 - number); + (number as u64, 200, page(&date, &models, 4, number)) + }) + .collect(); + assert!(run(pages).await.is_empty()); +} + +#[tokio::test] +async fn has_more_keeps_paging_until_done() { + let models = format!(r#"{{"fixture-alpha":{COUNTERS}}}"#); + let first = format!( + r#"{{"results":[{{"date":"2026-09-24","breakdown":{{"models":{models}}}}}],"metadata":{{"has_more":true,"page":1,"total_pages":1}}}}"# + ); + let second = format!( + r#"{{"results":[{{"date":"2026-09-23","breakdown":{{"models":{models}}}}}],"metadata":{{"has_more":false,"page":2,"total_pages":2}}}}"# + ); + let rows = run(vec![(1, 200, first), (2, 200, second)]).await; + assert_eq!(rows[0].value(), "60 tokens \u{b7} 2 requests"); +} + +#[tokio::test] +async fn repeated_page_metadata_is_omitted() { + let models = format!(r#"{{"fixture-alpha":{COUNTERS}}}"#); + let repeat = page("2026-09-24", &models, 2, 1); + let rows = run(vec![(1, 200, repeat.clone()), (2, 200, repeat)]).await; + assert!(rows.is_empty()); +} + +#[tokio::test] +async fn days_outside_the_window_are_omitted() { + for date in ["2026-08-25", "2026-09-25", "2026-9-24", "20260924"] { + let body = page(date, "{}", 1, 1); + assert!(run(vec![(1, 200, body)]).await.is_empty(), "{date}"); + } +} + +#[tokio::test] +async fn window_boundary_days_are_accepted() { + for date in ["2026-08-26", "2026-09-24"] { + let body = page(date, &format!(r#"{{"m":{COUNTERS}}}"#), 1, 1); + assert_eq!(run(vec![(1, 200, body)]).await.len(), 1, "{date}"); + } +} + +#[tokio::test] +async fn invalid_model_names_are_omitted() { + for name in ["", "bad\\nname"] { + let body = page("2026-09-24", &format!(r#"{{"{name}":{COUNTERS}}}"#), 1, 1); + assert!(run(vec![(1, 200, body)]).await.is_empty(), "{name:?}"); + } +} + +#[tokio::test] +async fn empty_history_has_no_section() { + let body = r#"{"results":[],"metadata":{"has_more":false}}"#; + assert!(run(vec![(1, 200, body.into())]).await.is_empty()); +} + +#[tokio::test] +async fn displayed_models_are_limited_to_twenty_by_total_tokens() { + let models = (0..25) + .map(|index| { + format!( + r#""fixture-{index}":{{"prompt_tokens":{index},"completion_tokens":0,"total_tokens":{index},"api_requests":1}}"# + ) + }) + .collect::>() + .join(","); + let body = page("2026-09-24", &format!("{{{models}}}"), 1, 1); + let rows = run(vec![(1, 200, body)]).await; + assert_eq!(rows.len(), 20); + assert_eq!(rows.first().map(|row| row.title()), Some("fixture-24")); + assert_eq!(rows.last().map(|row| row.title()), Some("fixture-5")); +} + +#[tokio::test] +async fn equal_totals_are_ordered_by_name() { + let models = format!(r#"{{"zeta":{COUNTERS},"alpha":{COUNTERS}}}"#); + let rows = run(vec![(1, 200, page("2026-09-24", &models, 1, 1))]).await; + let titles: Vec<_> = rows.iter().map(|row| row.title()).collect(); + assert_eq!(titles, ["alpha", "zeta"]); +} + +#[tokio::test] +async fn more_than_a_thousand_models_is_omitted() { + let models = (0..1001) + .map(|index| format!(r#""m{index}":{COUNTERS}"#)) + .collect::>() + .join(","); + let body = page("2026-09-24", &format!("{{{models}}}"), 1, 1); + assert!(run(vec![(1, 200, body)]).await.is_empty()); +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..9a5289145e 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -27,6 +27,7 @@ pub const CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID: &str = "metric:extra-claude-routi mod api_keys; mod manual_cookies; +mod optional_details; mod provider_workspace; mod raw; mod status; @@ -34,6 +35,7 @@ mod types; pub use api_keys::*; pub use manual_cookies::*; +pub use optional_details::provider_has_optional_details; pub use provider_workspace::*; use raw::RawSettings; pub use status::*; diff --git a/rust/src/settings/optional_details.rs b/rust/src/settings/optional_details.rs new file mode 100644 index 0000000000..943d84698c --- /dev/null +++ b/rust/src/settings/optional_details.rs @@ -0,0 +1,61 @@ +//! Opt-in extra provider breakdowns (LiteLLM model activity, Claude workspace +//! spend). Off by default because each one costs an extra request or widens +//! the Admin API query. + +use super::Settings; +use crate::core::ProviderId; + +/// Providers that expose exactly one optional detail breakdown. +pub fn provider_has_optional_details(id: ProviderId) -> bool { + matches!(id, ProviderId::LiteLLM | ProviderId::Claude) +} + +impl Settings { + /// Whether the user opted in to `id`'s optional breakdown. Always false + /// for providers without one. + pub fn optional_details_enabled(&self, id: ProviderId) -> bool { + provider_has_optional_details(id) + && self + .provider_configs + .get(&id) + .is_some_and(|config| config.optional_details_enabled) + } + + pub fn set_optional_details_enabled(&mut self, id: ProviderId, value: bool) { + if provider_has_optional_details(id) { + self.provider_config_mut(id).optional_details_enabled = value; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::settings::ProviderConfig; + + #[test] + fn optional_details_default_off_and_scoped_to_supporting_providers() { + let mut settings = Settings::default(); + assert!(!settings.optional_details_enabled(ProviderId::LiteLLM)); + assert!(!settings.optional_details_enabled(ProviderId::Claude)); + + settings.set_optional_details_enabled(ProviderId::LiteLLM, true); + assert!(settings.optional_details_enabled(ProviderId::LiteLLM)); + assert!(!settings.optional_details_enabled(ProviderId::Claude)); + + settings.set_optional_details_enabled(ProviderId::Codex, true); + assert!(!settings.optional_details_enabled(ProviderId::Codex)); + } + + #[test] + fn optional_details_round_trips_and_stays_out_of_default_json() { + let mut config = ProviderConfig::default(); + let encoded = serde_json::to_value(&config).unwrap(); + assert!(encoded.get("optional_details_enabled").is_none()); + + config.optional_details_enabled = true; + let encoded = serde_json::to_string(&config).unwrap(); + let decoded: ProviderConfig = serde_json::from_str(&encoded).unwrap(); + assert!(decoded.optional_details_enabled); + } +} diff --git a/rust/src/settings/types.rs b/rust/src/settings/types.rs index 6113969b15..b94c499f27 100644 --- a/rust/src/settings/types.rs +++ b/rust/src/settings/types.rs @@ -410,6 +410,11 @@ pub struct ProviderConfig { /// Disabled by default; the desktop shell never injects a prompt. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub auto_resume_after_quota_reset: bool, + /// LiteLLM/Claude-only: fetch one optional extra breakdown (LiteLLM model + /// activity, Claude workspace spend). Disabled by default; each provider + /// has exactly one such breakdown. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub optional_details_enabled: bool, /// Per-provider accent color override (hex, e.g. "#FF5733"). `None` /// means the shipped brand color is used (#2972). #[serde(skip_serializing_if = "Option::is_none")] From 2a5d336d75cde86f456050aab0f4bec1484a2398 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:05:17 +0700 Subject: [PATCH 132/407] Port upstream 0.67.0: add History window setting and period-aware spend surfaces Persist cost_reporting_period, add the Settings History window picker, and thread the selected period through usage spend, tray, menu card, float bar and share PNG. thirtyDay stays a fixed 30-day column and PowerToys stays 30 days. --- .../src-tauri/src/commands/bridge.rs | 4 + .../src-tauri/src/commands/chart.rs | 244 +++++++-------- .../src-tauri/src/commands/chart/tests.rs | 163 ++++++++++ .../src-tauri/src/commands/settings.rs | 40 ++- .../src-tauri/src/commands/spend_contract.rs | 22 +- .../src-tauri/src/commands/usage_spend.rs | 281 ++++++++++-------- .../src/commands/usage_spend/tests.rs | 116 ++++++++ apps/desktop-tauri/src-tauri/src/powertoys.rs | 7 + .../src/components/MenuCard.test.tsx | 34 +++ .../src/components/MenuCardDetails.tsx | 20 +- .../src/floatbar/FloatBar.test.tsx | 34 +++ apps/desktop-tauri/src/floatbar/FloatBar.tsx | 41 ++- apps/desktop-tauri/src/i18n/keys.ts | 18 ++ apps/desktop-tauri/src/lib/costPeriod.test.ts | 122 ++++++++ apps/desktop-tauri/src/lib/costPeriod.ts | 125 ++++++++ apps/desktop-tauri/src/lib/tauri.ts | 9 +- .../src/lib/usageSpendSharing.test.ts | 43 ++- .../src/lib/usageSpendSharing.ts | 54 +++- .../src/surfaces/TrayPanel.period.test.tsx | 193 ++++++++++++ apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 49 ++- .../settings/tabs/CostPeriodControl.tsx | 107 +++++++ .../settings/tabs/UsageSpendTab.test.tsx | 215 ++++++++++++++ .../surfaces/settings/tabs/UsageSpendTab.tsx | 73 +++-- apps/desktop-tauri/src/types/bridge.ts | 19 ++ rust/src/cost_reporting_period.rs | 64 ++++ rust/src/locale.rs | 18 ++ rust/src/locale/en-US.ftl | 16 + rust/src/spend_contract.rs | 59 +++- rust/src/spend_contract/tests.rs | 43 +++ 29 files changed, 1906 insertions(+), 327 deletions(-) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs create mode 100644 apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs create mode 100644 apps/desktop-tauri/src/lib/costPeriod.test.ts create mode 100644 apps/desktop-tauri/src/lib/costPeriod.ts create mode 100644 apps/desktop-tauri/src/surfaces/TrayPanel.period.test.tsx create mode 100644 apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx create mode 100644 apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.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..e10b94a0ab 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -716,6 +716,9 @@ pub struct SettingsSnapshot { cost_summary_display_style: &'static str, open_codex_usage_logs_enabled: bool, hide_native_codex_cost_when_open_codex_present: bool, + /// History window as its persisted raw form (`rolling:N`, + /// `month-to-date`, `all`). + cost_reporting_period: String, provider_accent_colors: std::collections::HashMap, } @@ -842,6 +845,7 @@ impl From for SettingsSnapshot { open_codex_usage_logs_enabled: settings.open_codex_usage_logs_enabled, hide_native_codex_cost_when_open_codex_present: settings .hide_native_codex_cost_when_open_codex_present, + cost_reporting_period: settings.cost_reporting_period.raw(), provider_accent_colors: settings .provider_configs .iter() diff --git a/apps/desktop-tauri/src-tauri/src/commands/chart.rs b/apps/desktop-tauri/src-tauri/src/commands/chart.rs index 4561dd9129..fa49349028 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/chart.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/chart.rs @@ -9,6 +9,7 @@ use crate::commands::bridge::RateWindowSnapshot; use crate::state::AppState; use chrono::{DateTime, Utc}; use codexbar::core::{OpenAIDashboardCacheStore, RateWindow}; +use codexbar::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; use codexbar::cost_scanner::{ CostScanner, CostSummary, get_daily_cost_history, get_daily_token_history, }; @@ -63,8 +64,18 @@ pub struct DailyUsageBreakdown { #[serde(rename_all = "camelCase")] pub struct ProviderLocalUsageSummary { pub today_cost: Option, + /// Always the trailing 30 days (the published PowerToys pipe reads it). pub thirty_day_cost: Option, pub thirty_day_tokens: Option, + /// Cost over the selected History window (`reporting_period`). + #[serde(default)] + pub period_cost: Option, + /// Tokens over the selected History window. + #[serde(default)] + pub period_tokens: Option, + /// Raw reporting period the `period_*` fields cover; empty when unknown. + #[serde(default)] + pub reporting_period: String, pub latest_tokens: Option, pub top_model: Option, pub estimate_note: String, @@ -194,10 +205,17 @@ fn build_provider_chart_data_with_cancel( .cloned() .map(|(date, tokens)| DailyTokenPoint { date, tokens }) .collect(); - let local_usage = snapshot - .local_summary - .as_ref() - .and_then(|summary| local_usage_summary_from_cost_summary(&provider_id, summary)); + let period = current_reporting_period(); + let local_usage = snapshot.local_summary.as_ref().and_then(|summary| { + let period_summary = + period_summary_for(&provider_id, period, summary, cancel.as_deref()); + local_usage_summary_from_cost_summary( + &provider_id, + summary, + period, + period_summary.as_ref(), + ) + }); store_local_usage_summary(&provider_id, local_usage.clone()); ( cost_history, @@ -221,7 +239,14 @@ fn build_provider_chart_data_with_cancel( tokens: day.total_tokens, }) .collect(); - let local_usage = muse_local_usage_summary(&report, locale::current_language()); + let period = current_reporting_period(); + let period_report = muse_period_report(period, &report, cancel.as_deref()); + let local_usage = muse_local_usage_summary( + &report, + &period_report, + period, + locale::current_language(), + ); ( Vec::new(), tokens_history, @@ -331,16 +356,40 @@ fn map_quota_window_history( } } +/// The History window the desktop is configured to report. +fn current_reporting_period() -> CostReportingPeriod { + codexbar::settings::Settings::load().cost_reporting_period +} + +/// The period summary for `provider_id`, reusing the 30-day scan when the +/// selected window is exactly the trailing 30 days. +fn period_summary_for( + provider_id: &str, + period: CostReportingPeriod, + thirty_day: &CostSummary, + cancel: Option<&AtomicBool>, +) -> Option { + if period == CostReportingPeriod::Rolling(30) { + return Some(thirty_day.clone()); + } + scan_local_cost(provider_id, period, cancel) +} + fn local_usage_summary_from_cost_summary( provider_id: &str, summary: &CostSummary, + period: CostReportingPeriod, + period_summary: Option<&CostSummary>, ) -> Option { - let total_tokens = total_tokens(summary); - let has_usage = summary.sessions_count > 0 || summary.total_cost_usd > 0.0 || total_tokens > 0; + let thirty_tokens = total_tokens(summary); + let has_usage = summary.sessions_count > 0 || summary.total_cost_usd > 0.0 || thirty_tokens > 0; has_usage.then(|| ProviderLocalUsageSummary { today_cost: None, thirty_day_cost: non_zero_f64(summary.total_cost_usd), - thirty_day_tokens: non_zero_u64(total_tokens), + thirty_day_tokens: non_zero_u64(thirty_tokens), + period_cost: period_summary.and_then(|s| non_zero_f64(s.total_cost_usd)), + period_tokens: period_summary.and_then(|s| non_zero_u64(total_tokens(s))), + reporting_period: period.raw(), latest_tokens: None, top_model: top_model(summary), estimate_note: localized_estimate_note(provider_id, locale::current_language()), @@ -388,27 +437,35 @@ fn load_local_usage_summary_with_unknown_models( provider_id: &str, cancel: Option<&AtomicBool>, ) -> (Option, HashSet) { + let period = current_reporting_period(); if provider_id == "muse" { let summary = if cancel.is_some_and(|flag| flag.load(Ordering::Relaxed)) { None } else { let report = muse_local_usage::scan(30, cancel); - muse_local_usage_summary(&report, locale::current_language()) + let period_report = muse_period_report(period, &report, cancel); + muse_local_usage_summary(&report, &period_report, period, locale::current_language()) }; return (summary, HashSet::new()); } - let Some(thirty_day) = scan_local_cost(provider_id, 30, cancel) else { + let Some(thirty_day) = scan_local_cost(provider_id, CostReportingPeriod::Rolling(30), cancel) + else { return (None, HashSet::new()); }; if cancel.is_some_and(|flag| flag.load(Ordering::Relaxed)) { return (None, HashSet::new()); } - let today = scan_local_cost(provider_id, 1, cancel).unwrap_or_default(); - let unknown_models = thirty_day + let today = + scan_local_cost(provider_id, CostReportingPeriod::Rolling(1), cancel).unwrap_or_default(); + let period_summary = period_summary_for(provider_id, period, &thirty_day, cancel); + let mut unknown_models: HashSet = thirty_day .unknown_models .union(&today.unknown_models) .cloned() .collect(); + if let Some(period_summary) = period_summary.as_ref() { + unknown_models.extend(period_summary.unknown_models.iter().cloned()); + } let thirty_day_tokens = total_tokens(&thirty_day); let latest_tokens = total_tokens(&today); @@ -424,6 +481,13 @@ fn load_local_usage_summary_with_unknown_models( today_cost: non_zero_f64(today.total_cost_usd), thirty_day_cost: non_zero_f64(thirty_day.total_cost_usd), thirty_day_tokens: non_zero_u64(thirty_day_tokens), + period_cost: period_summary + .as_ref() + .and_then(|s| non_zero_f64(s.total_cost_usd)), + period_tokens: period_summary + .as_ref() + .and_then(|s| non_zero_u64(total_tokens(s))), + reporting_period: period.raw(), latest_tokens: non_zero_u64(latest_tokens), top_model: top_model(&thirty_day), estimate_note: localized_estimate_note(provider_id, lang), @@ -433,8 +497,23 @@ fn load_local_usage_summary_with_unknown_models( ) } +/// The Muse report for the selected period, reusing the 30-day scan when the +/// selection is exactly the trailing 30 days. +fn muse_period_report( + period: CostReportingPeriod, + thirty_day: &muse_local_usage::Report, + cancel: Option<&AtomicBool>, +) -> muse_local_usage::Report { + if period == CostReportingPeriod::Rolling(30) { + return thirty_day.clone(); + } + muse_local_usage::scan(period.scan_days(Utc::now()), cancel) +} + fn muse_local_usage_summary( report: &muse_local_usage::Report, + period_report: &muse_local_usage::Report, + period: CostReportingPeriod, lang: codexbar::settings::Language, ) -> Option { if !report.is_available() || !report.is_complete() { @@ -445,6 +524,12 @@ fn muse_local_usage_summary( today_cost: None, thirty_day_cost: None, thirty_day_tokens: Some(total_tokens), + period_cost: None, + period_tokens: period_report + .is_complete() + .then_some(period_report.total_tokens) + .flatten(), + reporting_period: period.raw(), latest_tokens: report.today_tokens, top_model: report.top_model.clone(), estimate_note: locale::get_text(lang, LocaleKey::PanelEstimatedFromLocalLogsMuse), @@ -460,9 +545,16 @@ pub(crate) fn load_provider_local_usage_summary( struct CachedLocalUsage { loaded_at: Instant, + /// `CostReportingPeriod::identity` the entry was built for; a period or + /// month change makes the entry stale before its TTL. + period_identity: String, summary: Option, } +fn local_usage_period_identity() -> String { + current_reporting_period().identity(Utc::now(), CostTimeZone::Local) +} + fn local_usage_cache() -> &'static Mutex> { static CACHE: OnceLock>> = OnceLock::new(); CACHE.get_or_init(|| Mutex::new(HashMap::new())) @@ -547,8 +639,10 @@ fn load_local_usage_summary_cached( cancel: Option<&AtomicBool>, ) -> Option { let cache = local_usage_cache(); + let period_identity = local_usage_period_identity(); if let Ok(guard) = cache.lock() && let Some(entry) = guard.get(provider_id) + && entry.period_identity == period_identity && token_cost_cache_is_fresh(Some(entry.loaded_at), Instant::now(), LOCAL_USAGE_TTL) { return entry.summary.clone(); @@ -568,11 +662,13 @@ fn load_local_usage_summary_cached( } fn store_local_usage_summary(provider_id: &str, summary: Option) { + let period_identity = local_usage_period_identity(); if let Ok(mut guard) = local_usage_cache().lock() { guard.insert( provider_id.to_string(), CachedLocalUsage { loaded_at: Instant::now(), + period_identity, summary, }, ); @@ -585,11 +681,13 @@ fn record_local_usage_fetch_failure(provider_id: &str, failure: CostFetchFailure } else { Instant::now() }; + let period_identity = local_usage_period_identity(); if let Ok(mut guard) = local_usage_cache().lock() { guard.insert( provider_id.to_string(), CachedLocalUsage { loaded_at, + period_identity, summary: None, }, ); @@ -637,10 +735,10 @@ fn localized_estimate_note(provider_id: &str, lang: codexbar::settings::Language fn scan_local_cost( provider_id: &str, - days: u32, + period: CostReportingPeriod, cancel: Option<&AtomicBool>, ) -> Option { - let scanner = CostScanner::new(days); + let scanner = CostScanner::for_period(period); match provider_id { "codex" => Some(scanner.scan_codex_with_cancel(cancel)), "claude" => Some(scanner.scan_claude_with_cancel(cancel)), @@ -744,120 +842,4 @@ pub(crate) fn load_openai_dashboard_chart_data_for_test( } #[cfg(test)] -mod tests { - use super::{ - CostFetchFailure, ProviderLocalUsageSummary, cost_fetch_failure_allows_early_retry, - localized_estimate_note, muse_local_usage_summary, token_cost_cache_is_fresh, - }; - use crate::commands::is_provider_cache_fresh; - use codexbar::providers::muse::local_usage::{DailyUsage, Report}; - use codexbar::settings::Language; - use codexbar::spend_contract::LocalHistoryCoverage; - use std::time::{Duration, Instant}; - - #[test] - fn token_cost_age_does_not_use_provider_quota_age() { - let now = Instant::now(); - let token_loaded = now - Duration::from_secs(31); - let provider_updated = now; - assert!(!token_cost_cache_is_fresh( - Some(token_loaded), - now, - Duration::from_secs(30) - )); - assert!(is_provider_cache_fresh( - Some(provider_updated), - Duration::from_secs(30) - )); - } - - #[test] - fn fast_cost_failures_allow_the_next_pass_to_retry() { - assert!(cost_fetch_failure_allows_early_retry( - CostFetchFailure::Failed - )); - assert!(!cost_fetch_failure_allows_early_retry( - CostFetchFailure::TimedOut - )); - } - - #[test] - fn local_usage_summary_serializes_token_cost_timestamp() { - let summary = ProviderLocalUsageSummary { - today_cost: Some(1.0), - thirty_day_cost: Some(2.0), - thirty_day_tokens: Some(300), - latest_tokens: Some(40), - top_model: Some("gpt-5".to_string()), - estimate_note: "estimated".to_string(), - token_cost_updated_at_ms: 1234, - }; - - let json = serde_json::to_value(summary).expect("serialize summary"); - assert_eq!( - json.get("tokenCostUpdatedAtMs").and_then(|v| v.as_i64()), - Some(1234) - ); - } - - #[test] - fn muse_local_usage_summary_exposes_complete_tokens_without_cost() { - let report = Report { - daily: vec![DailyUsage { - day: "2026-09-20".to_string(), - input_tokens: 10, - output_tokens: 2, - cache_read_tokens: 0, - cache_write_tokens: 0, - reasoning_tokens: 0, - total_tokens: 12, - request_count: 1, - models: vec![("muse-spark-1.3".to_string(), 12)], - }], - total_tokens: Some(12), - today_tokens: Some(12), - session_count: 1, - top_model: Some("muse-spark-1.3".to_string()), - coverage: LocalHistoryCoverage::Complete, - }; - let summary = muse_local_usage_summary(&report, codexbar::settings::Language::default()) - .expect("complete history is visible"); - assert_eq!(summary.today_cost, None); - assert_eq!(summary.thirty_day_cost, None); - assert_eq!(summary.thirty_day_tokens, Some(12)); - assert_eq!(summary.latest_tokens, Some(12)); - assert_eq!(summary.top_model.as_deref(), Some("muse-spark-1.3")); - - let partial = Report { - coverage: LocalHistoryCoverage::Partial, - ..report - }; - assert!( - muse_local_usage_summary(&partial, codexbar::settings::Language::default()).is_none() - ); - } - - #[test] - fn japanese_estimate_note_is_localized() { - assert_eq!( - localized_estimate_note("codex", Language::Japanese), - "ローカルログから推定したもので、請求書と異なる場合があります" - ); - assert_eq!( - localized_estimate_note("claude", Language::Japanese), - "ClaudeのローカルログからAPIレートで推定したもので、トークン総数が請求書と異なる場合があります" - ); - } - - #[test] - fn english_estimate_note_is_localized() { - assert_eq!( - localized_estimate_note("codex", Language::English), - "Estimated from local logs; may differ from your bill" - ); - assert_eq!( - localized_estimate_note("claude", Language::English), - "Estimated from local Claude logs at API rates; token totals may differ from your bill" - ); - } -} +mod tests; diff --git a/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs new file mode 100644 index 0000000000..a9a49dbdf4 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/chart/tests.rs @@ -0,0 +1,163 @@ +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, +}; +use crate::commands::is_provider_cache_fresh; +use codexbar::cost_reporting_period::CostReportingPeriod; +use codexbar::cost_scanner::CostSummary; +use codexbar::providers::muse::local_usage::{DailyUsage, Report}; +use codexbar::settings::Language; +use codexbar::spend_contract::LocalHistoryCoverage; +use std::time::{Duration, Instant}; + +#[test] +fn token_cost_age_does_not_use_provider_quota_age() { + let now = Instant::now(); + let token_loaded = now - Duration::from_secs(31); + let provider_updated = now; + assert!(!token_cost_cache_is_fresh( + Some(token_loaded), + now, + Duration::from_secs(30) + )); + assert!(is_provider_cache_fresh( + Some(provider_updated), + Duration::from_secs(30) + )); +} + +#[test] +fn fast_cost_failures_allow_the_next_pass_to_retry() { + assert!(cost_fetch_failure_allows_early_retry( + CostFetchFailure::Failed + )); + assert!(!cost_fetch_failure_allows_early_retry( + CostFetchFailure::TimedOut + )); +} + +#[test] +fn local_usage_summary_serializes_token_cost_timestamp() { + let summary = ProviderLocalUsageSummary { + today_cost: Some(1.0), + thirty_day_cost: Some(2.0), + thirty_day_tokens: Some(300), + period_cost: Some(3.0), + period_tokens: Some(400), + reporting_period: "month-to-date".to_string(), + latest_tokens: Some(40), + top_model: Some("gpt-5".to_string()), + estimate_note: "estimated".to_string(), + token_cost_updated_at_ms: 1234, + }; + + let json = serde_json::to_value(summary).expect("serialize summary"); + assert_eq!( + json.get("tokenCostUpdatedAtMs").and_then(|v| v.as_i64()), + Some(1234) + ); +} + +#[test] +fn muse_local_usage_summary_exposes_complete_tokens_without_cost() { + let report = Report { + daily: vec![DailyUsage { + day: "2026-09-20".to_string(), + input_tokens: 10, + output_tokens: 2, + cache_read_tokens: 0, + cache_write_tokens: 0, + reasoning_tokens: 0, + total_tokens: 12, + request_count: 1, + models: vec![("muse-spark-1.3".to_string(), 12)], + }], + total_tokens: Some(12), + today_tokens: Some(12), + session_count: 1, + top_model: Some("muse-spark-1.3".to_string()), + coverage: LocalHistoryCoverage::Complete, + }; + let period = CostReportingPeriod::Rolling(30); + let summary = muse_local_usage_summary( + &report, + &report, + period, + codexbar::settings::Language::default(), + ) + .expect("complete history is visible"); + assert_eq!(summary.period_tokens, Some(12)); + assert_eq!(summary.reporting_period, "rolling:30"); + assert_eq!(summary.today_cost, None); + assert_eq!(summary.thirty_day_cost, None); + assert_eq!(summary.thirty_day_tokens, Some(12)); + assert_eq!(summary.latest_tokens, Some(12)); + assert_eq!(summary.top_model.as_deref(), Some("muse-spark-1.3")); + + let partial = Report { + coverage: LocalHistoryCoverage::Partial, + ..report + }; + assert!( + muse_local_usage_summary( + &partial, + &partial, + period, + codexbar::settings::Language::default() + ) + .is_none() + ); +} + +#[test] +fn local_usage_keeps_thirty_day_fixed_while_period_follows_selection() { + let thirty = CostSummary { + total_cost_usd: 30.0, + input_tokens: 3_000, + output_tokens: 0, + ..CostSummary::default() + }; + let month = CostSummary { + total_cost_usd: 5.0, + input_tokens: 400, + output_tokens: 100, + ..CostSummary::default() + }; + let summary = local_usage_summary_from_cost_summary( + "claude", + &thirty, + CostReportingPeriod::MonthToDate, + Some(&month), + ) + .expect("usage is visible"); + assert_eq!(summary.thirty_day_cost, Some(30.0)); + assert_eq!(summary.thirty_day_tokens, Some(3_000)); + assert_eq!(summary.period_cost, Some(5.0)); + assert_eq!(summary.period_tokens, Some(500)); + assert_eq!(summary.reporting_period, "month-to-date"); +} + +#[test] +fn japanese_estimate_note_is_localized() { + assert_eq!( + localized_estimate_note("codex", Language::Japanese), + "ローカルログから推定したもので、請求書と異なる場合があります" + ); + assert_eq!( + localized_estimate_note("claude", Language::Japanese), + "ClaudeのローカルログからAPIレートで推定したもので、トークン総数が請求書と異なる場合があります" + ); +} + +#[test] +fn english_estimate_note_is_localized() { + assert_eq!( + localized_estimate_note("codex", Language::English), + "Estimated from local logs; may differ from your bill" + ); + assert_eq!( + localized_estimate_note("claude", Language::English), + "Estimated from local Claude logs at API rates; token totals may differ from your bill" + ); +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..9cc8715582 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -85,6 +85,8 @@ pub struct SettingsUpdate { pub cost_summary_display_style: Option, pub open_codex_usage_logs_enabled: Option, pub hide_native_codex_cost_when_open_codex_present: Option, + /// History window: `rolling:N` (1..=365), `month-to-date`, or `all`. + pub cost_reporting_period: Option, } impl SettingsUpdate { @@ -103,6 +105,7 @@ impl SettingsUpdate { || self.low_power_mode_preference.is_some() || self.adaptive_refresh.is_some() || self.codex_custom_sessions_dirs.is_some() + || self.cost_reporting_period.is_some() || self.high_usage_threshold.is_some() || self.critical_usage_threshold.is_some() || self.provider_usage_thresholds.is_some() @@ -132,6 +135,7 @@ impl SettingsUpdate { || self.provider_hidden_usage_item_ids.is_some() || self.codex_spark_usage_visible.is_some() || self.copilot_seat_credit_entitlement.is_some() + || self.cost_reporting_period.is_some() || self.enabled_providers.is_some() || self.ui_language.is_some() } @@ -167,6 +171,13 @@ impl SettingsUpdate { if let Some(v) = self.refresh_all_providers_on_menu_open { settings.refresh_all_providers_on_menu_open = v; } + if let Some(v) = self + .cost_reporting_period + .as_deref() + .and_then(codexbar::cost_reporting_period::CostReportingPeriod::parse) + { + settings.cost_reporting_period = v; + } if let Some(v) = self.open_codex_usage_logs_enabled { settings.open_codex_usage_logs_enabled = v; } @@ -417,6 +428,11 @@ impl SettingsUpdate { { return Err(format!("Invalid low power mode preference: {value}")); } + if let Some(value) = self.cost_reporting_period.as_deref() + && codexbar::cost_reporting_period::CostReportingPeriod::parse(value).is_none() + { + return Err(format!("Invalid cost reporting period: {value}")); + } if let Some(value) = self.copilot_seat_credit_entitlement { settings.set_seat_credit_entitlement(codexbar::core::ProviderId::Copilot, value)?; } @@ -506,7 +522,8 @@ pub async fn update_settings( let mut settings = Settings::load(); let notify_float_bar = patch.notifies_float_bar(); let refresh_provider_data = patch.refreshes_provider_data(); - let clear_local_usage_cache = patch.codex_custom_sessions_dirs.is_some(); + let clear_local_usage_cache = + patch.codex_custom_sessions_dirs.is_some() || patch.cost_reporting_period.is_some(); let rebuild_tray_menu = patch.rebuilds_tray_menu(); let refresh_tray_presentation = patch.refreshes_tray_presentation(); let tray_promotion_changed = patch.changes_tray_promotion(); @@ -684,6 +701,27 @@ mod tests { assert_eq!(settings.seat_credit_entitlement(ProviderId::Copilot), None); } + #[test] + fn cost_reporting_period_update_is_validated_and_applied() { + use codexbar::cost_reporting_period::CostReportingPeriod; + + let mut settings = Settings::default(); + let bad: SettingsUpdate = + serde_json::from_str(r#"{"costReportingPeriod":"rolling:0"}"#).unwrap(); + let error = bad + .apply_to(&mut settings) + .expect_err("zero-day window must be rejected"); + assert_eq!(error, "Invalid cost reporting period: rolling:0"); + + let good: SettingsUpdate = + serde_json::from_str(r#"{"costReportingPeriod":"month-to-date"}"#).unwrap(); + good.apply_to(&mut settings).expect("valid period applies"); + assert_eq!( + settings.cost_reporting_period, + CostReportingPeriod::MonthToDate + ); + } + #[test] fn display_settings_that_affect_tray_trigger_presentation_refresh() { assert!( diff --git a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs index 4403ba7ad0..ee926ab736 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/spend_contract.rs @@ -1,13 +1,18 @@ //! Upstream 0.53 Usage & Spend accounting bridge. +use codexbar::cost_reporting_period::CostReportingPeriod; use codexbar::cost_scanner::CostScanner; use codexbar::settings::Settings; -use codexbar::spend_contract::{SpendContract, build_local_spend_contract_from_summary}; +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. #[tauri::command] pub async fn get_spend_contract( provider_id: String, history_days: Option, + period: Option, include_open_codex: Option, ) -> Result { let provider = provider_id.trim().to_ascii_lowercase(); @@ -16,11 +21,15 @@ pub async fn get_spend_contract( "Spend contract is unavailable for provider: {provider}" )); } - let days = history_days.unwrap_or(30); let include_import = include_open_codex.unwrap_or(false) && provider == "codex"; tauri::async_runtime::spawn_blocking(move || { - let history_days = if days == 0 { 365 } else { days.clamp(1, 365) }; - let scanner = CostScanner::new(history_days); + let settings = Settings::load(); + let period = CostReportingPeriod::resolve_request( + period.as_deref(), + history_days, + settings.cost_reporting_period, + ); + let scanner = CostScanner::for_period(period); let summary = match provider.as_str() { "codex" => scanner.scan_codex(), "claude" => scanner.scan_claude(), @@ -28,10 +37,9 @@ pub async fn get_spend_contract( "opencodego" => scanner.scan_opencodego_with_cancel(None), _ => unreachable!(), }; - let settings = Settings::load(); - build_local_spend_contract_from_summary( + build_contract_from_period_summary( &provider, - history_days, + period, include_import, settings.hide_native_codex_cost_when_open_codex_present && provider == "codex", settings.hide_personal_info, 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..4297905f10 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -1,8 +1,12 @@ -//! Usage & Spend settings tab: 7-day / 30-day local cost aggregates. +//! Usage & Spend settings tab: 7-day / 30-day compat columns plus the +//! selected reporting period (History window). -use codexbar::cost_scanner::{CostScanner, CostSummary}; +use chrono::Utc; +use codexbar::cost_reporting_period::{CostReportingPeriod, CostTimeZone}; +use codexbar::cost_scanner::CostScanner; use codexbar::spend_contract::{ - SpendContract, build_local_spend_contract, build_local_spend_contract_from_summary, + SpendContract, build_contract_from_period_summary, build_local_spend_contract, + build_local_spend_contract_for_period, build_local_spend_contract_from_summary, }; use serde::Serialize; use tauri::State; @@ -28,6 +32,10 @@ pub struct UsageSpendRow { pub thirty_day: Option, pub seven_day_tokens: Option, pub thirty_day_tokens: Option, + /// Cost over the selected reporting period (History window). + pub period_cost: Option, + /// Tokens over the selected reporting period. + pub period_tokens: Option, pub currency: String, pub source: String, /// Included in the shared Overview spend denominator. @@ -50,6 +58,8 @@ struct SpendValues { thirty_day: Option, seven_day_tokens: Option, thirty_day_tokens: Option, + period_cost: Option, + period_tokens: Option, source: String, refreshing: bool, stale_updated_at: Option, @@ -60,6 +70,8 @@ struct SpendValues { pub struct UsageSpendSummary { pub rows: Vec, pub contract: SpendContract, + /// Raw reporting period the rows' `period*` columns were built for. + pub reporting_period: String, pub reporting_day: String, pub dashboard_timezone: String, } @@ -182,6 +194,7 @@ struct BuiltUsageSpendSummary { pub async fn get_usage_spend_summary( state: State<'_, Mutex>, history_days: Option, + period: Option, force_refresh: Option, ) -> Result { let cached = { @@ -189,10 +202,16 @@ pub async fn get_usage_spend_summary( guard.provider_cache.clone() }; - let selected_days = history_days.unwrap_or(30); + // An explicit `period` wins, then the legacy `history_days` (0 = All), + // then 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); let built = tauri::async_runtime::spawn_blocking(move || { - build_usage_spend_summary_cached(&cached, selected_days, force_refresh) + build_usage_spend_summary_cached(&cached, period, force_refresh) }) .await .map_err(|e| format!("usage spend worker failed: {e}"))??; @@ -202,7 +221,7 @@ pub async fn get_usage_spend_summary( .map(|guard| guard.provider_cache.clone())?; let current_key = usage_spend_cache_key( ¤t_cached, - selected_days, + period, &codexbar::settings::Settings::load(), ); if current_key != built.key { @@ -231,11 +250,11 @@ pub fn write_usage_spend_export(path: String, payload: String) -> Result<(), Str fn build_usage_spend_summary_cached( cached: &[ProviderUsageSnapshot], - selected_days: u32, + period: CostReportingPeriod, force_refresh: bool, ) -> Result { let settings = codexbar::settings::Settings::load(); - let key = usage_spend_cache_key(cached, selected_days, &settings); + let key = usage_spend_cache_key(cached, period, &settings); { let guard = usage_spend_coordinator() .lock() @@ -257,7 +276,7 @@ fn build_usage_spend_summary_cached( .map_err(|error| error.to_string())?; coordinator.begin(key.clone()) }; - let summary = build_usage_spend_summary(cached, selected_days, &settings, force_refresh); + let summary = build_usage_spend_summary(cached, period, &settings, force_refresh); let refreshing = summary_is_refreshing(&summary); let codex_scan_pause_reason = codexbar::core::JsonlScanner::load_cache_status(codexbar::core::ProviderId::Codex, None) @@ -299,12 +318,12 @@ fn build_usage_spend_summary_cached( fn usage_spend_cache_key( cached: &[ProviderUsageSnapshot], - selected_days: u32, + period: CostReportingPeriod, settings: &codexbar::settings::Settings, ) -> String { usage_spend_cache_key_with_privacy( cached, - selected_days, + &period.identity(Utc::now(), CostTimeZone::Local), settings.open_codex_usage_logs_enabled, settings.hide_native_codex_cost_when_open_codex_present, settings.hide_personal_info, @@ -313,7 +332,7 @@ fn usage_spend_cache_key( fn usage_spend_cache_key_with_privacy( cached: &[ProviderUsageSnapshot], - selected_days: u32, + period_identity: &str, include_opencodex: bool, hide_native: bool, hide_personal_info: bool, @@ -347,7 +366,7 @@ fn usage_spend_cache_key_with_privacy( format!( "{}|{}|{}|{}|{}|{}", chrono::Local::now().date_naive(), - selected_days, + period_identity, include_opencodex, hide_native, hide_personal_info, @@ -357,10 +376,12 @@ fn usage_spend_cache_key_with_privacy( fn build_usage_spend_summary( cached: &[ProviderUsageSnapshot], - selected_days: u32, + period: CostReportingPeriod, settings: &codexbar::settings::Settings, force_refresh: bool, ) -> UsageSpendSummary { + let now = Utc::now(); + let period_scan_days = period.scan_days(now); 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") @@ -377,34 +398,51 @@ 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. let ( - (codex_7_summary, codex_30_summary), - (claude_7_summary, claude_30_summary), - (pi_7_summary, pi_30_summary), + (codex_7_summary, codex_30_summary, codex_period_summary), + (claude_7_summary, claude_30_summary, claude_period_summary), + (pi_7_summary, pi_30_summary, pi_period_summary), ) = std::thread::scope(|scope| { let codex = scope.spawn(move || { - ( - CostScanner::new(7) - .with_options(codex_scan_options) - .scan_codex(), - CostScanner::new(30) + let seven = CostScanner::new(7) + .with_options(codex_scan_options) + .scan_codex(); + 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) .with_options(codex_scan_options) .scan_codex(), - ) + }; + (seven, thirty, selected) }); let claude = scope.spawn(|| { - ( - CostScanner::new(7) - .scan_claude_with_cancel_and_pi_sessions(None, include_pi_in_native), - CostScanner::new(30) + let seven = CostScanner::new(7) + .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), - ) + }; + (seven, thirty, selected) }); let pi = scope.spawn(|| { - ( - CostScanner::new(7).scan_pi(), - CostScanner::new(30).scan_pi(), - ) + 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(), + }; + (seven, thirty, selected) }); ( codex.join().expect("Codex spend scan worker panicked"), @@ -438,6 +476,22 @@ fn build_usage_spend_summary( settings.hide_personal_info, codex_30_summary.clone(), ); + let codex_period_contract = build_contract_from_period_summary( + "codex", + period, + include_opencodex, + hide_native, + settings.hide_personal_info, + codex_period_summary, + ); + let pi_period_contract = build_contract_from_period_summary( + "pi", + period, + false, + false, + settings.hide_personal_info, + pi_period_summary, + ); let pi_7_contract = build_local_spend_contract_from_summary( "pi", 7, @@ -502,6 +556,8 @@ fn build_usage_spend_summary( thirty_day: codex_30_contract.known_cost_usd, seven_day_tokens: total_token_mix(&codex_7_contract.token_mix), thirty_day_tokens: total_token_mix(&codex_30_contract.token_mix), + period_cost: codex_period_contract.known_cost_usd, + period_tokens: total_token_mix(&codex_period_contract.token_mix), source: if include_opencodex && !codex_30_contract.imports.is_empty() { "local logs + OpenCodex".to_string() } else { @@ -523,6 +579,12 @@ fn build_usage_spend_summary( .input_tokens .saturating_add(claude_30_summary.output_tokens), ), + period_cost: Some(claude_period_summary.total_cost_usd), + period_tokens: Some( + claude_period_summary + .input_tokens + .saturating_add(claude_period_summary.output_tokens), + ), source: "local logs".to_string(), refreshing: false, stale_updated_at: None, @@ -532,6 +594,8 @@ fn build_usage_spend_summary( 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), + period_cost: pi_period_contract.known_cost_usd, + period_tokens: total_token_mix(&pi_period_contract.token_mix), source: "local Pi/OMP history".to_string(), refreshing: !pi_30_summary.history_coverage_established, stale_updated_at: None, @@ -540,11 +604,15 @@ fn build_usage_spend_summary( let seven = build_local_spend_contract(&provider_id, 7, true); let thirty = build_local_spend_contract(&provider_id, 30, true); if !thirty.imports.is_empty() { + let selected = + build_local_spend_contract_for_period(&provider_id, period, true); SpendValues { seven_day: seven.known_cost_usd, thirty_day: thirty.known_cost_usd, seven_day_tokens: total_token_mix(&seven.token_mix), thirty_day_tokens: total_token_mix(&thirty.token_mix), + period_cost: selected.known_cost_usd, + period_tokens: total_token_mix(&selected.token_mix), source: if provider_id == "opencodego" { "local logs + OpenCodex".to_string() } else { @@ -554,32 +622,39 @@ fn build_usage_spend_summary( stale_updated_at: None, } } else { - cached_spend(cached_snapshot) + cached_spend(cached_snapshot, period, now) } } "cursor" => { let seven = codexbar::providers::cursor::local_csv::summarize(7); let thirty = codexbar::providers::cursor::local_csv::summarize(30); if thirty.row_count > 0 { + let selected = + codexbar::providers::cursor::local_csv::summarize(period_scan_days); SpendValues { seven_day: (seven.row_count > 0).then_some(seven.total_cost_usd), thirty_day: Some(thirty.total_cost_usd), seven_day_tokens: (seven.row_count > 0).then_some(seven.total_tokens), thirty_day_tokens: Some(thirty.total_tokens), + period_cost: (selected.row_count > 0).then_some(selected.total_cost_usd), + period_tokens: (selected.row_count > 0).then_some(selected.total_tokens), source: "local Cursor tokscale cache".to_string(), refreshing: false, stale_updated_at: None, } } else { - cached_spend(cached_snapshot) + cached_spend(cached_snapshot, period, now) } } "grok" => { let seven = codexbar::providers::grok::local_sessions::summarize(7); let thirty = codexbar::providers::grok::local_sessions::summarize(30); - let mut spend = cached_spend(cached_snapshot); + let selected = + codexbar::providers::grok::local_sessions::summarize(period_scan_days); + let mut spend = cached_spend(cached_snapshot, period, now); spend.seven_day_tokens = (seven.session_count > 0).then_some(seven.total_tokens); spend.thirty_day_tokens = (thirty.session_count > 0).then_some(thirty.total_tokens); + spend.period_tokens = (selected.session_count > 0).then_some(selected.total_tokens); if thirty.session_count > 0 { spend.source = "local Grok sessions".to_string(); } @@ -589,17 +664,21 @@ fn build_usage_spend_summary( 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); + let selected = + codexbar::providers::antigravity::local_sessions::summarize(period_scan_days); + let mut spend = cached_spend(cached_snapshot, period, now); 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); + spend.period_tokens = matches!(selected.coverage, LocalHistoryCoverage::Complete) + .then_some(selected.total_tokens); if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { spend.source = "local Antigravity history".to_string(); } spend } - _ => cached_spend(cached_snapshot), + _ => cached_spend(cached_snapshot, period, now), }; let currency = cached_snapshot @@ -625,6 +704,8 @@ fn build_usage_spend_summary( thirty_day: spend.thirty_day, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, + period_cost: spend.period_cost, + period_tokens: spend.period_tokens, currency, source: spend.source, included_in_overview: include_in_shared_overview( @@ -638,31 +719,13 @@ fn build_usage_spend_summary( }); } - let history_days = if selected_days == 0 { - 365 - } else { - selected_days.clamp(1, 365) - }; - let selected_summary: CostSummary = match history_days { - 7 => codex_7_summary, - 30 => codex_30_summary, - days => CostScanner::new(days) - .with_options(codex_scan_options) - .scan_codex(), - }; - let contract = build_local_spend_contract_from_summary( - "codex", - history_days, - include_opencodex, - hide_native, - settings.hide_personal_info, - selected_summary, - ); + let contract = codex_period_contract; let reporting_day = last_included_reporting_day(&contract); let dashboard_timezone = codexbar::core::local_timezone_name(); UsageSpendSummary { rows, contract, + reporting_period: period.raw(), reporting_day, dashboard_timezone, } @@ -701,13 +764,42 @@ fn total_token_mix(mix: &codexbar::spend_contract::SpendTokenMix) -> Option saw.then_some(total) } -fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { +/// Provider-reported daily costs bucket by UTC day, so the selected period is +/// resolved in UTC here. `None` when no day falls inside the window. +fn period_cost_from_daily( + daily: &[super::bridge::CostDailyPointBridge], + period: CostReportingPeriod, + now: chrono::DateTime, +) -> Option { + let earliest = daily + .iter() + .filter_map(|point| chrono::NaiveDate::parse_from_str(&point.day, "%Y-%m-%d").ok()) + .min(); + let bounds = period.bounds(now, CostTimeZone::UTC, earliest); + let mut total = 0.0; + let mut saw = false; + for point in daily { + if bounds.contains_day_key(&point.day) { + total += point.amount; + saw = true; + } + } + saw.then_some(total) +} + +fn cached_spend( + snapshot: Option<&ProviderUsageSnapshot>, + reporting_period: CostReportingPeriod, + now: chrono::DateTime, +) -> SpendValues { let Some(snapshot) = snapshot else { return SpendValues { seven_day: None, thirty_day: None, seven_day_tokens: None, thirty_day_tokens: None, + period_cost: None, + period_tokens: None, source: "unavailable".to_string(), refreshing: false, stale_updated_at: None, @@ -719,6 +811,8 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { thirty_day: None, seven_day_tokens: None, thirty_day_tokens: None, + period_cost: None, + period_tokens: None, source: if snapshot.error.is_some() { "unavailable".to_string() } else { @@ -759,11 +853,21 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { } (saw_seven.then_some(seven), saw_thirty.then_some(thirty)) }; + let period_cost = if cost.daily.is_empty() { + // A provider-reported 30-day total only answers a 30-day selection. + (reporting_period == CostReportingPeriod::Rolling(30)) + .then_some(thirty_day) + .flatten() + } else { + period_cost_from_daily(&cost.daily, reporting_period, now) + }; SpendValues { seven_day, thirty_day, seven_day_tokens: None, thirty_day_tokens: None, + period_cost, + period_tokens: None, source: if period.is_empty() { snapshot.source_label.clone() } else { @@ -775,63 +879,4 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { } #[cfg(test)] -mod cache_key_tests { - use super::*; - - #[test] - fn invalidated_owner_clears_orphaned_indexing_activity() { - let mut coordinator = UsageSpendCoordinator::default(); - let owner = coordinator.begin("account:old".to_string()); - - assert!(coordinator.clear_if_indexing(&owner)); - assert!(!coordinator.is_current(&owner)); - } - - #[test] - fn old_owner_cleanup_cannot_clear_a_replacement() { - let mut coordinator = UsageSpendCoordinator::default(); - let old = coordinator.begin("account:old".to_string()); - let replacement = coordinator.begin("account:new".to_string()); - - assert!(!coordinator.clear_if_indexing(&old)); - assert!(coordinator.is_current(&replacement)); - } - - #[test] - fn settings_replacement_preserves_an_intentional_pause() { - let mut coordinator = UsageSpendCoordinator::default(); - let old = coordinator.begin("settings:old".to_string()); - let replacement = coordinator.begin("settings:new".to_string()); - let status = codexbar::core::CachedCostReadStatus { - codex_scan_pause_reason: Some(codexbar::core::CodexScanPauseReason::NoProgress), - ..Default::default() - }; - mark_refresh_paused_if_codex_scan_paused( - &mut coordinator, - &replacement, - true, - status.codex_scan_pause_reason.as_ref(), - ); - - assert!(!coordinator.clear_if_indexing(&old)); - assert_eq!( - coordinator.current.as_ref().map(|(_, phase)| *phase), - Some(UsageSpendRefreshPhase::Paused) - ); - } - - #[test] - fn privacy_mode_is_part_of_usage_spend_cache_identity() { - let public = usage_spend_cache_key_with_privacy(&[], 30, false, false, false); - 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)); - } -} +mod tests; 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 new file mode 100644 index 0000000000..20b39eadf2 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend/tests.rs @@ -0,0 +1,116 @@ +use super::*; + +#[test] +fn invalidated_owner_clears_orphaned_indexing_activity() { + let mut coordinator = UsageSpendCoordinator::default(); + let owner = coordinator.begin("account:old".to_string()); + + assert!(coordinator.clear_if_indexing(&owner)); + assert!(!coordinator.is_current(&owner)); +} + +#[test] +fn old_owner_cleanup_cannot_clear_a_replacement() { + let mut coordinator = UsageSpendCoordinator::default(); + let old = coordinator.begin("account:old".to_string()); + let replacement = coordinator.begin("account:new".to_string()); + + assert!(!coordinator.clear_if_indexing(&old)); + assert!(coordinator.is_current(&replacement)); +} + +#[test] +fn settings_replacement_preserves_an_intentional_pause() { + let mut coordinator = UsageSpendCoordinator::default(); + let old = coordinator.begin("settings:old".to_string()); + let replacement = coordinator.begin("settings:new".to_string()); + let status = codexbar::core::CachedCostReadStatus { + codex_scan_pause_reason: Some(codexbar::core::CodexScanPauseReason::NoProgress), + ..Default::default() + }; + mark_refresh_paused_if_codex_scan_paused( + &mut coordinator, + &replacement, + true, + status.codex_scan_pause_reason.as_ref(), + ); + + assert!(!coordinator.clear_if_indexing(&old)); + assert_eq!( + coordinator.current.as_ref().map(|(_, phase)| *phase), + Some(UsageSpendRefreshPhase::Paused) + ); +} + +#[test] +fn privacy_mode_is_part_of_usage_spend_cache_identity() { + let public = usage_spend_cache_key_with_privacy(&[], "rolling:30", false, false, false); + let private = usage_spend_cache_key_with_privacy(&[], "rolling:30", false, false, true); + assert_ne!(public, private); +} + +fn identity(period: CostReportingPeriod, now: &str) -> String { + let now = chrono::DateTime::parse_from_rfc3339(now) + .expect("valid instant") + .with_timezone(&Utc); + period.identity(now, CostTimeZone::UTC) +} + +#[test] +fn selected_period_is_part_of_usage_spend_cache_identity() { + let now = "2026-09-15T12:00:00Z"; + let month = identity(CostReportingPeriod::MonthToDate, now); + let rolling = identity(CostReportingPeriod::Rolling(30), now); + let all = identity(CostReportingPeriod::AllAvailable, now); + assert_ne!(month, rolling); + assert_ne!(rolling, all); + assert_ne!( + usage_spend_cache_key_with_privacy(&[], &month, false, false, false), + usage_spend_cache_key_with_privacy(&[], &rolling, false, false, false) + ); +} + +#[test] +fn month_to_date_identity_changes_across_a_month_boundary() { + let last_day = identity(CostReportingPeriod::MonthToDate, "2026-09-30T23:00:00Z"); + let first_day = identity(CostReportingPeriod::MonthToDate, "2026-10-01T01:00:00Z"); + assert_ne!(last_day, first_day); + assert_ne!( + usage_spend_cache_key_with_privacy(&[], &last_day, false, false, false), + usage_spend_cache_key_with_privacy(&[], &first_day, false, false, false) + ); +} + +#[test] +fn daily_period_cost_follows_the_selected_window() { + let daily = ["2026-09-01", "2026-09-10", "2026-08-31"] + .into_iter() + .map(|day| super::super::bridge::CostDailyPointBridge { + day: day.to_string(), + amount: 1.0, + }) + .collect::>(); + let now = chrono::DateTime::parse_from_rfc3339("2026-09-15T12:00:00Z") + .expect("valid instant") + .with_timezone(&Utc); + assert_eq!( + period_cost_from_daily(&daily, CostReportingPeriod::MonthToDate, now), + Some(2.0) + ); + assert_eq!( + period_cost_from_daily(&daily, CostReportingPeriod::AllAvailable, now), + Some(3.0) + ); + assert_eq!( + period_cost_from_daily(&daily, CostReportingPeriod::Rolling(3), now), + None + ); +} + +#[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-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 99012ae7c3..7e4865f8ac 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -33,8 +33,11 @@ pub struct PowerToysProviderSnapshot { secondary_label: Option, secondary: Option, today_cost: Option, + /// Published pipe contract: always the trailing 30 days, independent of + /// the History window setting (`cost_reporting_period`). thirty_day_cost: Option, latest_tokens: Option, + /// Published pipe contract: always the trailing 30 days. thirty_day_tokens: Option, top_model: Option, updated_at: String, @@ -227,6 +230,10 @@ mod tests { today_cost: Some(1.25), thirty_day_cost: Some(12.5), thirty_day_tokens: Some(42_000), + // A different History window must not change the pipe values. + period_cost: Some(99.0), + period_tokens: Some(99_000), + reporting_period: "month-to-date".to_string(), latest_tokens: Some(1_200), top_model: Some("gpt-5".to_string()), estimate_note: "cached".to_string(), diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index c8ec795546..ffd2f94677 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -129,6 +129,9 @@ describe("MenuCard", () => { PanelReserveSuffix: "in reserve", PanelThirtyDayCost: "30d cost", PanelThirtyDayTokens: "30d tokens", + PanelPeriodCost: "{} cost", + PanelPeriodTokens: "{} tokens", + CostPeriodShortMonthToDate: "MTD", PanelTodayBudget: "today", PanelUsedSuffix: "used", ResetsInHoursMinutes: "Resets in {}h {}m", @@ -601,6 +604,37 @@ describe("MenuCard", () => { expect(details.open).toBe(true); }); + it("labels local totals with the selected History window instead of 30 days", async () => { + tauriMocks.getProviderChartData.mockResolvedValue({ + providerId: "claude", + costHistory: [{ date: "2026-05-24", value: 1.23 }], + creditsHistory: [], + usageBreakdown: [], + localUsage: { + todayCost: null, + // The fixed 30-day fields must not leak into the period rows. + thirtyDayCost: 1.23, + thirtyDayTokens: 584_000, + periodCost: 7.5, + periodTokens: 2_000_000, + reportingPeriod: "month-to-date", + latestTokens: null, + topModel: "glim-4.6", + estimateNote: "Estimated from local logs", + tokenCostUpdatedAtMs: 1234, + }, + }); + + renderCard(provider(null)); + + expect(await screen.findByText("MTD cost")).toBeInTheDocument(); + expect(screen.getByText("MTD tokens")).toBeInTheDocument(); + expect(screen.getByText("$7.50")).toBeInTheDocument(); + expect(screen.getByText("2M")).toBeInTheDocument(); + expect(screen.queryByText("30d cost")).not.toBeInTheDocument(); + expect(screen.queryByText("584K")).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..930a04a0f9 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -18,6 +18,7 @@ import { type ResetTimeFormatMode, } from "../hooks/useFormattedResetTime"; import { formatEta } from "../lib/formatEta"; +import { periodCostLabel, periodTokensLabel } from "../lib/costPeriod"; import type { LocaleKey } from "../i18n/keys"; import { paceCategory } from "../surfaces/tray/paceCategory"; import { SimpleBarChart, StackedBarChart } from "./MiniBarChart"; @@ -114,6 +115,11 @@ function LocalUsageBlock({ const { t } = useLocale(); 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 visibleHistory = costHistory.slice(-30); const maxCost = Math.max( ...visibleHistory.flatMap((point) => (point.value == null ? [] : [point.value])), @@ -137,17 +143,19 @@ function LocalUsageBlock({ {!isMuse && (
- {t("PanelThirtyDayCost")} + + {period ? periodCostLabel(period, t) : t("PanelThirtyDayCost")} + - {summary.thirtyDayCost != null - ? formatCurrency(summary.thirtyDayCost, "USD") - : "—"} + {periodCost != null ? formatCurrency(periodCost, "USD") : "—"}
)}
- {t("PanelThirtyDayTokens")} - {formatCompactCount(summary.thirtyDayTokens)} + + {period ? periodTokensLabel(period, t) : t("PanelThirtyDayTokens")} + + {formatCompactCount(periodTokens ?? null)}
{!isMuse && (
diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index 88860e087d..d7b07da417 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -231,6 +231,8 @@ describe("FloatBar", () => { ProviderIssueLocalRuntimeOffline: "Local runtime offline", ProviderIssueUnknown: "Usage unavailable", FloatBarThirtyDayShort: "30d", + CostPeriodShortMonthToDate: "MTD", + CostPeriodShortDays: "{}d", FloatBarNoProviders: "No providers", FloatBarRemainingSuffix: "remaining", }), @@ -427,6 +429,38 @@ describe("FloatBar", () => { ); }); + it("shows the selected History window cost with its short label", async () => { + tauriMocks.getCachedProviders.mockResolvedValue([snapshot("codex", "Codex", 75)]); + tauriMocks.getSettingsSnapshot.mockResolvedValue( + settings({ floatBarShowCost: true, costReportingPeriod: "month-to-date" }), + ); + tauriMocks.getProviderLocalUsageSummary.mockResolvedValue({ + todayCost: 1.25, + // Fixed 30-day compat field stays different from the selected window. + thirtyDayCost: 12.5, + thirtyDayTokens: 1000, + periodCost: 4.75, + periodTokens: 300, + reportingPeriod: "month-to-date", + latestTokens: 200, + topModel: "gpt-5", + estimateNote: "Estimated from local logs", + tokenCostUpdatedAtMs: 1234, + }); + + const { container } = renderFloatBar( + bootstrap({ floatBarShowCost: true, costReportingPeriod: "month-to-date" }), + ); + + await waitFor(() => { + const items = Array.from(container.querySelectorAll(".floatbar__cost-item")).map( + (item) => item.textContent, + ); + expect(items).toEqual(["Today$1.25", "MTD$4.75"]); + }); + expect(container.textContent).not.toContain("$12.50"); + }); + it("does not scan local costs by default", async () => { tauriMocks.getCachedProviders.mockResolvedValue([ snapshot("codex", "Codex", 75), diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.tsx index e8d71af8fe..d067f49cb4 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.tsx @@ -19,6 +19,7 @@ import { } from "../lib/tauri"; import { ProviderIcon } from "../components/providers/ProviderIcon"; import { getProviderIcon } from "../components/providers/providerIcons"; +import { costPeriodShortLabel } from "../lib/costPeriod"; import { describeProviderState } from "../lib/providerState"; import type { BootstrapState, @@ -92,7 +93,10 @@ type FloatBarCostSummary = { providerId: string; displayName: string; todayCost: number | null; - thirtyDayCost: number | null; + /** Cost over the selected History window (30 days when the backend sends none). */ + periodCost: number | null; + /** Raw History window, or null for payloads that only carry 30 days. */ + period: string | null; }; type FloatBarCostTarget = { @@ -105,8 +109,13 @@ 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?.todayCost != null || summary?.thirtyDayCost != null; + return summary != null && (summary.todayCost != null || windowCost(summary) != null); } function formatUsd(value: number | null): string | null { @@ -118,22 +127,22 @@ function CostPill({ summary, scale, todayLabel, - thirtyDayLabel, + periodLabel, estimateLabel, }: { summary: FloatBarCostSummary; scale: number; todayLabel: string; - thirtyDayLabel: string; + periodLabel: string; estimateLabel: string; }) { const today = formatUsd(summary.todayCost); - const thirtyDay = formatUsd(summary.thirtyDayCost); + const periodCost = formatUsd(summary.periodCost); const iconSize = Math.round(10 * scale); const brand = getProviderIcon(summary.providerId).brandColor; const title = [ today ? `${todayLabel} ${today}` : null, - thirtyDay ? `${thirtyDayLabel} ${thirtyDay}` : null, + periodCost ? `${periodLabel} ${periodCost}` : null, ] .filter(Boolean) .join(" / "); @@ -159,13 +168,13 @@ function CostPill({ )} - {thirtyDay && ( + {periodCost && ( - {thirtyDayLabel} + {periodLabel} - {thirtyDay} + {periodCost} )} @@ -382,7 +391,8 @@ export default function FloatBar({ state }: { state: BootstrapState }) { providerId: target.providerId, displayName: target.displayName, todayCost: localUsage.todayCost, - thirtyDayCost: localUsage.thirtyDayCost, + periodCost: windowCost(localUsage), + period: localUsage.reportingPeriod || null, } satisfies FloatBarCostSummary; }), ) @@ -403,13 +413,14 @@ export default function FloatBar({ state }: { state: BootstrapState }) { return () => { cancelled = true; }; - }, [visibleCostTargets]); + // A History window change re-reads the local usage summaries. + }, [visibleCostTargets, settings.costReportingPeriod]); const visibleCosts = visible .map((provider) => localCosts[providerCostKey(provider)]) .filter((summary): summary is FloatBarCostSummary => Boolean(summary)); const visibleCostValuesKey = visibleCosts - .map((summary) => `${summary.key}:${summary.todayCost ?? ""}:${summary.thirtyDayCost ?? ""}`) + .map((summary) => `${summary.key}:${summary.todayCost ?? ""}:${summary.period ?? ""}:${summary.periodCost ?? ""}`) .join("|"); // Keep the native floatbar window fitted when late data/fonts/icons change layout. const lastResizeRef = useRef<{ w: number; h: number } | null>(null); @@ -522,7 +533,11 @@ export default function FloatBar({ state }: { state: BootstrapState }) { summary={summary} scale={scale} todayLabel={t("PanelToday")} - thirtyDayLabel={t("FloatBarThirtyDayShort")} + periodLabel={ + summary.period + ? costPeriodShortLabel(summary.period, t) + : t("FloatBarThirtyDayShort") + } estimateLabel={t("OverviewSpendEstimate")} /> ))} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 195620e3bf..1b6feb649a 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -914,6 +914,24 @@ export const ALL_LOCALE_KEYS = [ "ProviderAutoResumeTitle", "ProviderAutoResumeAfterQuotaReset", "ProviderAutoResumeAfterQuotaResetHelper", + + // History window for local cost surfaces (upstream 0.67.0 reporting periods) + "CostPeriodHistoryWindow", + "CostPeriodHistoryWindowHelper", + "CostPeriodMonthToDate", + "CostPeriodAll", + "CostPeriodToday", + "CostPeriodLastDays", + "CostPeriodCustom", + "CostPeriodCustomDays", + "CostPeriodShortMonthToDate", + "CostPeriodShortAll", + "CostPeriodShortDays", + "PanelPeriodCost", + "PanelPeriodTokens", + "OverviewSpendPeriodTitle", + "UsageSpendModelsPeriodCaption", + "UsageSpendProjectsPeriodCaption", ] as const; export type LocaleKey = (typeof ALL_LOCALE_KEYS)[number]; diff --git a/apps/desktop-tauri/src/lib/costPeriod.test.ts b/apps/desktop-tauri/src/lib/costPeriod.test.ts new file mode 100644 index 0000000000..2064aeb10a --- /dev/null +++ b/apps/desktop-tauri/src/lib/costPeriod.test.ts @@ -0,0 +1,122 @@ +import { describe, expect, it } from "vitest"; +import type { LocaleKey } from "../i18n/keys"; +import { + DEFAULT_COST_PERIOD, + PRESET_COST_PERIODS, + costPeriodLabel, + costPeriodShortLabel, + customPeriodRaw, + isPresetCostPeriod, + normalizeCostPeriod, + parseCostPeriod, + periodCostLabel, + periodFromLegacyDays, + periodTokensLabel, + rollingDays, +} from "./costPeriod"; + +const EN: Partial> = { + CostPeriodMonthToDate: "Month to date", + CostPeriodAll: "All available history", + CostPeriodToday: "Today", + CostPeriodLastDays: "Last {} days", + CostPeriodShortMonthToDate: "MTD", + CostPeriodShortAll: "All", + CostPeriodShortDays: "{}d", + PanelPeriodCost: "{} cost", + PanelPeriodTokens: "{} tokens", +}; +const t = (key: LocaleKey) => EN[key] ?? key; + +describe("parseCostPeriod", () => { + it("reads the persisted forms", () => { + expect(parseCostPeriod("month-to-date")).toEqual({ kind: "monthToDate" }); + expect(parseCostPeriod("all")).toEqual({ kind: "all" }); + expect(parseCostPeriod("rolling:7")).toEqual({ kind: "rolling", days: 7 }); + expect(parseCostPeriod(" rolling:90 ")).toEqual({ kind: "rolling", days: 90 }); + }); + + it("clamps counts above 365 and rejects zero and unreadable values", () => { + expect(parseCostPeriod("rolling:4000")).toEqual({ kind: "rolling", days: 365 }); + expect(parseCostPeriod("rolling:0")).toBeNull(); + expect(parseCostPeriod("rolling:-3")).toBeNull(); + expect(parseCostPeriod("rolling:seven")).toBeNull(); + expect(parseCostPeriod("weekly")).toBeNull(); + expect(parseCostPeriod("")).toBeNull(); + expect(parseCostPeriod(undefined)).toBeNull(); + }); +}); + +describe("normalizeCostPeriod", () => { + it("keeps valid periods and falls back to 30 days otherwise", () => { + expect(normalizeCostPeriod("month-to-date")).toBe("month-to-date"); + expect(normalizeCostPeriod("rolling:9999")).toBe("rolling:365"); + expect(normalizeCostPeriod("rolling:0")).toBe(DEFAULT_COST_PERIOD); + expect(normalizeCostPeriod(undefined)).toBe("rolling:30"); + }); +}); + +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"); + expect(customPeriodRaw(" 14 ")).toBe("rolling:14"); + expect(customPeriodRaw("365")).toBe("rolling:365"); + expect(customPeriodRaw("0")).toBeNull(); + expect(customPeriodRaw("366")).toBeNull(); + expect(customPeriodRaw("1.5")).toBeNull(); + expect(customPeriodRaw("-2")).toBeNull(); + expect(customPeriodRaw("")).toBeNull(); + expect(customPeriodRaw("abc")).toBeNull(); + }); +}); + +describe("presets", () => { + it("lists month to date, 1/7/30/90/365 days, and all", () => { + expect([...PRESET_COST_PERIODS]).toEqual([ + "month-to-date", + "rolling:1", + "rolling:7", + "rolling:30", + "rolling:90", + "rolling:365", + "all", + ]); + expect(isPresetCostPeriod("rolling:7")).toBe(true); + expect(isPresetCostPeriod("rolling:14")).toBe(false); + expect(rollingDays("rolling:14")).toBe(14); + expect(rollingDays("all")).toBeNull(); + }); +}); + +describe("labels", () => { + it("builds long labels", () => { + expect(costPeriodLabel("month-to-date", t)).toBe("Month to date"); + expect(costPeriodLabel("all", t)).toBe("All available history"); + expect(costPeriodLabel("rolling:1", t)).toBe("Today"); + expect(costPeriodLabel("rolling:14", t)).toBe("Last 14 days"); + expect(costPeriodLabel(undefined, t)).toBe("Last 30 days"); + }); + + it("builds short and menu labels", () => { + expect(costPeriodShortLabel("month-to-date", t)).toBe("MTD"); + expect(costPeriodShortLabel("all", t)).toBe("All"); + expect(costPeriodShortLabel("rolling:7", t)).toBe("7d"); + expect(periodCostLabel("month-to-date", t)).toBe("MTD cost"); + expect(periodTokensLabel("rolling:90", t)).toBe("90d tokens"); + }); +}); diff --git a/apps/desktop-tauri/src/lib/costPeriod.ts b/apps/desktop-tauri/src/lib/costPeriod.ts new file mode 100644 index 0000000000..c79b7a970e --- /dev/null +++ b/apps/desktop-tauri/src/lib/costPeriod.ts @@ -0,0 +1,125 @@ +import type { LocaleKey } from "../i18n/keys"; + +/** + * History window ("cost reporting period") helpers. The persisted form mirrors + * `codexbar::cost_reporting_period::CostReportingPeriod::raw`: + * `rolling:N` (1..=365), `month-to-date`, or `all`. + */ + +export const DEFAULT_COST_PERIOD = "rolling:30"; +export const MAX_ROLLING_DAYS = 365; + +export type ParsedCostPeriod = + | { kind: "rolling"; days: number } + | { kind: "monthToDate" } + | { kind: "all" }; + +type Translate = (key: LocaleKey) => string; + +/** Presets offered by the picker, in display order. Custom day counts are separate. */ +export const PRESET_COST_PERIODS = [ + "month-to-date", + "rolling:1", + "rolling:7", + "rolling:30", + "rolling:90", + "rolling:365", + "all", +] as const; + +/** Parse a persisted period. Unknown forms and `rolling:0` return null; counts above 365 clamp. */ +export function parseCostPeriod(raw: string | null | undefined): ParsedCostPeriod | null { + const value = raw?.trim(); + if (!value) return null; + if (value === "month-to-date") return { kind: "monthToDate" }; + if (value === "all") return { kind: "all" }; + const match = /^rolling:(\d+)$/.exec(value); + if (!match) return null; + const days = Number.parseInt(match[1], 10); + if (!Number.isSafeInteger(days) || days <= 0) return null; + return { kind: "rolling", days: Math.min(days, MAX_ROLLING_DAYS) }; +} + +function toRaw(parsed: ParsedCostPeriod): string { + switch (parsed.kind) { + case "monthToDate": + return "month-to-date"; + case "all": + return "all"; + case "rolling": + return `rolling:${parsed.days}`; + } +} + +/** Canonical persisted form, or the 30-day default when the value is missing or unreadable. */ +export function normalizeCostPeriod(raw: string | null | undefined): string { + const parsed = parseCostPeriod(raw); + 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(); + if (!/^\d+$/.test(text)) return null; + const days = Number.parseInt(text, 10); + if (!Number.isSafeInteger(days) || days < 1 || days > MAX_ROLLING_DAYS) return null; + return `rolling:${days}`; +} + +export function isPresetCostPeriod(raw: string): boolean { + return (PRESET_COST_PERIODS as readonly string[]).includes(raw); +} + +/** Day count of a rolling period, or null for month-to-date / all. */ +export function rollingDays(raw: string | null | undefined): number | null { + const parsed = parseCostPeriod(raw); + return parsed?.kind === "rolling" ? parsed.days : null; +} + +function fill(template: string, value: string): string { + return template.replace("{}", value); +} + +/** Long label, for headings and captions ("Month to date", "Last 7 days"). */ +export function costPeriodLabel(raw: string | null | undefined, t: Translate): string { + const parsed = parseCostPeriod(raw ?? DEFAULT_COST_PERIOD) ?? { kind: "rolling", days: 30 }; + switch (parsed.kind) { + case "monthToDate": + return t("CostPeriodMonthToDate"); + case "all": + return t("CostPeriodAll"); + case "rolling": + return parsed.days === 1 + ? t("CostPeriodToday") + : fill(t("CostPeriodLastDays"), String(parsed.days)); + } +} + +/** Compact label, for menu rows and the float bar ("MTD", "7d", "All"). */ +export function costPeriodShortLabel(raw: string | null | undefined, t: Translate): string { + const parsed = parseCostPeriod(raw ?? DEFAULT_COST_PERIOD) ?? { kind: "rolling", days: 30 }; + switch (parsed.kind) { + case "monthToDate": + return t("CostPeriodShortMonthToDate"); + case "all": + return t("CostPeriodShortAll"); + case "rolling": + return fill(t("CostPeriodShortDays"), String(parsed.days)); + } +} + +/** Menu row label for the selected window ("MTD cost"). */ +export function periodCostLabel(raw: string | null | undefined, t: Translate): string { + return fill(t("PanelPeriodCost"), costPeriodShortLabel(raw, t)); +} + +export function periodTokensLabel(raw: string | null | undefined, t: Translate): string { + return fill(t("PanelPeriodTokens"), costPeriodShortLabel(raw, t)); +} diff --git a/apps/desktop-tauri/src/lib/tauri.ts b/apps/desktop-tauri/src/lib/tauri.ts index 6ccd63c0de..4e1956fada 100644 --- a/apps/desktop-tauri/src/lib/tauri.ts +++ b/apps/desktop-tauri/src/lib/tauri.ts @@ -284,9 +284,10 @@ export function getProviderLocalUsageSummary( return invoke("get_provider_local_usage_summary", { providerId }); } -export function getUsageSpendSummary(options?: { historyDays?: number; forceRefresh?: boolean }): Promise { +/** `period` is a raw reporting period; omitted means the saved `costReportingPeriod`. */ +export function getUsageSpendSummary(options?: { period?: string; forceRefresh?: boolean }): Promise { return invoke("get_usage_spend_summary", { - historyDays: options?.historyDays ?? null, + period: options?.period ?? null, forceRefresh: options?.forceRefresh ?? null, }); } @@ -297,11 +298,11 @@ export function writeUsageSpendExport(path: string, payload: string): Promise { return invoke("get_spend_contract", { providerId, - historyDays: options?.historyDays ?? null, + period: options?.period ?? null, includeOpenCodex: options?.includeOpenCodex ?? null, }); } diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index bf693ce7b8..9f838100c2 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -4,6 +4,7 @@ import { formatUsageSpendReportingDay, filterUsageSpendSummaryForOverview, renderUsageSpendSharePng, + usageSpendShareColumns, usageSpendShareFooter, usageSpendSubscriptionCaption, } from "./usageSpendSharing"; @@ -86,7 +87,15 @@ describe("usage spend sharing", () => { // account emails or provider-identity secrets. The renderer draws exactly // these cells; keep the invariant comment on renderUsageSpendSharePng in // sync if this list ever grows. - const renderedCells = ["displayName", "sevenDay", "thirtyDay", "currency", "source"] as const; + const renderedCells = [ + "displayName", + "periodCost", + "periodTokens", + "sevenDay", + "thirtyDay", + "currency", + "source", + ] as const; const row: UsageSpendRow = { providerId: "codex", displayName: "Codex", @@ -104,4 +113,36 @@ describe("usage spend sharing", () => { expect(Object.keys(row).filter((key) => /email|org|token|account/i.test(key))).toEqual([]); expect(unsafeKeys).toEqual(["providerId", "includedInOverview"]); }); + + it("adds a labelled History window column while keeping the fixed 7 and 30 day columns", () => { + const row: UsageSpendRow = { + providerId: "codex", + displayName: "Codex", + sevenDay: 1, + thirtyDay: 2, + periodCost: 9, + periodTokens: 900, + currency: "USD", + 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", + "Month to date", + "7 days", + "30 days", + "Currency", + "Source", + ]); + const cells = withPeriod.cellsFor(row); + expect(cells).toHaveLength(withPeriod.headers.length); + expect(cells).toHaveLength(withPeriod.colW.length); + expect(cells[1]).toContain("900"); + expect(cells[2]).not.toContain("900"); + }); }); diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 2246e98505..5d80e3e4d9 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -139,6 +139,40 @@ export function formatSpendMetric( return parts.length > 0 ? parts.join(" · ") : "—"; } +/** + * 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. + */ +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], + cellsFor: (row) => [ + row.displayName, + ...(withPeriod + ? [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", + row.source, + ], + }; +} + /** * Render the sanitized share-card PNG. * @@ -148,12 +182,16 @@ 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, + periodLabel?: string, +): string { const rows = summary.rows; const pad = 24; const rowH = 28; const headerH = 48; - const colW = [160, 100, 100, 80, 160]; + const { headers, cellsFor, colW } = usageSpendShareColumns(periodLabel); const width = pad * 2 + colW.reduce((a, b) => a + b, 0); const height = pad * 2 + headerH + Math.max(1, rows.length) * rowH + 52; const canvas = document.createElement("canvas"); @@ -177,7 +215,6 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri ctx.font = "12px system-ui,Segoe UI,sans-serif"; ctx.fillText("Win-CodexBar · local estimates · no account emails", pad, pad + 36); - const headers = ["Provider", "7 days", "30 days", "Currency", "Source"]; let x = pad; const y0 = pad + headerH; ctx.fillStyle = "#9fb0c8"; @@ -200,13 +237,7 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri } else { rows.forEach((row, index) => { 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", - row.source, - ]; + const cells = cellsFor(row); let cellX = pad; cells.forEach((cell, cellIndex) => { ctx.fillStyle = cellIndex === 0 ? "#e7ecf3" : "#c5d0e0"; @@ -250,10 +281,11 @@ export function shareUsageSpendPng( summary: UsageSpendSummary | null, title: string, filename: string, + periodLabel?: string, ): string | null { if (!summary) return "UsageSpendShareEmpty"; try { - const dataUrl = renderUsageSpendSharePng(summary, title); + const dataUrl = renderUsageSpendSharePng(summary, title, periodLabel); if (!dataUrl) return "UsageSpendShareFailed"; downloadPng(dataUrl, filename); return null; diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.period.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.period.test.tsx new file mode 100644 index 0000000000..0f6263c089 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.period.test.tsx @@ -0,0 +1,193 @@ +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const tauriMocks = vi.hoisted(() => ({ + getCachedProviders: vi.fn(), + refreshProviders: vi.fn(), + refreshProvidersIfStale: vi.fn(), + getSettingsSnapshot: vi.fn(), + updateSettings: vi.fn(), + getUpdateState: vi.fn(), + checkForUpdates: vi.fn(), + downloadUpdate: vi.fn(), + applyUpdate: vi.fn(), + dismissUpdate: vi.fn(), + openReleasePage: vi.fn(), + setSurfaceMode: vi.fn(), + dismissTrayPanel: vi.fn(), + beginFlyoutGesture: vi.fn().mockResolvedValue(undefined), + endFlyoutGesture: vi.fn().mockResolvedValue(undefined), + openSettingsWindow: vi.fn(), + quitApp: vi.fn(), + getWorkAreaRect: vi.fn(), + reanchorTrayPanel: vi.fn(), + revealTrayPanelWindow: vi.fn(), + flyoutStoredSize: vi.fn().mockResolvedValue(null), + setFlyoutSize: vi.fn().mockResolvedValue(undefined), + openProviderDashboard: vi.fn(), + openProviderStatusPage: vi.fn(), + getProviderChartData: vi.fn(), + getCurrentSurfaceState: vi.fn(), + getLocaleStrings: vi.fn(), + setUiLanguage: vi.fn(), + getDeepSeekPricingStatus: vi.fn().mockResolvedValue(null), + getUsageSpendSummary: vi.fn(), + claudeReconciliationState: vi.fn().mockResolvedValue(null), +})); + +const eventMocks = vi.hoisted(() => ({ + listen: vi.fn(), +})); + +const windowMocks = vi.hoisted(() => ({ + getCurrentWindow: vi.fn(() => ({ + setSize: vi.fn().mockResolvedValue(undefined), + close: vi.fn().mockResolvedValue(undefined), + scaleFactor: vi.fn().mockResolvedValue(1), + onResized: vi.fn().mockResolvedValue(() => {}), + innerSize: vi.fn().mockResolvedValue({ width: 328, height: 200 }), + })), + LogicalSize: vi.fn((width: number, height: number) => ({ width, height })), + PhysicalSize: vi.fn((width: number, height: number) => ({ width, height })), +})); + +vi.mock("../lib/tauri", () => tauriMocks); +vi.mock("@tauri-apps/api/event", () => eventMocks); +vi.mock("@tauri-apps/api/window", () => windowMocks); + +import TrayPanel from "./TrayPanel"; +import { LocaleProvider } from "../i18n/LocaleProvider"; +import { buildBundle } from "../test/localeHarness"; +import type { ProviderUsageSnapshot, SettingsSnapshot } from "../types/bridge"; + +function rateWindow(used: number) { + return { + usedPercent: used, + remainingPercent: 100 - used, + windowMinutes: null, + resetsAt: null, + resetDescription: null, + isExhausted: false, + reservePercent: null, + reserveDescription: null, + }; +} + +const codex: ProviderUsageSnapshot = { + providerId: "codex", + displayName: "Codex", + primary: rateWindow(35), + selectedMetric: rateWindow(35), + primaryLabel: "Monthly", + secondary: null, + modelSpecific: null, + tertiary: null, + extraRateWindows: [], + cost: null, + planName: null, + accountEmail: null, + sourceLabel: "auto", + updatedAt: "2026-05-24T00:00:00Z", + error: null, + errorState: "ready", + pace: null, + accountOrganization: null, + trayStatusLabel: null, + fetchDurationMs: null, +}; + +const settings = { + enabledProviders: ["codex"], + refreshIntervalSecs: 300, + adaptiveRefresh: false, + refreshAllProvidersOnMenuOpen: false, + lowPowerMode: false, + highUsageThreshold: 70, + criticalUsageThreshold: 90, + trayIconMode: "single", + switcherShowsIcons: true, + showAsUsed: true, + resetTimeRelative: true, + overviewLayout: "detailed", + hidePersonalInfo: false, + uiLanguage: "english", + theme: "dark", + windowScalePercent: 125, + trayScalePercent: 100, + trayPanelAlwaysOnTop: false, + providerMetrics: {}, + costSummaryDisplayStyle: "compact", + costReportingPeriod: "rolling:7", + providerAccentColors: {}, +} as unknown as SettingsSnapshot; + +describe("TrayPanel History window", () => { + beforeEach(() => { + vi.clearAllMocks(); + tauriMocks.getCachedProviders.mockResolvedValue([codex]); + tauriMocks.getSettingsSnapshot.mockResolvedValue(settings); + tauriMocks.updateSettings.mockResolvedValue(settings); + tauriMocks.refreshProviders.mockResolvedValue(undefined); + tauriMocks.refreshProvidersIfStale.mockResolvedValue(undefined); + tauriMocks.dismissTrayPanel.mockResolvedValue(undefined); + tauriMocks.reanchorTrayPanel.mockResolvedValue(undefined); + tauriMocks.getWorkAreaRect.mockResolvedValue({ x: 0, y: 0, width: 1440, height: 900 }); + tauriMocks.getCurrentSurfaceState.mockResolvedValue({ + mode: "trayPanel", + target: { kind: "summary" }, + }); + tauriMocks.getUpdateState.mockResolvedValue({ + status: "idle", + version: null, + error: null, + progress: null, + releaseUrl: null, + canDownload: false, + canApply: false, + lastCheckedAt: null, + }); + tauriMocks.getProviderChartData.mockResolvedValue({ + providerId: "codex", + costHistory: [], + creditsHistory: [], + usageBreakdown: [], + localUsage: null, + }); + tauriMocks.getLocaleStrings.mockResolvedValue(buildBundle({})); + eventMocks.listen.mockResolvedValue(() => {}); + }); + + it("totals the Overview spend over the selected History window", async () => { + tauriMocks.getUsageSpendSummary.mockResolvedValue({ + contract: {}, + reportingPeriod: "rolling:7", + reportingDay: "2026-09-29", + dashboardTimezone: "UTC", + rows: [ + { + providerId: "codex", + displayName: "Codex", + sevenDay: 9, + thirtyDay: 2, + periodCost: 9, + periodTokens: 900, + currency: "USD", + source: "local", + includedInOverview: true, + }, + ], + }); + + render( + + + , + ); + + expect(await screen.findByText("OverviewSpendPeriodTitle")).toBeInTheDocument(); + expect(screen.getByText("$9.00")).toBeInTheDocument(); + expect(screen.queryByText("$2.00")).not.toBeInTheDocument(); + // The scan carries no explicit window: the backend resolves the saved one. + expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledWith(); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..920f5db80b 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -1,7 +1,13 @@ 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 { + BootstrapState, + ProviderUsageSnapshot, + UsageSpendRow, + UsageSpendSummary, +} from "../types/bridge"; import type { LocaleKey } from "../i18n/keys"; +import { costPeriodLabel, costPeriodShortLabel } from "../lib/costPeriod"; import { beginFlyoutGesture, getUsageSpendSummary, @@ -206,7 +212,11 @@ export default function TrayPanel({ state }: { state: BootstrapState }) { />
{selectedProviderId === null && ( - provider.providerId)} t={t} /> + provider.providerId)} + period={settings.costReportingPeriod} + t={t} + /> )}
{useWideColumns @@ -336,45 +346,66 @@ function TrayResizeHandles() { ); } -function OverviewSpendSummary({ providerIds, t }: { providerIds: string[]; t: (key: LocaleKey) => string }) { +function OverviewSpendSummary({ + providerIds, + period, + t, +}: { + providerIds: string[]; + /** Saved History window; a change rescans. Omitted means the backend default. */ + period?: string; + t: (key: LocaleKey) => string; +}) { const [summary, setSummary] = useState(null); const [shareError, setShareError] = useState(null); useEffect(() => { let cancelled = false; - void getUsageSpendSummary({ historyDays: 30 }) + // No explicit period: the backend resolves the saved History window. + void getUsageSpendSummary() .then((value) => { if (!cancelled) setSummary(value); }) .catch(() => { if (!cancelled) setSummary(null); }); return () => { cancelled = true; }; - }, [providerIds.join("|")]); + }, [providerIds.join("|"), period]); // Overview consumes the same backend spend catalog as Usage & Spend. Do not // restrict accounting to whichever cards happen to be rendered in this tray. 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"); const onShare = () => { setShareError(null); const error = shareUsageSpendPng( overviewSummary, - t("OverviewSpendTitle"), + title, `codexbar-overview-usage-${overviewSummary.reportingDay}.png`, + summaryPeriod ? costPeriodLabel(summaryPeriod, t) : undefined, ); 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) => row.thirtyDay != null && Number.isFinite(row.thirtyDay)); + const known = summable.filter((row) => { + const spend = spendOf(row); + return spend != null && Number.isFinite(spend); + }); if (known.length === 0) return null; - const total = known.reduce((sum, row) => sum + (row.thirtyDay ?? 0), 0); + const total = known.reduce((sum, row) => sum + (spendOf(row) ?? 0), 0); const partial = known.length < rows.length; const formatter = new Intl.NumberFormat(undefined, { style: "currency", currency: "USD", maximumFractionDigits: 2 }); return (
- {t("OverviewSpendTitle")} + {title} {partial ? "~" : ""}{formatter.format(total)}
diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx new file mode 100644 index 0000000000..6229e9ede0 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/CostPeriodControl.tsx @@ -0,0 +1,107 @@ +import { useEffect, useRef, useState } from "react"; +import type { LocaleKey } from "../../../i18n/keys"; +import { + MAX_ROLLING_DAYS, + PRESET_COST_PERIODS, + costPeriodLabel, + customPeriodRaw, + isPresetCostPeriod, + rollingDays, +} from "../../../lib/costPeriod"; + +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. + */ +export default function CostPeriodControl({ + value, + onChange, + disabled = false, + t, +}: { + value: string; + onChange: (raw: string) => void; + disabled?: boolean; + t: (key: LocaleKey) => string; +}) { + 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 { + setCustomMode(true); + setCustomText(String(rollingDays(value) ?? "")); + } + }, [value]); + + const customInvalid = customMode && customText !== "" && customPeriodRaw(customText) === null; + + return ( +
+ +

{t("CostPeriodHistoryWindowHelper")}

+
+ + {customMode && ( + { + const text = event.target.value; + setCustomText(text); + const raw = customPeriodRaw(text); + if (raw && raw !== value) { + typedRaw.current = raw; + onChange(raw); + } + }} + /> + )} +
+
+ ); +} diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx new file mode 100644 index 0000000000..f4221cf77f --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.test.tsx @@ -0,0 +1,215 @@ +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const tauriMocks = vi.hoisted(() => ({ + getSettingsSnapshot: vi.fn(), + getUsageSpendSummary: vi.fn(), + updateSettings: vi.fn(), + writeUsageSpendExport: vi.fn(), +})); + +vi.mock("../../../lib/tauri", () => tauriMocks); +vi.mock("@tauri-apps/plugin-dialog", () => ({ save: vi.fn() })); +vi.mock("@tauri-apps/api/event", () => ({ + listen: vi.fn().mockResolvedValue(() => {}), +})); + +const EN: Record = { + CostPeriodMonthToDate: "Month to date", + CostPeriodAll: "All available history", + CostPeriodToday: "Today", + CostPeriodLastDays: "Last {} days", + CostPeriodCustom: "Custom days", + UsageSpendModelsPeriodCaption: "Models for: {}", + UsageSpendProjectsPeriodCaption: "Projects for: {}", +}; +vi.mock("../../../hooks/useLocale", () => ({ + useLocale: () => ({ t: (key: string) => EN[key] ?? key }), +})); + +import UsageSpendTab from "./UsageSpendTab"; +import type { TabProps } from "../settingsTabs"; +import type { SpendContract, UsageSpendSummary } from "../../../types/bridge"; + +function contract(reportingPeriod: string): SpendContract { + return { + providerId: "codex", + historyDays: 30, + reportingPeriod, + knownCostUsd: 12, + knownZero: false, + provenance: "listPriceEstimate", + priceCoverage: { priced: 1, unpriced: 0 }, + priceCoverageRatio: 1, + historyCoverageEstablished: true, + tokenMix: {}, + conversationCount: 1, + models: [], + projects: [], + conversations: [], + daily: [], + hourlyActivity: [], + projectSourceStatus: null, + imports: [], + customPricingActive: false, + } as unknown as SpendContract; +} + +function summary(reportingPeriod: string): UsageSpendSummary { + return { + rows: [ + { + providerId: "codex", + displayName: "Codex", + sevenDay: 1, + thirtyDay: 2, + periodCost: 9, + periodTokens: 900, + currency: "USD", + source: "local", + includedInOverview: true, + }, + ], + contract: contract(reportingPeriod), + reportingPeriod, + reportingDay: "2026-09-29", + dashboardTimezone: "UTC", + } as unknown as UsageSpendSummary; +} + +const props = {} as TabProps; + +function lastScanPeriod(): string | undefined { + const calls = tauriMocks.getUsageSpendSummary.mock.calls; + return calls[calls.length - 1]?.[0]?.period; +} + +describe("UsageSpendTab History window", () => { + beforeEach(() => { + vi.clearAllMocks(); + tauriMocks.getSettingsSnapshot.mockResolvedValue({ + costReportingPeriod: "month-to-date", + costSummaryDisplayStyle: "compact", + }); + tauriMocks.updateSettings.mockResolvedValue({}); + tauriMocks.getUsageSpendSummary.mockImplementation( + async (options?: { period?: string }) => summary(options?.period ?? "rolling:30"), + ); + }); + + it("scans once with the saved window and shows it selected", async () => { + render(); + + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledTimes(1)); + expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledWith({ + period: "month-to-date", + forceRefresh: false, + }); + const select = screen.getByLabelText("CostPeriodHistoryWindow") as HTMLSelectElement; + expect(select.value).toBe("month-to-date"); + // The period column is labelled with the window and shows its numbers. + expect(await screen.findByRole("columnheader", { name: "Month to date" })).toBeTruthy(); + expect(screen.getByText(/900/)).toBeTruthy(); + expect(screen.getByText("Models for: Month to date")).toBeTruthy(); + }); + + it("migrates a settings payload without a saved window to 30 days", async () => { + tauriMocks.getSettingsSnapshot.mockResolvedValue({ costSummaryDisplayStyle: "compact" }); + render(); + + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalled()); + expect(lastScanPeriod()).toBe("rolling:30"); + expect((screen.getByLabelText("CostPeriodHistoryWindow") as HTMLSelectElement).value).toBe( + "rolling:30", + ); + }); + + it("falls back to 30 days when the saved window is unreadable", async () => { + tauriMocks.getSettingsSnapshot.mockResolvedValue({ + costReportingPeriod: "rolling:0", + costSummaryDisplayStyle: "compact", + }); + render(); + + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalled()); + expect(lastScanPeriod()).toBe("rolling:30"); + }); + + it("persists a preset choice and rescans with it", async () => { + render(); + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledTimes(1)); + + fireEvent.change(screen.getByLabelText("CostPeriodHistoryWindow"), { + target: { value: "rolling:90" }, + }); + + await waitFor(() => + expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:90" }), + ); + await waitFor(() => expect(lastScanPeriod()).toBe("rolling:90")); + expect(await screen.findByRole("columnheader", { name: "Last 90 days" })).toBeTruthy(); + }); + + it("offers a custom day count and only persists values in 1..=365", 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: "400" } }); + expect(input.getAttribute("aria-invalid")).toBe("true"); + expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); + + fireEvent.change(input, { target: { value: "0" } }); + expect(tauriMocks.updateSettings).not.toHaveBeenCalled(); + + // Typing "1" on the way to "14" keeps the custom input open. + fireEvent.change(input, { target: { value: "1" } }); + await waitFor(() => + expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:1" }), + ); + expect(screen.getByLabelText("CostPeriodCustomDays")).toBe(input); + + fireEvent.change(input, { target: { value: "14" } }); + await waitFor(() => + expect(tauriMocks.updateSettings).toHaveBeenCalledWith({ costReportingPeriod: "rolling:14" }), + ); + 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("starts in custom mode for a saved count that is not a preset", async () => { + tauriMocks.getSettingsSnapshot.mockResolvedValue({ + costReportingPeriod: "rolling:14", + costSummaryDisplayStyle: "compact", + }); + render(); + + const input = (await screen.findByLabelText("CostPeriodCustomDays")) as HTMLInputElement; + expect(input.value).toBe("14"); + expect((screen.getByLabelText("CostPeriodHistoryWindow") as HTMLSelectElement).value).toBe( + "custom", + ); + }); + + it("reverts the picker and shows the error when saving fails", async () => { + tauriMocks.updateSettings.mockRejectedValue(new Error("Invalid cost reporting period: x")); + render(); + await waitFor(() => expect(tauriMocks.getUsageSpendSummary).toHaveBeenCalledTimes(1)); + + fireEvent.change(screen.getByLabelText("CostPeriodHistoryWindow"), { + target: { value: "all" }, + }); + + expect(await screen.findByText("Invalid cost reporting period: x")).toBeTruthy(); + await waitFor(() => + expect((screen.getByLabelText("CostPeriodHistoryWindow") as HTMLSelectElement).value).toBe( + "month-to-date", + ), + ); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index 0f7cc7ee89..7a66a469bf 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -15,12 +15,22 @@ import { } from "../../../lib/usageSpendSharing"; import type { CostSummaryDisplayStyle, SettingsSnapshot, SpendContract, UsageSpendSummary } from "../../../types/bridge"; import type { LocaleKey } from "../../../i18n/keys"; +import { + DEFAULT_COST_PERIOD, + costPeriodLabel, + normalizeCostPeriod, + periodFromLegacyDays, +} from "../../../lib/costPeriod"; import type { TabProps } from "../settingsTabs"; +import CostPeriodControl from "./CostPeriodControl"; export default function UsageSpendTab(_props: TabProps) { const { t } = useLocale(); const [summary, setSummary] = useState(null); - const [selectedDays, setSelectedDays] = useState<0 | 7 | 30>(30); + // Null until the saved History window is known, so the first scan uses it + // instead of scanning the default and then rescanning. + const [period, setPeriod] = useState(null); + const [periodError, setPeriodError] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); const [shareError, setShareError] = useState(null); @@ -36,26 +46,43 @@ export default function UsageSpendTab(_props: TabProps) { .then((settings) => { setIncludeOpenCodex(settings.openCodexUsageLogsEnabled ?? false); setHideNativeCodex(settings.hideNativeCodexCostWhenOpenCodexPresent ?? false); + setPeriod(normalizeCostPeriod(settings.costReportingPeriod)); }) .catch(() => { - // Keep the safe default (off) if settings cannot be loaded. + // Keep the safe defaults (off, 30 days) if settings cannot be loaded. + setPeriod(DEFAULT_COST_PERIOD); }); }, []); const load = useCallback(async (forceRefresh = false, silent = false) => { + if (period === null) return; if (!silent) { setLoading(true); setError(null); } try { - const data = await getUsageSpendSummary({ historyDays: selectedDays, forceRefresh }); + const data = await getUsageSpendSummary({ period, forceRefresh }); setSummary(data); } catch (err: unknown) { if (!silent) setError(err instanceof Error ? err.message : String(err)); } finally { if (!silent) setLoading(false); } - }, [hideNativeCodex, includeOpenCodex, selectedDays]); + }, [hideNativeCodex, includeOpenCodex, period]); + + const onPeriodChange = useCallback( + (next: string) => { + const previous = period; + setPeriodError(null); + setPeriod(next); + void updateSettings({ costReportingPeriod: next }).catch((err: unknown) => { + setPeriod(previous); + // Kept apart from `error`: the rescan for the restored window clears it. + setPeriodError(err instanceof Error ? err.message : String(err)); + }); + }, + [period], + ); useEffect(() => { load(false); @@ -84,15 +111,20 @@ export default function UsageSpendTab(_props: TabProps) { }; }, [load]); + // Label the period column with the window the summary was built for, so it + // stays accurate while a rescan for a newly picked window is in flight. + const periodTitle = costPeriodLabel(summary?.reportingPeriod ?? period, t); + const onShare = useCallback(() => { setShareError(null); const error = shareUsageSpendPng( summary, t("UsageSpendTitle"), `codexbar-usage-spend-${summary?.reportingDay ?? "unknown"}.png`, + periodTitle, ); if (error) setShareError(t(error as LocaleKey)); - }, [summary, t]); + }, [summary, t, periodTitle]); const onCopyJson = useCallback(async () => { setShareError(null); @@ -168,19 +200,10 @@ export default function UsageSpendTab(_props: TabProps) {
+ {period !== null && } +
- {([7, 30, 0] as const).map((days) => ( - - ))} -
); diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 8e9be16dc2..297c12b609 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -8,6 +8,7 @@ use super::usage::{OutputFormat, ProviderSelection}; use crate::codex_costs::{ CodexHostCostReport, CodexHostCostsArgs, CodexHostOutcome, run_codex_host_costs, }; +use crate::codex_workspaces::short_session_id; use crate::core::{CostScanOptions, ProviderId}; use crate::cost_scanner::{CostScanner, CostSummary}; use crate::settings::Settings; @@ -436,23 +437,6 @@ fn print_codex_session_output(result: &CostResult, days: u32) { println!(" Not a subscription bill or plan value · local usage × public API prices"); } -fn short_session_id(value: &str) -> String { - let trimmed = value.trim(); - if trimmed.chars().count() <= 12 { - return trimmed.to_string(); - } - let prefix: String = trimmed.chars().take(4).collect(); - let suffix: String = trimmed - .chars() - .rev() - .take(8) - .collect::>() - .into_iter() - .rev() - .collect(); - format!("{prefix}...{suffix}") -} - /// Print JSON output fn build_json_payloads(results: &[CostResult], days: u32) -> Vec { let settings = Settings::load(); @@ -664,12 +648,6 @@ mod tests { assert_eq!(CostGroupBy::from_arg(Some("session")), CostGroupBy::Session); } - #[test] - fn short_session_id_is_privacy_conscious() { - assert_eq!(short_session_id("abc"), "abc"); - assert_eq!(short_session_id("1234567890abcdef"), "1234...90abcdef"); - } - #[test] fn remote_failure_retains_local_report() { let local_summary = CodexCostSummary::from_summaries_at( diff --git a/rust/src/codex_workspaces/indexer.rs b/rust/src/codex_workspaces/indexer.rs index 993b695412..68d26748c3 100644 --- a/rust/src/codex_workspaces/indexer.rs +++ b/rust/src/codex_workspaces/indexer.rs @@ -12,9 +12,10 @@ use crate::codex_costs::codex_period_start; use crate::core::{CostUsageDayRange, CostUsagePricing, JsonlScanner, sha256_hex}; use super::sidecar::{SidecarError, WorkspaceUsageSidecar}; +use super::thread_names::{SessionRef, thread_titles}; use super::types::{ CodexLocalProjectUsageSnapshot, CostEstimate, DailyPoint, Progress, ProgressPhase, - ProjectUsage, SessionUsage, SourceStatus, UsageTotals, + ProjectUsage, SessionUsage, SourceStatus, UsageTotals, untitled_session_label, }; use super::{CHATS_DISPLAY_NAME, CHATS_PROJECT_ID}; @@ -214,6 +215,7 @@ impl CodexWorkspacesIndex { } } + apply_thread_titles(&mut session_buckets, &scope.codex_home); let mut projects = build_projects(&session_buckets); projects.sort_by(|a, b| { b.latest_activity @@ -245,11 +247,7 @@ impl CodexWorkspacesIndex { .values() .map(SessionBucket::to_session_usage) .collect(); - sessions.sort_by(|a, b| { - b.latest_activity - .cmp(&a.latest_activity) - .then_with(|| a.id.cmp(&b.id)) - }); + sessions.sort_by(SessionUsage::rank_cmp); let snapshot = CodexLocalProjectUsageSnapshot { updated_at: Utc::now(), @@ -359,7 +357,7 @@ impl SessionBucket { .title .clone() .filter(|t| !t.trim().is_empty()) - .unwrap_or_else(|| "Local Codex chat".to_string()), + .unwrap_or_else(|| untitled_session_label(&self.id)), cwd: self.cwd.clone(), started_at: self.started_at, latest_activity: self.latest_activity, @@ -378,6 +376,30 @@ struct DailyAcc { unknown_tokens: u64, } +/// Overlay Codex thread names (session index, then `threads.title`) after the scan. +/// Titles never touch totals, cost, or project grouping. +fn apply_thread_titles(buckets: &mut HashMap, codex_home: &Path) { + let titles = { + let refs: Vec> = buckets + .values() + .map(|bucket| SessionRef { + id: &bucket.id, + cwd: bucket.cwd.as_deref(), + }) + .collect(); + thread_titles( + codex_home, + non_empty_env("CODEX_SQLITE_HOME").as_deref(), + &refs, + ) + }; + for (id, title) in titles { + if let Some(bucket) = buckets.get_mut(&id) { + bucket.title = Some(title); + } + } +} + fn build_projects(sessions: &HashMap) -> Vec { let mut by_project: HashMap> = HashMap::new(); for session in sessions.values() { @@ -412,11 +434,10 @@ fn build_projects(sessions: &HashMap) -> Vec = + buckets.iter().map(|b| b.to_session_usage()).collect(); + top_sessions.sort_by(SessionUsage::rank_cmp); + top_sessions.truncate(5); ProjectUsage { id: first.project_id.clone(), display_name: first.project_display_name.clone(), @@ -955,10 +976,7 @@ mod tests { snap.redact_for_privacy(); assert_eq!(snap.projects[0].display_name, "Workspace"); assert!(snap.projects[0].path.is_none()); - assert_eq!( - snap.projects[0].top_sessions[0].display_title, - "Local Codex chat" - ); + assert_eq!(snap.projects[0].top_sessions[0].display_title, "Session s1"); assert!(snap.projects[0].top_sessions[0].cwd.is_none()); } } diff --git a/rust/src/codex_workspaces/mod.rs b/rust/src/codex_workspaces/mod.rs index 1f5115da8d..a0db3043c6 100644 --- a/rust/src/codex_workspaces/mod.rs +++ b/rust/src/codex_workspaces/mod.rs @@ -5,13 +5,18 @@ mod indexer; mod sidecar; +mod thread_names; mod types; +#[cfg(test)] +mod tests; + pub use indexer::{CodexLocalDataScope, CodexWorkspacesIndex, IndexError, project_identity}; pub use sidecar::{PAYLOAD_FORMAT_VERSION, SCHEMA_VERSION, SidecarError, WorkspaceUsageSidecar}; pub use types::{ CodexLocalProjectUsageSnapshot, CostEstimate, DailyPoint, Progress, ProgressPhase, - ProjectUsage, SessionUsage, SourceStatus, UsageTotals, + ProjectUsage, SessionUsage, SourceStatus, UsageTotals, short_session_id, + untitled_session_label, }; pub const CHATS_PROJECT_ID: &str = "chats"; diff --git a/rust/src/codex_workspaces/tests.rs b/rust/src/codex_workspaces/tests.rs new file mode 100644 index 0000000000..f91cd48821 --- /dev/null +++ b/rust/src/codex_workspaces/tests.rs @@ -0,0 +1,179 @@ +//! End-to-end checks for thread names and session ranking in the workspaces snapshot. + +use std::fs::{self, File}; +use std::io::Write; +use std::path::Path; + +use chrono::Local; +use rusqlite::Connection; +use tempfile::TempDir; + +use super::{CodexLocalProjectUsageSnapshot, CodexWorkspacesIndex, short_session_id}; + +fn write_session(home: &Path, id: &str, cwd: &Path, input: i32, output: i32) { + let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); + let parts: Vec<_> = day.split('-').collect(); + let folder = home + .join("sessions") + .join(parts[0]) + .join(parts[1]) + .join(parts[2]); + fs::create_dir_all(&folder).unwrap(); + let mut file = File::create(folder.join(format!("{id}.jsonl"))).unwrap(); + writeln!( + file, + r#"{{"timestamp":"{day}T10:00:00.000Z","type":"session_meta","payload":{{"session_id":"{id}","cwd":"{cwd}","originator":"codex_exec","source":"cli"}}}}"#, + cwd = cwd.to_string_lossy().replace('\\', "\\\\"), + ) + .unwrap(); + writeln!( + file, + r#"{{"timestamp":"{day}T10:00:01.000Z","type":"turn_context","payload":{{"model":"gpt-5"}}}}"# + ) + .unwrap(); + writeln!( + file, + r#"{{"timestamp":"{day}T10:00:02.000Z","type":"event_msg","payload":{{"type":"token_count","info":{{"last_token_usage":{{"input_tokens":{input},"cached_input_tokens":0,"output_tokens":{output}}}}}}}}}"# + ) + .unwrap(); +} + +/// Three sessions in one project; `cheap` < `middle` < `dear` by cost. +fn build_home(root: &Path, named: bool) -> CodexLocalProjectUsageSnapshot { + let home = root.join("codex"); + let project = root.join("project"); + fs::create_dir_all(&project).unwrap(); + let short = "short-id"; + let long = "019f79b9-1790-7921-8d6f-258a1e92b191"; + write_session(&home, short, &project, 1_000, 100); + write_session(&home, "middle-session", &project, 20_000, 1_000); + write_session(&home, long, &project, 900_000, 90_000); + + if named { + // Pin the SQLite home so a developer's CODEX_SQLITE_HOME cannot interfere. + fs::write( + home.join("config.toml"), + format!( + "sqlite_home = '{}'\n", + home.canonicalize().unwrap().display() + ), + ) + .unwrap(); + fs::write( + home.join("session_index.jsonl"), + "{\"id\":\"middle-session\",\"thread_name\":\"Fix the icon\"}\n", + ) + .unwrap(); + let conn = Connection::open(home.join("state_5.sqlite")).unwrap(); + conn.execute_batch( + "CREATE TABLE threads (id TEXT PRIMARY KEY, title TEXT); + INSERT INTO threads (id, title) VALUES ('short-id', 'Database title'); + INSERT INTO threads (id, title) VALUES ('middle-session', 'Ignored, index wins');", + ) + .unwrap(); + } + + CodexWorkspacesIndex::new(30) + .with_codex_home(&home) + .with_sidecar_path(root.join("sidecar.sqlite")) + .load_snapshot(true, |_| {}) + .expect("snapshot") +} + +#[test] +fn sessions_get_thread_names_and_rank_by_cost() { + let tmp = TempDir::new().unwrap(); + let snapshot = build_home(tmp.path(), true); + + let titles: Vec<_> = snapshot + .sessions + .iter() + .map(|s| s.display_title.as_str()) + .collect(); + // Most expensive first; an untitled session shows its short id. + assert_eq!( + titles, + ["Session 019f...1e92b191", "Fix the icon", "Database title"] + ); + let top: Vec<_> = snapshot.projects[0] + .top_sessions + .iter() + .map(|s| s.display_title.as_str()) + .collect(); + assert_eq!(top, titles); +} + +#[test] +fn untitled_sessions_show_the_short_id() { + let tmp = TempDir::new().unwrap(); + let snapshot = build_home(tmp.path(), false); + let titles: Vec<_> = snapshot + .sessions + .iter() + .map(|s| s.display_title.clone()) + .collect(); + assert_eq!( + titles, + [ + format!( + "Session {}", + short_session_id("019f79b9-1790-7921-8d6f-258a1e92b191") + ), + "Session midd...-session".to_string(), + "Session short-id".to_string(), + ] + ); +} + +#[test] +fn naming_and_ranking_do_not_change_totals_or_costs() { + let plain = tempfile::tempdir().unwrap(); + let named = tempfile::tempdir().unwrap(); + let plain = build_home(plain.path(), false); + let named = build_home(named.path(), true); + + assert_eq!(plain.total, named.total); + assert_eq!(plain.daily, named.daily); + assert_eq!(plain.projects.len(), named.projects.len()); + assert_eq!( + plain.projects[0].cost_estimate, + named.projects[0].cost_estimate + ); + let costs = |snapshot: &CodexLocalProjectUsageSnapshot| { + snapshot + .sessions + .iter() + .map(|s| (s.id.clone(), s.cost_estimate, s.totals)) + .collect::>() + }; + assert_eq!(costs(&plain), costs(&named)); +} + +#[test] +fn hide_personal_info_masks_names_with_the_short_id() { + let tmp = TempDir::new().unwrap(); + let mut snapshot = build_home(tmp.path(), true); + snapshot.redact_for_privacy(); + + let titles: Vec<_> = snapshot + .sessions + .iter() + .map(|s| s.display_title.as_str()) + .collect(); + assert_eq!( + titles, + [ + "Session 019f...1e92b191", + "Session midd...-session", + "Session short-id" + ] + ); + assert!(snapshot.sessions.iter().all(|s| s.cwd.is_none())); + assert!(snapshot.projects.iter().all(|p| p.path.is_none())); + assert!( + snapshot.projects[0] + .top_sessions + .iter() + .all(|s| s.display_title.starts_with("Session ")) + ); +} diff --git a/rust/src/codex_workspaces/thread_names.rs b/rust/src/codex_workspaces/thread_names.rs new file mode 100644 index 0000000000..95e3bd1be1 --- /dev/null +++ b/rust/src/codex_workspaces/thread_names.rs @@ -0,0 +1,248 @@ +//! Codex thread names for indexed sessions (upstream 0.68.0 `CodexThreadMetadataReader`). +//! +//! Codex keeps thread names outside the rollout files, so they are overlaid after +//! the usage scan. Order: `session_index.jsonl` in the Codex home (latest line +//! wins), then `threads.title` in the newest `state_.sqlite`. Everything is +//! read-only and best effort: any failure leaves the session untitled. + +use std::collections::{HashMap, HashSet}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader, Read}; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use rusqlite::{Connection, OpenFlags, OptionalExtension}; +use serde::Deserialize; + +/// Lines longer than this in `session_index.jsonl` are skipped. +const SESSION_INDEX_MAX_LINE_BYTES: u64 = 64 * 1024; +/// Only the head of `config.toml` is read when looking for `sqlite_home`. +const CONFIG_PREFIX_BYTES: u64 = 256 * 1024; +const BUSY_TIMEOUT: Duration = Duration::from_millis(100); +const DEFAULT_STATE_DATABASE: &str = "state_5.sqlite"; + +/// A session whose thread name is wanted: its id and original working directory. +pub(super) struct SessionRef<'a> { + pub id: &'a str, + pub cwd: Option<&'a str>, +} + +#[derive(Deserialize)] +struct SessionIndexEntry { + id: String, + thread_name: String, +} + +/// Thread names keyed by session id for the given sessions. +/// +/// `env_sqlite_home` is the raw `CODEX_SQLITE_HOME` value, passed in so callers +/// (and tests) control the environment. +pub(super) fn thread_titles( + codex_home: &Path, + env_sqlite_home: Option<&str>, + sessions: &[SessionRef<'_>], +) -> HashMap { + if sessions.is_empty() { + return HashMap::new(); + } + let ids: HashSet<&str> = sessions.iter().map(|s| s.id).collect(); + let mut titles = indexed_thread_names(&codex_home.join("session_index.jsonl"), &ids); + + let configured_home = configured_sqlite_home(codex_home); + let mut newest_by_dir: HashMap = HashMap::new(); + let mut ids_by_database: HashMap> = HashMap::new(); + for session in sessions { + let dir = sqlite_home( + codex_home, + configured_home.as_deref(), + env_sqlite_home, + session.cwd, + ); + let database = newest_by_dir + .entry(dir.clone()) + .or_insert_with(|| newest_state_database(&dir)) + .clone(); + ids_by_database + .entry(database) + .or_default() + .push(session.id); + } + for (database, session_ids) in ids_by_database { + for (id, title) in database_titles(&database, &session_ids) { + // The session index is the explicit name; the database is the fallback. + titles.entry(id).or_insert(title); + } + } + titles +} + +/// `session_index.jsonl`: `{"id": ..., "thread_name": ...}` per line, latest wins. +fn indexed_thread_names(path: &Path, ids: &HashSet<&str>) -> HashMap { + let mut names = HashMap::new(); + let Ok(file) = File::open(path) else { + return names; + }; + let mut reader = BufReader::new(file); + let mut line = Vec::new(); + loop { + line.clear(); + let Ok(read) = (&mut reader) + .take(SESSION_INDEX_MAX_LINE_BYTES + 1) + .read_until(b'\n', &mut line) + else { + break; + }; + if read == 0 { + break; + } + if line.last() != Some(&b'\n') && read as u64 > SESSION_INDEX_MAX_LINE_BYTES { + // Over-long line: drop the remainder without buffering it. + if !skip_to_line_end(&mut reader) { + break; + } + continue; + } + let Ok(entry) = serde_json::from_slice::(&line) else { + continue; + }; + let name = entry.thread_name.trim(); + if !name.is_empty() && ids.contains(entry.id.as_str()) { + names.insert(entry.id, name.to_string()); + } + } + names +} + +/// Consume input up to and including the next newline. Returns `false` at EOF or on error. +fn skip_to_line_end(reader: &mut R) -> bool { + loop { + let Ok(buffer) = reader.fill_buf() else { + return false; + }; + if buffer.is_empty() { + return false; + } + match buffer.iter().position(|byte| *byte == b'\n') { + Some(index) => { + reader.consume(index + 1); + return true; + } + None => { + let len = buffer.len(); + reader.consume(len); + } + } + } +} + +/// Directory holding `state_.sqlite` for one session. +fn sqlite_home( + codex_home: &Path, + configured_home: Option<&Path>, + env_sqlite_home: Option<&str>, + cwd: Option<&str>, +) -> PathBuf { + if let Some(configured) = configured_home { + return configured.to_path_buf(); + } + if let Some(raw) = env_sqlite_home.map(str::trim).filter(|v| !v.is_empty()) { + let path = Path::new(raw); + if path.is_absolute() { + return path.to_path_buf(); + } + if let Some(cwd) = cwd.map(str::trim).filter(|c| !c.is_empty()) { + return Path::new(cwd).join(path); + } + } + codex_home.to_path_buf() +} + +/// Absolute top-level `sqlite_home` from `/config.toml`. +/// +/// Only the user config is read: project, managed and profile layers are not +/// persisted with the rollout and an untrusted checkout must not redirect the +/// lookup. Only the first 256 KiB is read. +fn configured_sqlite_home(codex_home: &Path) -> Option { + let file = File::open(codex_home.join("config.toml")).ok()?; + let mut bytes = Vec::new(); + file.take(CONFIG_PREFIX_BYTES) + .read_to_end(&mut bytes) + .ok()?; + let contents = String::from_utf8_lossy(&bytes); + let raw = top_level_sqlite_home(&contents)?; + let path = PathBuf::from(raw); + path.is_absolute().then_some(path) +} + +fn top_level_sqlite_home(contents: &str) -> Option { + let table = contents.parse::().ok().or_else(|| { + // A truncated or malformed table section must not hide a valid + // top-level key, so retry with everything before the first table header. + let mut offset = 0usize; + let mut end = contents.len(); + for line in contents.split_inclusive('\n') { + if line.trim_start().starts_with('[') { + end = offset; + break; + } + offset += line.len(); + } + contents[..end].parse::().ok() + })?; + match table.get("sqlite_home")? { + toml::Value::String(value) => Some(value.clone()), + _ => None, + } +} + +/// Highest `state_.sqlite` in `dir`, else the default `state_5.sqlite` path. +fn newest_state_database(dir: &Path) -> PathBuf { + let newest = fs::read_dir(dir).ok().and_then(|entries| { + entries + .flatten() + .filter_map(|entry| { + let name = entry.file_name().into_string().ok()?; + let version = name + .strip_prefix("state_")? + .strip_suffix(".sqlite")? + .parse::() + .unwrap_or(0); + Some((version, entry.path())) + }) + .max_by(|a, b| a.0.cmp(&b.0).then_with(|| b.1.cmp(&a.1))) + .map(|(_, path)| path) + }); + newest.unwrap_or_else(|| dir.join(DEFAULT_STATE_DATABASE)) +} + +/// `SELECT title FROM threads WHERE id = ?1 LIMIT 1` for each id, read-only. +fn database_titles(database: &Path, ids: &[&str]) -> Vec<(String, String)> { + if !database.is_file() { + return Vec::new(); + } + let Ok(conn) = Connection::open_with_flags( + database, + OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) else { + return Vec::new(); + }; + if conn.busy_timeout(BUSY_TIMEOUT).is_err() { + return Vec::new(); + } + let Ok(mut statement) = conn.prepare("SELECT title FROM threads WHERE id = ?1 LIMIT 1") else { + return Vec::new(); + }; + ids.iter() + .filter_map(|id| { + let title = statement + .query_row([id], |row| row.get::<_, Option>(0)) + .optional() + .ok()??; + let title = title?.trim().to_string(); + (!title.is_empty()).then(|| ((*id).to_string(), title)) + }) + .collect() +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/codex_workspaces/thread_names/tests.rs b/rust/src/codex_workspaces/thread_names/tests.rs new file mode 100644 index 0000000000..287cf639bc --- /dev/null +++ b/rust/src/codex_workspaces/thread_names/tests.rs @@ -0,0 +1,239 @@ +use super::*; +use tempfile::TempDir; + +fn session<'a>(id: &'a str, cwd: Option<&'a str>) -> SessionRef<'a> { + SessionRef { id, cwd } +} + +fn write_state(dir: &Path, file: &str, rows: &[(&str, Option<&str>)]) { + fs::create_dir_all(dir).unwrap(); + let conn = Connection::open(dir.join(file)).unwrap(); + conn.execute_batch("CREATE TABLE threads (id TEXT PRIMARY KEY, title TEXT)") + .unwrap(); + for (id, title) in rows { + conn.execute( + "INSERT INTO threads (id, title) VALUES (?1, ?2)", + (id, title), + ) + .unwrap(); + } +} + +fn write_index(home: &Path, contents: &str) { + fs::create_dir_all(home).unwrap(); + fs::write(home.join("session_index.jsonl"), contents).unwrap(); +} + +fn titles(home: &Path, env: Option<&str>, sessions: &[SessionRef<'_>]) -> HashMap { + thread_titles(home, env, sessions) +} + +fn toml_path(path: &Path) -> String { + path.to_string_lossy().replace('\\', "\\\\") +} + +#[test] +fn session_index_latest_line_wins_and_blank_names_are_ignored() { + let tmp = TempDir::new().unwrap(); + write_index( + tmp.path(), + concat!( + "{\"id\":\"a\",\"thread_name\":\"First\"}\n", + "{\"id\":\"a\",\"thread_name\":\" Renamed \"}\n", + "{\"id\":\"b\",\"thread_name\":\" \"}\n", + "{\"id\":\"c\",\"thread_name\":\"Not requested\"}\n", + "not json\n", + "{\"id\":\"d\",\"thread_name\":\"No trailing newline\"}", + ), + ); + let found = titles( + tmp.path(), + None, + &[session("a", None), session("b", None), session("d", None)], + ); + assert_eq!(found.get("a").map(String::as_str), Some("Renamed")); + assert_eq!( + found.get("d").map(String::as_str), + Some("No trailing newline") + ); + assert!(!found.contains_key("b")); + assert!(!found.contains_key("c")); +} + +#[test] +fn session_index_skips_lines_over_64_kib() { + let tmp = TempDir::new().unwrap(); + let huge = "x".repeat(70 * 1024); + write_index( + tmp.path(), + &format!( + "{{\"id\":\"a\",\"thread_name\":\"{huge}\"}}\n{{\"id\":\"b\",\"thread_name\":\"After\"}}\n" + ), + ); + let found = titles(tmp.path(), None, &[session("a", None), session("b", None)]); + assert!(!found.contains_key("a")); + assert_eq!(found.get("b").map(String::as_str), Some("After")); +} + +#[test] +fn database_title_is_the_fallback_and_the_index_wins() { + let tmp = TempDir::new().unwrap(); + write_state( + tmp.path(), + "state_5.sqlite", + &[ + ("a", Some("Database name")), + ("b", Some(" Trimmed ")), + ("c", None), + ("d", Some(" ")), + ], + ); + write_index( + tmp.path(), + "{\"id\":\"a\",\"thread_name\":\"Index name\"}\n", + ); + let found = titles( + tmp.path(), + None, + &[ + session("a", None), + session("b", None), + session("c", None), + session("d", None), + session("missing", None), + ], + ); + assert_eq!(found.get("a").map(String::as_str), Some("Index name")); + assert_eq!(found.get("b").map(String::as_str), Some("Trimmed")); + assert_eq!(found.len(), 2); +} + +#[test] +fn newest_state_database_version_is_used() { + let tmp = TempDir::new().unwrap(); + write_state(tmp.path(), "state_5.sqlite", &[("a", Some("old"))]); + write_state(tmp.path(), "state_10.sqlite", &[("a", Some("newest"))]); + write_state(tmp.path(), "state_7.sqlite", &[("a", Some("middle"))]); + let found = titles(tmp.path(), None, &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("newest")); +} + +#[test] +fn missing_or_corrupt_database_yields_no_titles() { + let tmp = TempDir::new().unwrap(); + assert!(titles(tmp.path(), None, &[session("a", None)]).is_empty()); + fs::write(tmp.path().join("state_5.sqlite"), b"not a database").unwrap(); + assert!(titles(tmp.path(), None, &[session("a", None)]).is_empty()); +} + +#[test] +fn config_sqlite_home_must_be_absolute_and_top_level() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + let other = tmp.path().join("other"); + fs::create_dir_all(&home).unwrap(); + write_state(&home, "state_5.sqlite", &[("a", Some("from home"))]); + write_state(&other, "state_5.sqlite", &[("a", Some("from config"))]); + + let config = home.join("config.toml"); + fs::write( + &config, + format!("sqlite_home = \"{}\"\n", toml_path(&other)), + ) + .unwrap(); + let found = titles(&home, None, &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("from config")); + + fs::write(&config, "sqlite_home = \"relative/dir\"\n").unwrap(); + let found = titles(&home, None, &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("from home")); + + fs::write( + &config, + format!( + "model = \"x\"\n[projects.p]\nsqlite_home = \"{}\"\n", + toml_path(&other) + ), + ) + .unwrap(); + let found = titles(&home, None, &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("from home")); +} + +#[test] +fn config_top_level_key_survives_a_malformed_table_section() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + let other = tmp.path().join("other"); + fs::create_dir_all(&home).unwrap(); + write_state(&other, "state_5.sqlite", &[("a", Some("from config"))]); + fs::write( + home.join("config.toml"), + format!( + "sqlite_home = \"{}\"\r\n[broken\r\nkey = \r\n", + toml_path(&other) + ), + ) + .unwrap(); + let found = titles(&home, None, &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("from config")); +} + +#[test] +fn config_beats_environment_which_beats_codex_home() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + let from_config = tmp.path().join("cfg"); + let from_env = tmp.path().join("env"); + fs::create_dir_all(&home).unwrap(); + write_state(&home, "state_5.sqlite", &[("a", Some("home"))]); + write_state(&from_env, "state_5.sqlite", &[("a", Some("env"))]); + write_state(&from_config, "state_5.sqlite", &[("a", Some("config"))]); + let env = from_env.to_string_lossy().into_owned(); + + let found = titles(&home, Some(&env), &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("env")); + + fs::write( + home.join("config.toml"), + format!("sqlite_home = '{}'\n", from_config.display()), + ) + .unwrap(); + let found = titles(&home, Some(&env), &[session("a", None)]); + assert_eq!(found.get("a").map(String::as_str), Some("config")); +} + +#[test] +fn relative_environment_home_resolves_against_the_session_cwd() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("home"); + let cwd_one = tmp.path().join("repo-one"); + let cwd_two = tmp.path().join("repo-two"); + fs::create_dir_all(&home).unwrap(); + write_state(&home, "state_5.sqlite", &[("c", Some("home title"))]); + write_state(&cwd_one.join("db"), "state_5.sqlite", &[("a", Some("one"))]); + write_state(&cwd_two.join("db"), "state_5.sqlite", &[("b", Some("two"))]); + let one = cwd_one.to_string_lossy().into_owned(); + let two = cwd_two.to_string_lossy().into_owned(); + + let found = titles( + &home, + Some("db"), + &[ + session("a", Some(&one)), + session("b", Some(&two)), + // No working directory: a relative path cannot be resolved, use the home. + session("c", None), + ], + ); + assert_eq!(found.get("a").map(String::as_str), Some("one")); + assert_eq!(found.get("b").map(String::as_str), Some("two")); + assert_eq!(found.get("c").map(String::as_str), Some("home title")); +} + +#[test] +fn no_sessions_reads_nothing() { + let tmp = TempDir::new().unwrap(); + write_index(tmp.path(), "{\"id\":\"a\",\"thread_name\":\"x\"}\n"); + assert!(titles(tmp.path(), None, &[]).is_empty()); +} diff --git a/rust/src/codex_workspaces/types.rs b/rust/src/codex_workspaces/types.rs index e96af15da3..2849069bc3 100644 --- a/rust/src/codex_workspaces/types.rs +++ b/rust/src/codex_workspaces/types.rs @@ -1,5 +1,7 @@ //! Snapshot DTOs for Codex local Workspaces indexing. +use std::cmp::Ordering; + use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -132,6 +134,48 @@ pub struct SessionUsage { pub top_model: Option, } +impl SessionUsage { + /// Cost used for ranking: `None` when nothing could be priced, so unpriced + /// sessions sort after priced ones (a priced zero stays priced). + pub fn ranking_cost_usd(&self) -> Option { + let cost = self.cost_estimate; + (cost.unknown_tokens == 0 || cost.known_usd > 0.0).then_some(cost.known_usd) + } + + /// Ranking order for the sessions panel: cost descending with unpriced + /// last, then tokens descending, latest activity descending, id ascending. + pub fn rank_cmp(&self, other: &Self) -> Ordering { + match (self.ranking_cost_usd(), other.ranking_cost_usd()) { + (Some(left), Some(right)) => right.total_cmp(&left), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } + .then_with(|| other.totals.total_tokens.cmp(&self.totals.total_tokens)) + .then_with(|| other.latest_activity.cmp(&self.latest_activity)) + .then_with(|| self.id.cmp(&other.id)) + } +} + +/// Short display form of a session id: ids over 12 characters keep the first 4 +/// and last 8 (`019f...1e92b191`). +pub fn short_session_id(value: &str) -> String { + let trimmed = value.trim(); + let count = trimmed.chars().count(); + if count <= 12 { + return trimmed.to_string(); + } + let prefix: String = trimmed.chars().take(4).collect(); + let suffix: String = trimmed.chars().skip(count - 8).collect(); + format!("{prefix}...{suffix}") +} + +/// Label for a session without a thread name, and the masked label under +/// Hide personal information. The session id is not personal data. +pub fn untitled_session_label(session_id: &str) -> String { + format!("Session {}", short_session_id(session_id)) +} + /// Per-project usage aggregate. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -170,7 +214,7 @@ impl CodexLocalProjectUsageSnapshot { /// Does not rewrite the sidecar. pub fn redact_for_privacy(&mut self) { for session in &mut self.sessions { - session.display_title = "Local Codex chat".to_string(); + session.display_title = untitled_session_label(&session.id); session.cwd = None; } for project in &mut self.projects { @@ -181,9 +225,12 @@ impl CodexLocalProjectUsageSnapshot { } project.path = None; for session in &mut project.top_sessions { - session.display_title = "Local Codex chat".to_string(); + session.display_title = untitled_session_label(&session.id); session.cwd = None; } } } } + +#[cfg(test)] +mod tests; diff --git a/rust/src/codex_workspaces/types/tests.rs b/rust/src/codex_workspaces/types/tests.rs new file mode 100644 index 0000000000..25df8bbdd2 --- /dev/null +++ b/rust/src/codex_workspaces/types/tests.rs @@ -0,0 +1,82 @@ +use super::*; + +fn session(id: &str, cost: Option, tokens: u64, minutes: i64) -> SessionUsage { + SessionUsage { + id: id.to_string(), + project_id: "p".to_string(), + display_title: id.to_string(), + cwd: None, + started_at: None, + latest_activity: Some(DateTime::::UNIX_EPOCH + chrono::Duration::minutes(minutes)), + totals: UsageTotals::from_parts(tokens, 0, 0), + cost_estimate: match cost { + Some(known_usd) => CostEstimate { + known_usd, + unknown_tokens: 0, + }, + None => CostEstimate { + known_usd: 0.0, + unknown_tokens: tokens, + }, + }, + top_model: None, + } +} + +fn ranked(mut sessions: Vec) -> Vec { + sessions.sort_by(SessionUsage::rank_cmp); + sessions.into_iter().map(|s| s.id).collect() +} + +#[test] +fn ranking_orders_by_cost_then_tokens_then_activity_then_id_with_unpriced_last() { + let sessions = vec![ + session("b", Some(1.0), 10, 5), + session("a", Some(1.0), 10, 5), + session("older", Some(1.0), 10, 1), + session("more-tokens", Some(1.0), 20, 5), + session("free", Some(0.0), 10, 5), + session("unpriced", None, 999, 5), + session("expensive", Some(5.0), 1, 0), + ]; + let expected = [ + "expensive", + "more-tokens", + "a", + "b", + "older", + "free", + "unpriced", + ]; + for offset in 0..sessions.len() { + let mut rotated = sessions.clone(); + rotated.rotate_left(offset); + assert_eq!(ranked(rotated.clone()), expected); + rotated.reverse(); + assert_eq!(ranked(rotated), expected); + } +} + +#[test] +fn partially_priced_sessions_rank_by_their_known_cost() { + let mut partial = session("partial", Some(2.0), 10, 0); + partial.cost_estimate.unknown_tokens = 4; + assert_eq!(partial.ranking_cost_usd(), Some(2.0)); + let unpriced = session("unpriced", None, 10, 0); + assert_eq!(unpriced.ranking_cost_usd(), None); +} + +#[test] +fn short_session_id_keeps_first_four_and_last_eight_of_long_ids() { + assert_eq!(short_session_id("abc"), "abc"); + assert_eq!(short_session_id(" exactly-12ch "), "exactly-12ch"); + assert_eq!( + short_session_id("019f79b9-1790-7921-8d6f-258a1e92b191"), + "019f...1e92b191" + ); + assert_eq!(short_session_id("fixture-session"), "fixt...-session"); + assert_eq!( + untitled_session_label("019f79b9-1790-7921-8d6f-258a1e92b191"), + "Session 019f...1e92b191" + ); +} From de19e48ca4493a6fc3d043492ceb3a38b73bfee8 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:24:06 +0700 Subject: [PATCH 141/407] Port upstream 0.65.0: Zed opt-in browser billing --- .../src-tauri/src/commands/providers.rs | 6 + .../src-tauri/src/commands/tests.rs | 49 ++ .../sections/UsageSourceSection.test.tsx | 6 + .../providers/sections/usageSourcePolicy.ts | 7 + rust/src/core/provider.rs | 11 +- rust/src/providers/zed/mod.rs | 195 ++++--- rust/src/providers/zed/snapshot.rs | 314 +++++++++++ rust/src/providers/zed/tests.rs | 492 ++++++++++++++++++ 8 files changed, 1011 insertions(+), 69 deletions(-) create mode 100644 rust/src/providers/zed/snapshot.rs create mode 100644 rust/src/providers/zed/tests.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 779f9c801c..5dfef7ea78 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -101,6 +101,12 @@ pub(crate) fn build_fetch_context( (source_mode, None, false) } else { match cookie_source { + // Opt-in web providers keep their default credential lane + // unless the usage source is explicitly Web; a stored or + // browser cookie must not turn Auto into Web. + _ if provider.web_is_opt_in() && usage_source != SourceMode::Web => { + (usage_source, None, false) + } // #433: an explicitly selected, non-empty Claude manual cookie is // authoritative. Do not let an active OAuth token account silently // replace it; this keeps tray refresh behavior aligned with diagnose, diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 94ac348f11..2a34fdfa3c 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -465,6 +465,55 @@ fn fetch_context_grok_explicit_web_still_uses_manual_cookie() { ); } +#[test] +fn fetch_context_zed_default_and_stored_cookie_keep_editor_credential_lane() { + // Zed browser billing is opt-in: neither the default manual cookie source + // nor a stored cookie may turn Auto into Web. + let ctx = super::build_fetch_context( + ProviderId::Zed, + &Settings::default(), + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + assert_eq!(ctx.source_mode, SourceMode::Auto); + assert!(ctx.manual_cookie_header.is_none()); + assert!(!ctx.manual_cookie_missing); + + let mut cookies = ManualCookies::default(); + cookies.set("zed", "zed.session=stored"); + let ctx = super::build_fetch_context( + ProviderId::Zed, + &Settings::default(), + &cookies, + &ApiKeys::default(), + &HashMap::new(), + ); + assert_eq!(ctx.source_mode, SourceMode::Auto); + assert!(ctx.manual_cookie_header.is_none()); +} + +#[test] +fn fetch_context_zed_explicit_web_uses_stored_cookie() { + let mut settings = Settings::default(); + settings.set_usage_source(ProviderId::Zed, "web"); + let mut cookies = ManualCookies::default(); + cookies.set("zed", "zed.session=stored"); + let ctx = super::build_fetch_context( + ProviderId::Zed, + &settings, + &cookies, + &ApiKeys::default(), + &HashMap::new(), + ); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert_eq!( + ctx.manual_cookie_header.as_deref(), + Some("zed.session=stored") + ); +} + #[test] fn fetch_context_opencode_empty_manual_remaps_to_web() { let settings = Settings::default(); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.test.tsx index 77087ce037..76b32b2691 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSourceSection.test.tsx @@ -32,6 +32,12 @@ describe("usage source policy", () => { "oauth", "web", ]); + expect(usageSourcePolicy("zed")?.options.map((option) => option.value)).toEqual([ + "auto", + "oauth", + "web", + ]); + expect(usageSourcePolicy("zed")?.options[0].description).toContain("editor credential"); expect(usageSourcePolicy("antigravity")?.options[0].description).toContain( "skips agy reports without account identity", ); 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 57ab1c71d5..25e76fddf3 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts @@ -53,6 +53,13 @@ const POLICIES: Readonly> = { { value: "oauth", label: "API", description: "Uses the configured Charm Hyper API key only." }, ], }, + 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." }, + ], + }, gitkraken: { options: [ { value: "auto", label: "Auto", description: "Uses the configured GitKraken access token." }, diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 868b66fde6..75c1c1f655 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -443,7 +443,7 @@ impl ProviderId { ProviderId::LiteLLM => None, ProviderId::Poe => None, ProviderId::Devin => None, - ProviderId::Zed => None, + ProviderId::Zed => Some("zed.dev"), ProviderId::CrossModel => None, ProviderId::Sub2Api => None, ProviderId::Wayfinder => None, @@ -909,6 +909,14 @@ pub trait Provider: Send + Sync { false } + /// Whether the web lane is used only when the usage source is explicitly + /// `web`. A cookie domain otherwise lets the shell turn Auto into Web + /// (manual cookie present or browser import), which would replace a + /// provider's default non-web credential. + fn web_is_opt_in(&self) -> bool { + false + } + /// How the shell should treat a failed refresh when a prior good snapshot exists. fn last_good_failure_policy(&self, _error: &str) -> LastGoodFailurePolicy { LastGoodFailurePolicy::Replace @@ -1343,6 +1351,7 @@ mod tests { Some("aistudio.google.com") ); assert_eq!(ProviderId::Kiro.cookie_domain(), Some("kiro.dev")); + assert_eq!(ProviderId::Zed.cookie_domain(), Some("zed.dev")); assert_eq!(ProviderId::Kimi.cookie_domain(), Some("kimi.moonshot.cn")); assert_eq!(ProviderId::OpenCode.cookie_domain(), Some("opencode.ai")); assert_eq!(ProviderId::Venice.cookie_domain(), Some("venice.ai")); diff --git a/rust/src/providers/zed/mod.rs b/rust/src/providers/zed/mod.rs index 3e4f3e50f5..75e64808ed 100644 --- a/rust/src/providers/zed/mod.rs +++ b/rust/src/providers/zed/mod.rs @@ -1,22 +1,46 @@ +//! Zed provider: editor credential lane (default) and opt-in browser billing. + +mod snapshot; + use async_trait::async_trait; -use reqwest::Client; -use serde_json::Value; +use reqwest::{Client, StatusCode}; use crate::core::{ FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + ProviderStateKind, SourceMode, }; +use crate::providers::{BoundedBodyError, read_bounded_response}; const CREDENTIAL_TARGET: &str = "codexbar-zed"; const DEFAULT_URL: &str = "https://cloud.zed.dev/client/users/me"; +const BILLING_URL: &str = "https://cloud.zed.dev/frontend/billing/usage"; +const COOKIE_DOMAIN: &str = "zed.dev"; +const MAX_RESPONSE_BYTES: usize = 512 * 1024; +const SESSION_EXPIRED: &str = + "Zed browser session expired. Sign in to zed.dev in Chrome or update the Cookie header."; +const MISSING_SESSION: &str = + "Sign in to zed.dev in a supported browser or paste a Cookie header to read token spend."; +const COOKIES_DISABLED: &str = + "Enable Zed browser cookies or paste a Cookie header to read token spend."; pub struct ZedProvider { metadata: ProviderMetadata, client: Client, + billing_url: String, } impl ZedProvider { pub fn new() -> Self { + Self::with_client( + BILLING_URL, + crate::core::credentialed_http_client_builder() + .timeout(std::time::Duration::from_secs(15)) + .build() + .unwrap_or_else(|_| Client::new()), + ) + } + + fn with_client(billing_url: impl Into, client: Client) -> Self { Self { metadata: ProviderMetadata { id: ProviderId::Zed, @@ -31,12 +55,37 @@ impl ZedProvider { 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()), + client, + billing_url: billing_url.into(), } } + + async fn fetch_editor(&self, ctx: &FetchContext) -> Result { + let key = crate::providers::resolve_api_key( + ctx.api_key.as_deref(), + CREDENTIAL_TARGET, + &["ZED_API_KEY", "ZED_CREDENTIALS"], + )?; + let url = ctx.workspace_id.as_deref().unwrap_or(DEFAULT_URL); + let request = self + .client + .get(url) + .header(reqwest::header::AUTHORIZATION, key.trim()) + .header(reqwest::header::ACCEPT, "application/json"); + let body = read_body(request, false).await?; + snapshot::editor_result(&body, chrono::Utc::now()) + } + + async fn fetch_web(&self, ctx: &FetchContext) -> Result { + let cookie = web_cookie(ctx)?; + let request = self + .client + .get(&self.billing_url) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::COOKIE, cookie); + let body = read_body(request, true).await?; + snapshot::web_result(&body) + } } impl Default for ZedProvider { @@ -57,77 +106,87 @@ impl Provider for ZedProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => { - let key = crate::providers::resolve_api_key( - ctx.api_key.as_deref(), - CREDENTIAL_TARGET, - &["ZED_API_KEY", "ZED_CREDENTIALS"], - )?; - let url = ctx.workspace_id.as_deref().unwrap_or(DEFAULT_URL); - let response = self - .client - .get(url) - .header("Authorization", key.trim()) - .header("Accept", "application/json") - .send() - .await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); - } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Zed usage returned status {}", - response.status() - ))); - } - let value: Value = response - .json() - .await - .map_err(|e| ProviderError::Parse(format!("Failed to parse Zed usage: {e}")))?; - Ok(ProviderFetchResult::new(snapshot_from_user(&value), "api")) - } - SourceMode::Web | SourceMode::Cli => { - Err(ProviderError::UnsupportedSource(ctx.source_mode)) - } + // The editor credential is the default lane. Browser billing is + // opt-in (Web only) and the two lanes are never combined. + SourceMode::Auto | SourceMode::OAuth => self.fetch_editor(ctx).await, + SourceMode::Web => self.fetch_web(ctx).await, + // The shell maps a disabled cookie source to `Cli`; it must not + // import cookies or send a request. + SourceMode::Cli => Err(ProviderError::NotInstalled(COOKIES_DISABLED.into())), } } fn available_sources(&self) -> Vec { - vec![SourceMode::Auto, SourceMode::OAuth] + vec![SourceMode::Auto, SourceMode::Web, SourceMode::OAuth] } -} -fn snapshot_from_user(value: &Value) -> UsageSnapshot { - let plan = value.get("plan").unwrap_or(value); - let usage = plan.get("usage").unwrap_or(plan); - let edits = usage - .pointer("/edit_predictions") - .or_else(|| usage.pointer("/editPredictions")) - .unwrap_or(usage); - let used = number(edits, &["used"]).unwrap_or(0.0); - let limit = number(edits, &["limit"]); - let percent = limit - .filter(|v| *v > 0.0) - .map_or(0.0, |limit| used / limit * 100.0); - UsageSnapshot::new(RateWindow::new(percent)).with_login_method("Zed") + fn supports_web(&self) -> bool { + true + } + + fn web_is_opt_in(&self) -> bool { + true + } + + fn error_state_kind(&self, error: &ProviderError) -> ProviderStateKind { + match error { + ProviderError::Other(message) if message == SESSION_EXPIRED => { + ProviderStateKind::ExpiredSession + } + other => other.state_kind(), + } + } } -fn number(value: &Value, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_f64)) +async fn read_body(request: reqwest::RequestBuilder, web: bool) -> Result, ProviderError> { + let response = request.send().await?; + let status = response.status(); + if status != StatusCode::OK { + return Err(status_error(status, web)); + } + read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::Read(error) => ProviderError::Network(error), + BoundedBodyError::TooLarge => ProviderError::Parse(format!( + "Zed usage response exceeded {MAX_RESPONSE_BYTES} bytes." + )), + }) } -#[cfg(test)] -mod tests { - use super::*; +/// Cookie for the browser lane: a pasted header, else the browser session. +/// A manual source with no usable header never falls back to the browser. +fn web_cookie(ctx: &FetchContext) -> Result { + let missing = || ProviderError::NotInstalled(MISSING_SESSION.into()); + if let Some(raw) = ctx.manual_cookie_header.as_deref() { + return crate::providers::normalize_cookie_header(raw).ok_or_else(missing); + } + if ctx.manual_cookie_missing { + return Err(missing()); + } + match crate::providers::browser_cookie_header(&[COOKIE_DOMAIN]) { + Ok(header) => crate::providers::normalize_cookie_header(&header).ok_or_else(missing), + Err(ProviderError::NoCookies) => Err(missing()), + Err(error) => Err(error), + } +} - #[test] - fn parses_edit_predictions() { - let snapshot = snapshot_from_user( - &serde_json::json!({"plan":{"usage":{"editPredictions":{"used":50,"limit":200}}}}), - ); - assert_eq!(snapshot.primary.used_percent, 25.0); +fn status_error(status: StatusCode, web: bool) -> ProviderError { + match status { + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN if web => { + ProviderError::Other(SESSION_EXPIRED.into()) + } + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => ProviderError::AuthRequired, + StatusCode::TOO_MANY_REQUESTS => { + ProviderError::Other("Zed usage requests are rate limited.".into()) + } + status if status.is_server_error() => ProviderError::Other(format!( + "Zed cloud API is unavailable (HTTP {}).", + status.as_u16() + )), + status => ProviderError::Other(format!("Zed cloud API returned HTTP {}.", status.as_u16())), } } + +#[cfg(test)] +mod tests; diff --git a/rust/src/providers/zed/snapshot.rs b/rust/src/providers/zed/snapshot.rs new file mode 100644 index 0000000000..1c71d7803f --- /dev/null +++ b/rust/src/providers/zed/snapshot.rs @@ -0,0 +1,314 @@ +//! Zed response validation and snapshot mapping shared by the editor and +//! browser-billing lanes (upstream v0.65.0 `zed.js` `fetchUsage`). +//! +//! Every value is validated; drifted shapes fail without publishing totals. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Deserializer, de::Error as _}; +use serde_json::Value; + +use crate::core::{ + CostSnapshot, NamedRateWindow, ProviderDisplayDetail, ProviderError, ProviderFetchResult, + RateWindow, SubscriptionMetadata, UsageSnapshot, +}; + +/// `Number.MAX_SAFE_INTEGER`, the upper bound upstream accepts for counts and cents. +const MAX_SAFE_INTEGER: f64 = 9_007_199_254_740_991.0; +const PARSE_MESSAGE: &str = "Could not parse Zed usage response. Its format may have changed."; +const SPEND_PERIOD: &str = "Current billing period"; + +fn parse_error() -> ProviderError { + ProviderError::Parse(PARSE_MESSAGE.into()) +} + +/// A non-negative safe integer (upstream `count`). Stored as `f64` because a +/// safe integer is exactly representable and no truncating cast is needed. +#[derive(Debug, Clone, Copy)] +struct Count(f64); + +impl<'de> Deserialize<'de> for Count { + fn deserialize>(deserializer: D) -> Result { + let value = f64::deserialize(deserializer)?; + if value.fract() == 0.0 && (0.0..=MAX_SAFE_INTEGER).contains(&value) { + Ok(Self(value)) + } else { + Err(D::Error::custom("not a safe non-negative integer")) + } + } +} + +/// A finite, non-negative amount of cents no larger than the safe integer +/// range (upstream `cents`), converted to dollars on read. +#[derive(Debug, Clone, Copy)] +struct Cents(f64); + +impl<'de> Deserialize<'de> for Cents { + fn deserialize>(deserializer: D) -> Result { + let value = f64::deserialize(deserializer)?; + if value.is_finite() && (0.0..=MAX_SAFE_INTEGER).contains(&value) { + Ok(Self(value)) + } else { + Err(D::Error::custom("not a safe non-negative amount")) + } + } +} + +impl Cents { + fn dollars(self) -> f64 { + self.0 / 100.0 + } +} + +/// A non-blank string (upstream `text`). +#[derive(Debug, Clone)] +struct Text(String); + +impl<'de> Deserialize<'de> for Text { + fn deserialize>(deserializer: D) -> Result { + let value = String::deserialize(deserializer)?; + if value.trim().is_empty() { + Err(D::Error::custom("blank text")) + } else { + Ok(Self(value)) + } + } +} + +#[derive(Debug, Clone, Copy)] +enum RawLimit { + Unlimited, + Null, + Count(Count), +} + +#[derive(Deserialize)] +struct LimitedObject { + limited: Count, +} + +/// `"unlimited"`, `null`, an integer, or `{"limited": n}`; a missing key fails. +fn deserialize_limit<'de, D: Deserializer<'de>>(deserializer: D) -> Result { + match Value::deserialize(deserializer)? { + Value::String(text) if text == "unlimited" => Ok(RawLimit::Unlimited), + Value::Null => Ok(RawLimit::Null), + value @ Value::Object(_) => LimitedObject::deserialize(value) + .map(|object| RawLimit::Count(object.limited)) + .map_err(D::Error::custom), + value => Count::deserialize(value) + .map(RawLimit::Count) + .map_err(D::Error::custom), + } +} + +#[derive(Deserialize)] +struct EditPredictions { + used: Count, + #[serde(deserialize_with = "deserialize_limit")] + limit: RawLimit, +} + +#[derive(Deserialize)] +struct TokenSpend { + spend_in_cents: Cents, + #[serde(default)] + limit_in_cents: Option, +} + +#[derive(Deserialize)] +struct BillingUsage { + edit_predictions: EditPredictions, + token_spend: TokenSpend, +} + +#[derive(Deserialize)] +struct BillingResponse { + plan: Text, + current_usage: BillingUsage, +} + +#[derive(Deserialize)] +struct EditorUsage { + edit_predictions: EditPredictions, +} + +#[derive(Deserialize)] +struct SubscriptionPeriod { + started_at: Text, + ended_at: Text, +} + +#[derive(Deserialize)] +struct EditorPlan { + plan_v3: Text, + usage: EditorUsage, + has_overdue_invoices: bool, + #[serde(default)] + subscription_period: Option, +} + +#[derive(Deserialize)] +struct EditorUser { + #[serde(rename = "id")] + _id: Count, + github_login: String, + #[serde(default)] + name: Option, +} + +#[derive(Deserialize)] +struct EditorResponse { + user: EditorUser, + plan: EditorPlan, +} + +fn parse_root Deserialize<'de>>(body: &[u8]) -> Result { + let value: Value = serde_json::from_slice(body).map_err(|_| parse_error())?; + if !value.is_object() { + return Err(parse_error()); + } + serde_json::from_value(value).map_err(|_| parse_error()) +} + +/// Browser billing lane: `GET /frontend/billing/usage`. +pub(super) fn web_result(body: &[u8]) -> Result { + let response: BillingResponse = parse_root(body)?; + let usage = response.current_usage; + let primary = primary_window(&usage.edit_predictions, true)?; + let snapshot = UsageSnapshot::new(primary).with_login_method(plan_label(&response.plan.0)); + + let spent = usage.token_spend.spend_in_cents.dollars(); + let cap = usage.token_spend.limit_in_cents.map(Cents::dollars); + let mut result = ProviderFetchResult::new(snapshot, "web") + .with_display_detail(ProviderDisplayDetail::new( + "token-spend", + "Token spend", + format_usd(spent), + )) + .with_display_detail(ProviderDisplayDetail::new( + "token-spend-limit", + "Spend limit", + cap.map_or_else(|| "Not reported".to_string(), format_usd), + )); + if let Some(cap) = cap { + result = result + .with_cost(CostSnapshot::new(spent, "USD", SPEND_PERIOD).with_limit(cap)) + .with_display_detail(ProviderDisplayDetail::new( + "token-spend-remaining", + "Remaining budget", + format_usd((cap - spent).max(0.0)), + )); + } + Ok(result) +} + +/// Editor lane: `GET /client/users/me` with the editor credential. +pub(super) fn editor_result( + body: &[u8], + now: DateTime, +) -> Result { + let response: EditorResponse = parse_root(body)?; + let EditorResponse { user, plan } = response; + let primary = primary_window(&plan.usage.edit_predictions, false)?; + let mut snapshot = UsageSnapshot::new(primary).with_login_method(plan_label(&plan.plan_v3.0)); + if !user.github_login.trim().is_empty() { + snapshot = snapshot.with_email(user.github_login); + } + if let Some(name) = user.name.filter(|name| !name.trim().is_empty()) { + snapshot = snapshot.with_organization(name); + } + if plan.has_overdue_invoices { + snapshot.extra_rate_windows.push( + NamedRateWindow::new( + "zed.overdue-invoices", + "Billing", + RateWindow::with_details(100.0, None, None, Some("Overdue invoices".into())), + ) + .with_usage_known(false), + ); + } + if let Some(period) = plan.subscription_period { + let start = parse_instant(&period.started_at.0)?; + let end = parse_instant(&period.ended_at.0)?; + let elapsed_percent = if end > start { + (now - start).num_milliseconds() as f64 / (end - start).num_milliseconds() as f64 + * 100.0 + } else { + 0.0 + }; + snapshot = snapshot + .with_secondary(RateWindow::with_details( + elapsed_percent, + None, + Some(end), + Some(cycle_description((end - now).num_seconds())), + )) + .with_subscription(Some(SubscriptionMetadata::new(None, None, Some(end)))); + } + Ok(ProviderFetchResult::new(snapshot, "api")) +} + +fn parse_instant(value: &str) -> Result, ProviderError> { + DateTime::parse_from_rfc3339(value) + .map(|instant| instant.with_timezone(&Utc)) + .map_err(|_| parse_error()) +} + +/// Edit-prediction lane. `null` means unlimited only in the browser payload. +fn primary_window( + predictions: &EditPredictions, + null_is_unlimited: bool, +) -> Result { + let used = predictions.used.0; + let limit = match predictions.limit { + RawLimit::Unlimited => None, + RawLimit::Null if null_is_unlimited => None, + RawLimit::Null => return Err(parse_error()), + RawLimit::Count(limit) => Some(limit.0), + }; + Ok(match limit { + None => RateWindow::with_details(0.0, None, None, Some("Unlimited".into())), + Some(limit) if limit > 0.0 => RateWindow::with_details( + used / limit * 100.0, + None, + None, + Some(format!("{} / {limit} predictions", used.min(limit))), + ), + Some(_) => RateWindow::informational("No edit predictions included"), + }) +} + +/// `zed_pro_trial` -> `Zed Pro Trial`. +fn plan_label(plan: &str) -> String { + plan.split(['_', ' ']) + .filter(|word| !word.is_empty()) + .map(|word| { + let mut chars = word.chars(); + chars.next().map_or_else(String::new, |first| { + first + .to_uppercase() + .chain(chars.flat_map(char::to_lowercase)) + .collect() + }) + }) + .collect::>() + .join(" ") +} + +fn cycle_description(remaining_seconds: i64) -> String { + if remaining_seconds <= 0 { + return "Cycle ended".into(); + } + let hours = remaining_seconds / 3600; + let minutes = remaining_seconds % 3600 / 60; + if hours >= 24 { + format!("Cycle ends in {}d {}h", hours / 24, hours % 24) + } else if hours > 0 { + format!("Cycle ends in {hours}h {minutes}m") + } else { + format!("Cycle ends in {minutes}m") + } +} + +fn format_usd(value: f64) -> String { + CostSnapshot::new(value, "USD", SPEND_PERIOD).format_used() +} diff --git a/rust/src/providers/zed/tests.rs b/rust/src/providers/zed/tests.rs new file mode 100644 index 0000000000..ed3173c68a --- /dev/null +++ b/rust/src/providers/zed/tests.rs @@ -0,0 +1,492 @@ +use chrono::{TimeZone, Utc}; +use reqwest::redirect::Policy; + +use super::snapshot::{editor_result, web_result}; +use super::*; + +// Fixtures mirror upstream v0.65.0 `ZedPluginTests.swift` / `ZedStatusProbeTests.swift`. +const BILLING: &str = r#"{"plan":"zed_pro","current_usage":{ + "token_spend":{"spend_in_cents":250,"limit_in_cents":1000}, + "edit_predictions":{"used":12,"limit":100}}}"#; + +fn editor_body(plan: &str, used: u32, limit: &str, overdue: bool) -> String { + format!( + r#"{{"user":{{"id":4242,"github_login":"octocat","name":"The Octocat"}}, + "feature_flags":[], + "plan":{{"plan_v3":"{plan}", + "subscription_period":{{"started_at":"2026-05-13T00:00:00.000Z","ended_at":"2026-06-13T00:00:00.000Z"}}, + "usage":{{"edit_predictions":{{"used":{used},"limit":{limit}}}}}, + "has_overdue_invoices":{overdue}}}}}"# + ) +} + +fn detail_value<'a>(result: &'a ProviderFetchResult, id: &str) -> Option<&'a str> { + result + .display_details() + .iter() + .find(|row| row.id() == id) + .map(|row| row.value()) +} + +fn now() -> chrono::DateTime { + Utc.with_ymd_and_hms(2026, 6, 1, 0, 0, 0).unwrap() +} + +#[test] +fn browser_billing_reports_dollar_spend_and_plan() { + let result = web_result(BILLING.as_bytes()).unwrap(); + assert_eq!(result.source_label, "web"); + let cost = result.cost.as_ref().unwrap(); + assert_eq!(cost.used, 2.5); + assert_eq!(cost.limit, Some(10.0)); + assert_eq!(cost.currency_code, "USD"); + assert_eq!(cost.period, "Current billing period"); + assert_eq!(cost.balance, None); + assert_eq!(detail_value(&result, "token-spend"), Some("$2.50")); + assert_eq!(detail_value(&result, "token-spend-limit"), Some("$10.00")); + assert_eq!( + detail_value(&result, "token-spend-remaining"), + Some("$7.50") + ); + assert_eq!(result.usage.primary.used_percent, 12.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("12 / 100 predictions") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("Zed Pro")); + assert!(result.usage.account_email.is_none()); +} + +#[test] +fn zero_and_overage_spend_keep_exact_amounts_and_floor_remaining() { + for cents in ["0", "12.5", "250", "1500"] { + let body = BILLING.replace("250", cents); + let result = web_result(body.as_bytes()).unwrap(); + let spent = cents.parse::().unwrap() / 100.0; + assert_eq!(result.cost.as_ref().unwrap().used, spent); + assert_eq!( + detail_value(&result, "token-spend-remaining"), + Some(format!("${:.2}", (10.0 - spent).max(0.0)).as_str()) + ); + } +} + +#[test] +fn missing_or_null_spend_cap_and_null_limit_do_not_invent_limits() { + for cap in [ + BILLING.replace("1000", "null"), + BILLING.replace(",\"limit_in_cents\":1000", ""), + ] { + let body = cap.replace("\"limit\":100", "\"limit\":null"); + let result = web_result(body.as_bytes()).unwrap(); + assert!(result.cost.is_none()); + assert_eq!(detail_value(&result, "token-spend"), Some("$2.50")); + assert_eq!( + detail_value(&result, "token-spend-limit"), + Some("Not reported") + ); + assert_eq!(detail_value(&result, "token-spend-remaining"), None); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Unlimited") + ); + assert_eq!(result.usage.primary.used_percent, 0.0); + } +} + +#[test] +fn drifted_and_unsafe_billing_values_fail_without_publishing_totals() { + let bodies = [ + "{}".to_string(), + "[]".to_string(), + "login".to_string(), + BILLING.replace("250", "-1"), + BILLING.replace("250", "true"), + BILLING.replace("250", "9007199254740992"), + BILLING.replace("1000", "-1"), + BILLING.replace("\"plan\":\"zed_pro\"", "\"plan\":{}"), + BILLING.replace("\"plan\":\"zed_pro\"", "\"plan\":\" \""), + BILLING.replace("\"used\":12", "\"used\":1.5"), + BILLING.replace("\"used\":12", "\"used\":-1"), + BILLING.replace("\"limit\":100", "\"limit\":\"lots\""), + BILLING.replace("\"limit\":100", "\"other\":100"), + BILLING.replace("\"limit\":100", "\"limit\":[5]"), + BILLING.replace("\"limit\":100", "\"limit\":{\"limited\":-2}"), + BILLING.replace("edit_predictions", "edit_prediction"), + ]; + for body in bodies { + assert!( + matches!(web_result(body.as_bytes()), Err(ProviderError::Parse(_))), + "{body}" + ); + } +} + +#[test] +fn edit_prediction_limit_accepts_unlimited_null_integer_and_limited_object() { + let unlimited = BILLING.replace("\"limit\":100", "\"limit\":\"unlimited\""); + let limited = BILLING.replace("\"limit\":100", "\"limit\":{\"limited\":20}"); + let overage = BILLING.replace("\"used\":12,\"limit\":100", "\"used\":150,\"limit\":100"); + let zero = BILLING.replace("\"limit\":100", "\"limit\":0"); + let primary = |body: &str| web_result(body.as_bytes()).unwrap().usage.primary; + + assert_eq!( + primary(&unlimited).reset_description.as_deref(), + Some("Unlimited") + ); + let limited = primary(&limited); + assert_eq!(limited.used_percent, 60.0); + assert_eq!( + limited.reset_description.as_deref(), + Some("12 / 20 predictions") + ); + let overage = primary(&overage); + assert_eq!(overage.used_percent, 100.0); + assert_eq!( + overage.reset_description.as_deref(), + Some("100 / 100 predictions") + ); + assert!(primary(&zero).is_informational); +} + +#[test] +fn editor_payload_retains_identity_quota_billing_dates_and_overdue_warning() { + let body = editor_body("zed_pro_trial", 10, r#"{"limited":20}"#, true); + let result = editor_result(body.as_bytes(), now()).unwrap(); + let usage = &result.usage; + assert_eq!(result.source_label, "api"); + assert_eq!(usage.account_email.as_deref(), Some("octocat")); + assert_eq!(usage.account_organization.as_deref(), Some("The Octocat")); + assert_eq!(usage.login_method.as_deref(), Some("Zed Pro Trial")); + assert_eq!(usage.primary.used_percent, 50.0); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("10 / 20 predictions") + ); + let cycle = usage.secondary.as_ref().unwrap(); + let end = Utc.with_ymd_and_hms(2026, 6, 13, 0, 0, 0).unwrap(); + assert_eq!(cycle.resets_at, Some(end)); + assert_eq!( + cycle.reset_description.as_deref(), + Some("Cycle ends in 12d 0h") + ); + // 19 of 31 days elapsed. + assert!((cycle.used_percent - 19.0 / 31.0 * 100.0).abs() < 1e-9); + assert_eq!(usage.subscription.as_ref().unwrap().renews_at, Some(end)); + let overdue = &usage.extra_rate_windows[0]; + assert_eq!(overdue.id, "zed.overdue-invoices"); + assert_eq!(overdue.title, "Billing"); + assert!(!overdue.usage_known); + assert_eq!( + overdue.window.reset_description.as_deref(), + Some("Overdue invoices") + ); + assert!(result.cost.is_none()); +} + +#[test] +fn editor_cycle_reports_ended_and_short_countdowns() { + let body = editor_body("zed_pro", 0, "\"unlimited\"", false); + let describe = |now| { + editor_result(body.as_bytes(), now) + .unwrap() + .usage + .secondary + .unwrap() + .reset_description + .unwrap() + }; + let end = Utc.with_ymd_and_hms(2026, 6, 13, 0, 0, 0).unwrap(); + assert_eq!(describe(end + chrono::Duration::seconds(1)), "Cycle ended"); + assert_eq!( + describe(end - chrono::Duration::minutes(150)), + "Cycle ends in 2h 30m" + ); + assert_eq!( + describe(end - chrono::Duration::minutes(45)), + "Cycle ends in 45m" + ); +} + +#[test] +fn editor_plans_retain_names_and_unlimited_predictions() { + for (plan, label) in [ + ("zed_free", "Zed Free"), + ("zed_pro", "Zed Pro"), + ("zed_pro_trial", "Zed Pro Trial"), + ("zed_student", "Zed Student"), + ("zed_business", "Zed Business"), + ] { + let body = editor_body(plan, 3, "\"unlimited\"", false); + let usage = editor_result(body.as_bytes(), now()).unwrap().usage; + assert_eq!(usage.login_method.as_deref(), Some(label)); + assert_eq!(usage.primary.used_percent, 0.0); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("Unlimited") + ); + assert!(usage.extra_rate_windows.is_empty()); + } +} + +#[test] +fn editor_payload_rejects_drifted_shapes() { + let good = editor_body("zed_pro", 1, "5", false); + let bodies = [ + good.replace("\"id\":4242", "\"id\":\"4242\""), + good.replace("\"github_login\":\"octocat\"", "\"github_login\":7"), + good.replace("\"name\":\"The Octocat\"", "\"name\":7"), + good.replace( + "\"has_overdue_invoices\":false", + "\"has_overdue_invoices\":\"no\"", + ), + good.replace("2026-06-13T00:00:00.000Z", "not-a-date"), + good.replace("\"plan_v3\":\"zed_pro\"", "\"plan_v3\":\"\""), + // `null` is unlimited only in the browser payload. + editor_body("zed_pro", 1, "null", false), + // The legacy lenient shape is not an upstream wire shape. + r#"{"plan":{"usage":{"editPredictions":{"used":50,"limit":200}}}}"#.to_string(), + ]; + for body in bodies { + assert!( + matches!( + editor_result(body.as_bytes(), now()), + Err(ProviderError::Parse(_)) + ), + "{body}" + ); + } +} + +#[test] +fn editor_optional_identity_and_period_may_be_absent_or_blank() { + let body = r#"{"user":{"id":1,"github_login":" ","name":null}, + "plan":{"plan_v3":"zed_free","usage":{"edit_predictions":{"used":0,"limit":50}}, + "has_overdue_invoices":false}}"#; + let usage = editor_result(body.as_bytes(), now()).unwrap().usage; + assert!(usage.account_email.is_none()); + assert!(usage.account_organization.is_none()); + assert!(usage.secondary.is_none()); + assert!(usage.subscription.is_none()); + assert!(usage.extra_rate_windows.is_empty()); +} + +#[test] +fn classifies_statuses_per_lane() { + assert!(matches!( + status_error(StatusCode::UNAUTHORIZED, true), + ProviderError::Other(message) if message == SESSION_EXPIRED + )); + assert!(matches!( + status_error(StatusCode::FORBIDDEN, true), + ProviderError::Other(message) if message == SESSION_EXPIRED + )); + assert!(matches!( + status_error(StatusCode::FORBIDDEN, false), + ProviderError::AuthRequired + )); + assert!(matches!( + status_error(StatusCode::TOO_MANY_REQUESTS, true), + ProviderError::Other(message) if message.contains("rate limited") + )); + assert!(matches!( + status_error(StatusCode::BAD_GATEWAY, true), + ProviderError::Other(message) if message.contains("unavailable") && message.contains("502") + )); + assert!(matches!( + status_error(StatusCode::NOT_FOUND, true), + ProviderError::Other(message) if message.contains("404") + )); +} + +#[test] +fn expired_browser_session_is_classified_as_expired_not_unknown() { + let provider = ZedProvider::new(); + assert_eq!( + provider.error_state_kind(&status_error(StatusCode::UNAUTHORIZED, true)), + ProviderStateKind::ExpiredSession + ); + assert_eq!( + provider.error_state_kind(&ProviderError::Other("other".into())), + ProviderStateKind::Unknown + ); +} + +#[test] +fn sources_are_auto_web_and_api_with_web_opt_in() { + let provider = ZedProvider::new(); + assert_eq!( + provider.available_sources(), + vec![SourceMode::Auto, SourceMode::Web, SourceMode::OAuth] + ); + assert!(provider.web_is_opt_in()); +} + +#[tokio::test] +async fn web_source_sends_only_the_zed_cookie_and_no_authorization() { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/frontend/billing/usage") + .match_header("accept", "application/json") + .match_header("cookie", "zed.session=fixture-session") + .match_header("authorization", mockito::Matcher::Missing) + .with_status(200) + .with_body(BILLING) + .create_async() + .await; + let provider = test_provider(&server.url()); + let ctx = FetchContext { + source_mode: SourceMode::Web, + manual_cookie_header: Some("Cookie: zed.session=fixture-session".into()), + // An editor credential must never be sent by, or substituted for, Web. + api_key: Some("4242 editor-token".into()), + ..FetchContext::default() + }; + + let result = provider.fetch_usage(&ctx).await.unwrap(); + mock.assert_async().await; + assert_eq!(result.source_label, "web"); + assert_eq!(result.cost.unwrap().used, 2.5); +} + +#[tokio::test] +async fn web_source_maps_expired_and_failing_responses_without_editor_fallback() { + for (status, expected) in [ + (401, SESSION_EXPIRED), + (403, SESSION_EXPIRED), + (429, "Zed usage requests are rate limited."), + (503, "Zed cloud API is unavailable (HTTP 503)."), + (404, "Zed cloud API returned HTTP 404."), + ] { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/frontend/billing/usage") + .with_status(status) + .with_body("login") + .expect(1) + .create_async() + .await; + let ctx = FetchContext { + source_mode: SourceMode::Web, + manual_cookie_header: Some("zed.session=fixture-session".into()), + api_key: Some("4242 editor-token".into()), + // Would be hit if Web fell back to the editor credential. + workspace_id: Some(format!("{}/client/users/me", server.url())), + ..FetchContext::default() + }; + + let error = test_provider(&server.url()) + .fetch_usage(&ctx) + .await + .err() + .unwrap(); + mock.assert_async().await; + assert!( + matches!(&error, ProviderError::Other(message) if message == expected), + "{status}: {error}" + ); + } +} + +#[tokio::test] +async fn web_source_without_a_usable_cookie_fails_before_any_request() { + for ctx in [ + FetchContext { + source_mode: SourceMode::Web, + manual_cookie_header: Some(" ".into()), + ..FetchContext::default() + }, + FetchContext { + source_mode: SourceMode::Web, + manual_cookie_header: Some("zed.session=bad\nvalue".into()), + ..FetchContext::default() + }, + FetchContext { + source_mode: SourceMode::Web, + manual_cookie_missing: true, + ..FetchContext::default() + }, + // Cookie source Off: the shell maps it to Cli. + FetchContext { + source_mode: SourceMode::Cli, + manual_cookie_header: Some("zed.session=fixture-session".into()), + ..FetchContext::default() + }, + ] { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", mockito::Matcher::Any) + .expect(0) + .create_async() + .await; + let error = test_provider(&server.url()) + .fetch_usage(&ctx) + .await + .err() + .unwrap(); + mock.assert_async().await; + assert!(matches!(error, ProviderError::NotInstalled(_)), "{error}"); + } +} + +#[tokio::test] +async fn auto_and_api_sources_use_the_editor_credential_and_never_the_browser() { + for source_mode in [SourceMode::Auto, SourceMode::OAuth] { + let mut server = mockito::Server::new_async().await; + let editor = server + .mock("GET", "/client/users/me") + .match_header("authorization", "4242 fixture-token") + .match_header("cookie", mockito::Matcher::Missing) + .with_status(200) + .with_body(editor_body("zed_pro", 10, r#"{"limited":20}"#, false)) + .create_async() + .await; + let billing = server + .mock("GET", "/frontend/billing/usage") + .expect(0) + .create_async() + .await; + let ctx = FetchContext { + source_mode, + api_key: Some("4242 fixture-token".into()), + manual_cookie_header: Some("zed.session=fixture-session".into()), + workspace_id: Some(format!("{}/client/users/me", server.url())), + ..FetchContext::default() + }; + + let result = test_provider(&server.url()) + .fetch_usage(&ctx) + .await + .unwrap(); + editor.assert_async().await; + billing.assert_async().await; + assert_eq!(result.source_label, "api"); + assert_eq!(result.usage.account_email.as_deref(), Some("octocat")); + } +} + +#[tokio::test] +async fn editor_lane_keeps_auth_required_for_rejected_credentials() { + let mut server = mockito::Server::new_async().await; + let _mock = server + .mock("GET", "/client/users/me") + .with_status(401) + .create_async() + .await; + let ctx = FetchContext { + source_mode: SourceMode::Auto, + api_key: Some("4242 stale".into()), + workspace_id: Some(format!("{}/client/users/me", server.url())), + ..FetchContext::default() + }; + + assert!(matches!( + test_provider(&server.url()).fetch_usage(&ctx).await, + Err(ProviderError::AuthRequired) + )); +} + +fn test_provider(base_url: &str) -> ZedProvider { + let client = Client::builder().redirect(Policy::none()).build().unwrap(); + ZedProvider::with_client(format!("{base_url}/frontend/billing/usage"), client) +} From 3fb049ab0f402220309f62eddc8dfe681f0960ff Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:25:43 +0700 Subject: [PATCH 142/407] Port upstream 0.65.0: Bifrost identity, unknown rate limits, model labels --- rust/src/providers/bifrost/mod.rs | 517 ++++++++------------- rust/src/providers/bifrost/model_labels.rs | 160 +++++++ rust/src/providers/bifrost/tests.rs | 390 ++++++++++++++++ 3 files changed, 731 insertions(+), 336 deletions(-) create mode 100644 rust/src/providers/bifrost/model_labels.rs create mode 100644 rust/src/providers/bifrost/tests.rs diff --git a/rust/src/providers/bifrost/mod.rs b/rust/src/providers/bifrost/mod.rs index d872103c12..cddce9474f 100644 --- a/rust/src/providers/bifrost/mod.rs +++ b/rust/src/providers/bifrost/mod.rs @@ -10,7 +10,7 @@ use serde_json::Value; use std::{net::IpAddr, time::Duration}; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, + CostSnapshot, FetchContext, NamedRateWindow, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; @@ -18,8 +18,11 @@ const CREDENTIAL_TARGET: &str = "codexbar-bifrost"; const API_KEY_ENV: &str = "BIFROST_API_KEY"; const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; +const MAX_MODEL_ROWS: usize = 5; const QUOTA_PATH: &str = "/api/governance/virtual-keys/quota"; +mod model_labels; + pub struct BifrostProvider { metadata: ProviderMetadata, client: Option, @@ -38,7 +41,6 @@ struct Budget { scope: Scope, source: Option, used: f64, - usage_known: bool, limit: f64, reset: ResetTiming, models: Vec, @@ -285,8 +287,8 @@ fn parse_usage(value: &Value, now: DateTime) -> Result) -> Result) -> Result) -> Result, + key_name: Option, budgets: Vec, limits: Vec<(Scope, Value, usize)>, } @@ -406,7 +409,7 @@ fn add_scopes( } fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { - let mut budget_windows = usage + let budget_windows = usage .budgets .iter() .filter(|budget| budget.limit > 0.0) @@ -426,8 +429,7 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { budget.reset.window_minutes, budget.reset.resets_at, Some(budget_description(budget)), - ) - .with_usage_known(budget.usage_known); + ); (budget, title, window) }) .collect::>(); @@ -441,90 +443,28 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { .unwrap_or_else(|| RateWindow::informational("No Bifrost budget quota reported")); let secondary = root_windows.get(1).map(|(_, _, window)| (*window).clone()); let mut snapshot = UsageSnapshot::new(primary).with_login_method("API"); + if let Some(key_name) = &usage.key_name { + snapshot = snapshot.with_email(key_name.clone()); + } + if let Some(source) = first_root_budget(&usage).and_then(|first| first.source.as_deref()) { + snapshot = snapshot.with_organization(source); + } if let Some(secondary) = secondary { snapshot = snapshot.with_secondary(secondary); } let mut extra_windows = root_windows .iter() .skip(2) - .map(|(_, title, window)| { - ( - format!( - "bifrost-budget-{}", - window.reset_description.as_deref().unwrap_or("extra") - ), - title.clone(), - window.clone(), - ) - }) + .map(|(budget, title, window)| budget_named_window(budget, title, window)) .collect::>(); extra_windows.extend( budget_windows - .drain(..) + .iter() .filter(|(budget, _, _)| !budget.scope.id.is_empty()) - .map(|(budget, title, window)| { - ( - format!("bifrost-{}budget-{}", budget.scope.id, budget.id), - title, - window, - ) - }), + .map(|(budget, title, window)| budget_named_window(budget, title, window)), ); - - for (scope, limit, index) in &usage.limits { - let source = optional_text(limit, "source_name").ok().flatten(); - for (key, title) in [("token", "Tokens"), ("request", "Requests")] { - let max = optional_number(limit, &format!("{key}_max_limit")) - .ok() - .flatten(); - let reset_raw = optional_text(limit, &format!("{key}_reset_duration")) - .ok() - .flatten(); - let Some(max_or_reset) = max - .filter(|max| *max > 0.0) - .or_else(|| reset_raw.as_ref().map(|_| 0.0)) - else { - continue; - }; - let used = optional_number(limit, &format!("{key}_current_usage")) - .ok() - .flatten() - .unwrap_or(0.0); - let reset = reset_timing( - reset_raw.as_deref(), - optional_text(limit, &format!("{key}_last_reset")) - .ok() - .flatten() - .as_deref(), - Utc::now(), - ); - let title = bounded( - &[scope.title.as_deref(), source.as_deref(), Some(title)] - .into_iter() - .flatten() - .collect::>() - .join(" "), - ); - let window = RateWindow::with_details( - if max_or_reset > 0.0 { - percent(used, max_or_reset) - } else { - 0.0 - }, - reset.window_minutes, - reset.resets_at, - reset.label.map(str::to_owned), - ); - extra_windows.push(( - format!("bifrost-{}{}s-{index}", scope.id, key), - title, - window, - )); - } - } - for (id, title, window) in extra_windows { - snapshot = snapshot.with_extra_rate_window(id, title, window); - } + extra_windows.extend(rate_limit_windows(&usage.limits)); + snapshot.extra_rate_windows.extend(extra_windows); if usage.active == Some(false) { snapshot = snapshot.with_extra_rate_window( "bifrost-key-inactive", @@ -534,11 +474,7 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { } let mut result = ProviderFetchResult::new(snapshot, "api"); - if let Some(first) = usage - .budgets - .iter() - .find(|budget| budget.scope.id.is_empty()) - { + if let Some(first) = first_root_budget(&usage) { let mut cost = CostSnapshot::new( first.used, "USD", @@ -585,59 +521,175 @@ fn result_from_usage(usage: ParsedUsage) -> ProviderFetchResult { result = result.with_display_detail(detail); } } - if let Some(first) = usage + if let Some(first) = first_root_budget(&usage) { + for detail in model_details(&first.models) { + result = result.with_display_detail(detail); + } + } + result +} + +fn first_root_budget(usage: &ParsedUsage) -> Option<&Budget> { + usage .budgets .iter() .find(|budget| budget.scope.id.is_empty()) - { - let mut models = first - .models - .iter() - .filter_map(|model| { - let name = optional_text(model, "model") +} + +fn budget_named_window(budget: &Budget, title: &str, window: &RateWindow) -> NamedRateWindow { + NamedRateWindow::new( + format!("bifrost-{}budget-{}", budget.scope.id, budget.id), + title, + window.clone(), + ) +} + +/// Token and request windows for every rate limit. A dimension with a reset +/// duration but no positive max is reported as an unknown-usage window; a +/// dimension with neither is unconfigured and skipped. +fn rate_limit_windows(limits: &[(Scope, Value, usize)]) -> Vec { + let mut windows = Vec::new(); + for (scope, limit, index) in limits { + let source = optional_text(limit, "source_name").ok().flatten(); + for (key, title) in [("token", "Tokens"), ("request", "Requests")] { + let max = optional_number(limit, &format!("{key}_max_limit")) + .ok() + .flatten() + .filter(|max| *max > 0.0); + let reset_raw = optional_text(limit, &format!("{key}_reset_duration")) + .ok() + .flatten(); + if max.is_none() && reset_raw.is_none() { + continue; + } + let used = optional_number(limit, &format!("{key}_current_usage")) + .ok() + .flatten() + .unwrap_or(0.0); + let reset = reset_timing( + reset_raw.as_deref(), + optional_text(limit, &format!("{key}_last_reset")) .ok() .flatten() - .unwrap_or_else(|| "Model".into()); - let provider = optional_text(model, "provider").ok().flatten(); - let cost = optional_number(model, "total_cost").ok().flatten(); - let tokens = optional_number(model, "total_tokens").ok().flatten(); - (cost.unwrap_or(0.0) != 0.0 || tokens.unwrap_or(0.0) != 0.0) - .then_some((name, provider, cost, tokens)) - }) - .collect::>(); - models.sort_by(|a, b| { - b.2.unwrap_or(0.0) - .total_cmp(&a.2.unwrap_or(0.0)) - .then_with(|| b.3.unwrap_or(0.0).total_cmp(&a.3.unwrap_or(0.0))) - .then_with(|| a.0.cmp(&b.0)) - }); - for (index, (name, provider, cost, tokens)) in models.iter().take(5).enumerate() { - let label = bounded( - &[provider.as_deref(), Some(name.as_str())] + .as_deref(), + Utc::now(), + ); + let title = bounded( + &[scope.title.as_deref(), source.as_deref(), Some(title)] .into_iter() .flatten() .collect::>() - .join(" · "), + .join(" "), + ); + let known = max.is_some(); + let window = RateWindow::with_details( + max.map_or(0.0, |max| percent(used, max)), + reset.window_minutes, + reset.resets_at, + reset.label.map(str::to_owned), + ) + .with_usage_known(known); + windows.push( + NamedRateWindow::new( + format!("bifrost-{}{}s-{index}", scope.id, key), + title, + window, + ) + .with_usage_known(known), ); - let value = cost.map(usd).unwrap_or_else(|| "—".into()); - let mut detail = - ProviderDisplayDetail::new(format!("bifrost-model-{index}"), label, value); - if let Some(tokens) = tokens { - detail = detail.and_then(|row| { - row.with_secondary_value(format!("{} tokens", format_count(*tokens))) - }); - } - result = result.with_display_detail(detail); - } - if models.len() > 5 { - result = result.with_display_detail(ProviderDisplayDetail::new( - "bifrost-model-other", - "Other models", - (models.len() - 5).to_string(), - )); } } - result + windows +} + +struct ModelRow { + raw: String, + provider: Option, + cost: Option, + tokens: Option, +} + +/// Spend rows for the first root budget, most expensive first, capped at +/// `MAX_MODEL_ROWS` plus an "Other models" count. +fn model_details(models: &[Value]) -> Vec> { + let mut rows = models + .iter() + .filter_map(|model| { + let cost = optional_number(model, "total_cost").ok().flatten(); + // Upstream truncates the count and drops values outside i64. + let tokens = optional_number(model, "total_tokens") + .ok() + .flatten() + .filter(|tokens| tokens.abs() < 9_223_372_036_854_775_808.0) + .map(f64::trunc); + (cost.unwrap_or(0.0) != 0.0 || tokens.unwrap_or(0.0) != 0.0).then(|| ModelRow { + raw: optional_text(model, "model") + .ok() + .flatten() + .unwrap_or_else(|| "Model".into()), + provider: optional_text(model, "provider").ok().flatten(), + cost, + tokens, + }) + }) + .collect::>(); + rows.sort_by(|a, b| { + b.cost + .unwrap_or(0.0) + .total_cmp(&a.cost.unwrap_or(0.0)) + .then_with(|| b.tokens.unwrap_or(0.0).total_cmp(&a.tokens.unwrap_or(0.0))) + .then_with(|| a.raw.cmp(&b.raw)) + }); + let total = rows.len(); + let visible = &rows[..total.min(MAX_MODEL_ROWS)]; + let names = visible + .iter() + .map(|row| model_labels::model_name(&row.raw)) + .collect::>(); + // The provider prefix only helps when the visible rows mix providers. + let mut providers = visible.iter().filter_map(|row| row.provider.as_deref()); + let mixed = providers + .next() + .is_some_and(|first| providers.any(|provider| provider != first)); + + let mut details = visible + .iter() + .zip(&names) + .enumerate() + .map(|(index, (row, name))| { + // Two rows sharing a short name fall back to the raw model id. + let duplicate = names.iter().filter(|other| *other == name).count() > 1; + let label = bounded( + &[ + mixed.then_some(row.provider.as_deref()).flatten(), + Some(if duplicate { &row.raw } else { name }), + ] + .into_iter() + .flatten() + .collect::>() + .join(" · "), + ); + let value = row.cost.map(usd).unwrap_or_else(|| "—".into()); + let detail = ProviderDisplayDetail::new(format!("bifrost-model-{index}"), label, value); + match row.tokens { + Some(tokens) => detail.and_then(|detail| { + detail.with_secondary_value(format!( + "{} tokens", + model_labels::token_count(tokens) + )) + }), + None => detail, + } + }) + .collect::>(); + if total > MAX_MODEL_ROWS { + details.push(ProviderDisplayDetail::new( + "bifrost-model-other", + "Other models", + (total - MAX_MODEL_ROWS).to_string(), + )); + } + details } fn budget_description(budget: &Budget) -> String { @@ -820,13 +872,6 @@ fn usd(value: f64) -> String { fn bounded(value: &str) -> String { value.chars().take(120).collect() } -fn format_count(value: f64) -> String { - if value.fract() == 0.0 && value.abs() <= i64::MAX as f64 { - format!("{value:.0}") - } else { - format!("{value:.2}") - } -} fn parse_error(reason: &str) -> ProviderError { ProviderError::Parse(format!( "Bifrost returned an unrecognized quota response ({reason})." @@ -834,204 +879,4 @@ fn parse_error(reason: &str) -> ProviderError { } #[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, - time::timeout, - }; - - fn now() -> DateTime { - DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc) - } - - #[test] - fn component_rate_limits_take_precedence_over_aggregate_compatibility_field() { - let value = json!({ - "rate_limit": { "token_max_limit": 900, "token_current_usage": 450, "token_reset_duration": "1h" }, - "rate_limits": [ - { "source_name": "provider-a", "token_max_limit": 100, "token_current_usage": 25, "token_reset_duration": "1h" }, - { "source_name": "provider-b", "request_max_limit": 20, "request_current_usage": 5, "request_reset_duration": "1d" } - ] - }); - let parsed = parse_usage(&value, now()).unwrap(); - assert_eq!(parsed.limits.len(), 2); - assert_eq!(parsed.limits[0].1["token_max_limit"], 100); - assert_eq!(parsed.limits[1].1["request_max_limit"], 20); - assert!( - !parsed - .limits - .iter() - .any(|(_, limit, _)| limit["token_max_limit"] == 900) - ); - } - - #[test] - fn aggregate_rate_limit_is_fallback_when_components_are_absent() { - let parsed = parse_usage(&json!({ - "rate_limit": { "token_max_limit": 50, "token_current_usage": 10, "token_reset_duration": "1h" } - }), now()).unwrap(); - assert_eq!(parsed.limits.len(), 1); - assert_eq!(parsed.limits[0].1["token_max_limit"], 50); - } - - #[test] - fn missing_current_usage_at_a_positive_rate_limit_is_known_zero() { - for (usage_field, id) in [ - ("token_current_usage", "bifrost-tokens-0"), - ("request_current_usage", "bifrost-requests-0"), - ] { - let mut limit = json!({ - "token_max_limit": 100, - "token_current_usage": 25, - "request_max_limit": 20, - "request_current_usage": 5 - }); - limit.as_object_mut().unwrap().remove(usage_field); - let result = - result_from_usage(parse_usage(&json!({ "rate_limit": limit }), now()).unwrap()); - let named = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == id) - .unwrap_or_else(|| panic!("missing named rate window {id}")); - - assert_eq!(named.window.used_percent, 0.0, "{usage_field}"); - assert!(named.usage_known, "{usage_field}"); - assert!(named.window.usage_known(), "{usage_field}"); - } - } - - #[test] - fn validates_gateway_before_request_url_is_built() { - assert!( - quota_url_for_test("https://bifrost.example.com/base/") - .unwrap() - .as_str() - .starts_with("https://bifrost.example.com/base/api/governance/virtual-keys/quota") - ); - assert!(quota_url_for_test("http://10.1.2.3:8080").is_ok()); - assert!(quota_url_for_test("http://bifrost.example.com").is_err()); - assert!(quota_url_for_test("https://user:secret@bifrost.example.com").is_err()); - assert!(quota_url_for_test("ftp://10.1.2.3").is_err()); - } - - #[tokio::test] - async fn rejects_public_http_before_resolving_any_credential() { - let provider = BifrostProvider::new(); - let ctx = FetchContext { - gateway_url: Some("http://public.example.com".into()), - ..FetchContext::default() - }; - let error = provider.fetch_api(&ctx).await.unwrap_err(); - assert!(error.to_string().contains("must use HTTPS")); - } - - #[tokio::test] - async fn does_not_forward_virtual_key_through_gateway_redirects() { - let redirect_target = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let redirect_target_addr = redirect_target.local_addr().unwrap(); - let gateway = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let gateway_addr = gateway.local_addr().unwrap(); - let gateway_task = tokio::spawn(async move { - let (mut stream, _) = gateway.accept().await.unwrap(); - let mut request = vec![0; 4096]; - let read = stream.read(&mut request).await.unwrap(); - let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); - assert!(request.contains("x-bf-vk: test-virtual-key")); - let response = format!( - "HTTP/1.1 302 Found\r\nLocation: http://{redirect_target_addr}/redirected\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" - ); - stream.write_all(response.as_bytes()).await.unwrap(); - }); - let redirect_target_task = tokio::spawn(async move { - timeout(Duration::from_millis(250), redirect_target.accept()).await - }); - - let provider = BifrostProvider::new(); - let ctx = FetchContext { - gateway_url: Some(format!("http://{gateway_addr}")), - api_key: Some("test-virtual-key".into()), - ..FetchContext::default() - }; - let error = provider.fetch_api(&ctx).await.unwrap_err(); - assert!(error.to_string().contains("HTTP 302")); - gateway_task.await.unwrap(); - assert!(redirect_target_task.await.unwrap().is_err()); - } - - #[test] - fn parses_fixed_reset_and_calendar_label_without_inventing_calendar_time() { - let fixed = reset_timing(Some("1h"), Some("2025-12-31T23:30:00Z"), now()); - assert_eq!(fixed.window_minutes, Some(60)); - assert_eq!( - fixed.resets_at.unwrap().to_rfc3339(), - "2026-01-01T00:30:00+00:00" - ); - let calendar = reset_timing(Some("1M"), Some("2025-12-01T00:00:00Z"), now()); - assert_eq!(calendar.label, Some("Monthly")); - assert_eq!(calendar.resets_at, None); - assert_eq!(calendar.window_minutes, None); - assert_eq!(parse_duration("1h30m"), Some(5_400.0)); - } - - #[test] - fn budgets_and_overrides_are_mapped_to_percent_and_spend() { - let parsed = parse_usage( - &json!({ - "virtual_key_name": "Build key", - "budgets": [{ "id": "b1", "max_limit": 10, "current_usage": 5, - "override_amount": 5, "override_mode": "forever", "source_name": "Team" }] - }), - now(), - ) - .unwrap(); - let result = result_from_usage(parsed); - assert!((result.usage.primary.used_percent - (100.0 / 3.0)).abs() < 0.001); - assert_eq!(result.cost.unwrap().limit, Some(15.0)); - } - - #[test] - fn shortest_root_budget_is_primary_and_cost_without_summing_budgets() { - let parsed = parse_usage( - &json!({ - "budgets": [ - { "id": "monthly", "max_limit": 1_000, "current_usage": 200, - "reset_duration": "1M" }, - { "id": "daily", "max_limit": 100, "current_usage": 10, - "reset_duration": "1d" } - ] - }), - now(), - ) - .unwrap(); - let result = result_from_usage(parsed); - - assert_eq!(result.usage.primary.used_percent, 10.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Daily · $10.00 / $100.00") - ); - let secondary = result.usage.secondary.as_ref().unwrap(); - assert_eq!(secondary.used_percent, 20.0); - assert_eq!( - secondary.reset_description.as_deref(), - Some("Monthly · $200.00 / $1000.00") - ); - - let cost = result.cost.unwrap(); - assert_eq!(cost.used, 10.0); - assert_eq!(cost.limit, Some(100.0)); - assert_eq!(cost.period, "Daily"); - } - - #[test] - fn malformed_optional_scope_collection_fails_closed() { - assert!(parse_usage(&json!({ "provider_configs": {} }), now()).is_err()); - } -} +mod tests; diff --git a/rust/src/providers/bifrost/model_labels.rs b/rust/src/providers/bifrost/model_labels.rs new file mode 100644 index 0000000000..432361d0a4 --- /dev/null +++ b/rust/src/providers/bifrost/model_labels.rs @@ -0,0 +1,160 @@ +//! Display formatting for Bifrost per-model spend rows, matching the tagged +//! upstream `modelName()` and `tokenCount()` helpers. + +const REGIONS: [&str; 5] = ["us-gov", "us", "eu", "apac", "global"]; +const VENDORS: [&str; 13] = [ + "ai21", + "amazon", + "anthropic", + "cohere", + "deepseek", + "luma", + "meta", + "mistral", + "openai", + "qwen", + "stability", + "twelvelabs", + "writer", +]; + +/// Strip a leading region and vendor token, a Bedrock `-vN:N` revision, and a +/// trailing date so a versioned gateway model id reads as its short name. +pub(super) fn model_name(raw: &str) -> String { + let name = strip_dotted_prefix(raw, ®IONS); + 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() +} + +fn strip_dotted_prefix<'a>(value: &'a str, tokens: &[&str]) -> &'a str { + for token in tokens { + if let Some(head) = value.get(..token.len()) + && head.eq_ignore_ascii_case(token) + && let Some(rest) = value[token.len()..].strip_prefix('.') + { + return rest; + } + } + value +} + +/// Remove a trailing `-v:`. +fn strip_revision(value: &str) -> &str { + let Some((head, tail)) = value.rsplit_once("-v") else { + return value; + }; + match tail.split_once(':') { + Some((major, minor)) if all_digits(major) && all_digits(minor) => head, + _ => value, + } +} + +/// Remove a trailing `-`/whitespace separator plus `YYYYMMDD` or `YYYY-MM-DD`. +fn strip_date(value: &str) -> &str { + let bytes = value.as_bytes(); + for len in [10usize, 8] { + let Some(start) = bytes.len().checked_sub(len + 1) else { + continue; + }; + let (separator, date) = (bytes[start], &bytes[start + 1..]); + let dated = if len == 10 { + date[4] == b'-' + && date[7] == b'-' + && [0, 1, 2, 3, 5, 6, 8, 9] + .iter() + .all(|&i| date[i].is_ascii_digit()) + } else { + date.iter().all(u8::is_ascii_digit) + }; + if dated && (separator == b'-' || separator.is_ascii_whitespace()) { + return &value[..start]; + } + } + value +} + +fn all_digits(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) +} + +/// Compact token count: `K`/`M`/`B` above 1,000 / 999,500 / 999,500,000 with +/// one decimal below 10 and no trailing `.0`. +pub(super) fn token_count(value: f64) -> String { + let magnitude = value.abs(); + let sign = if value < 0.0 { "-" } else { "" }; + for (threshold, divisor, unit) in [ + (999_500_000.0, 1e9, "B"), + (999_500.0, 1e6, "M"), + (1_000.0, 1e3, "K"), + ] { + if magnitude >= threshold { + let scaled = magnitude / divisor; + let fixed = to_fixed(scaled, usize::from(scaled < 10.0)); + return format!("{sign}{}{unit}", fixed.strip_suffix(".0").unwrap_or(&fixed)); + } + } + format!("{value:.0}") +} + +/// `Number.prototype.toFixed`: exact ties round up, where Rust's formatter +/// rounds them to even. +fn to_fixed(value: f64, digits: usize) -> String { + let exact = format!("{value:.40}"); + let fraction = exact.split_once('.').map_or("", |(_, fraction)| fraction); + let tie = fraction + .get(digits..) + .is_some_and(|tail| tail.starts_with('5') && tail[1..].bytes().all(|b| b == b'0')); + let value = if tie { value.next_up() } else { value }; + format!("{value:.digits$}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_names_drop_region_vendor_revision_and_date() { + for (raw, expected) in [ + ( + "us.anthropic.claude-sonnet-4-20250514-v1:0", + "claude-sonnet-4", + ), + ("EU.Amazon.nova-pro-v1:0", "nova-pro"), + ("us-gov.anthropic.claude-3-haiku", "claude-3-haiku"), + ("gpt-4o-2024-08-06", "gpt-4o"), + ("claude 20250514", "claude"), + ("gpt-4o", "gpt-4o"), + ("global.qwen.qwen3-32b", "qwen3-32b"), + ("usa.model", "usa.model"), + ("anthropic.", "anthropic."), + ("-20250514", "-20250514"), + ("model-v1:", "model-v1:"), + ("20250514", "20250514"), + ] { + assert_eq!(model_name(raw), expected, "{raw}"); + } + } + + #[test] + fn token_counts_use_upstream_thresholds_and_half_up_ties() { + for (value, expected) in [ + (0.0, "0"), + (999.0, "999"), + (1_000.0, "1K"), + (1_250.0, "1.3K"), + (1_500.0, "1.5K"), + (12_500.0, "13K"), + (999_499.0, "999K"), + (999_500.0, "1M"), + (2_340_000.0, "2.3M"), + (999_499_999.0, "999M"), + (999_500_000.0, "1B"), + (-1_250.0, "-1.3K"), + ] { + assert_eq!(token_count(value), expected, "{value}"); + } + } +} diff --git a/rust/src/providers/bifrost/tests.rs b/rust/src/providers/bifrost/tests.rs new file mode 100644 index 0000000000..331c0dedad --- /dev/null +++ b/rust/src/providers/bifrost/tests.rs @@ -0,0 +1,390 @@ +use super::*; +use serde_json::json; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + time::timeout, +}; + +fn now() -> DateTime { + DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc) +} + +#[test] +fn component_rate_limits_take_precedence_over_aggregate_compatibility_field() { + let value = json!({ + "rate_limit": { "token_max_limit": 900, "token_current_usage": 450, "token_reset_duration": "1h" }, + "rate_limits": [ + { "source_name": "provider-a", "token_max_limit": 100, "token_current_usage": 25, "token_reset_duration": "1h" }, + { "source_name": "provider-b", "request_max_limit": 20, "request_current_usage": 5, "request_reset_duration": "1d" } + ] + }); + let parsed = parse_usage(&value, now()).unwrap(); + assert_eq!(parsed.limits.len(), 2); + assert_eq!(parsed.limits[0].1["token_max_limit"], 100); + assert_eq!(parsed.limits[1].1["request_max_limit"], 20); + assert!( + !parsed + .limits + .iter() + .any(|(_, limit, _)| limit["token_max_limit"] == 900) + ); +} + +#[test] +fn aggregate_rate_limit_is_fallback_when_components_are_absent() { + let parsed = parse_usage(&json!({ + "rate_limit": { "token_max_limit": 50, "token_current_usage": 10, "token_reset_duration": "1h" } + }), now()).unwrap(); + assert_eq!(parsed.limits.len(), 1); + assert_eq!(parsed.limits[0].1["token_max_limit"], 50); +} + +#[test] +fn missing_current_usage_at_a_positive_rate_limit_is_known_zero() { + for (usage_field, id) in [ + ("token_current_usage", "bifrost-tokens-0"), + ("request_current_usage", "bifrost-requests-0"), + ] { + let mut limit = json!({ + "token_max_limit": 100, + "token_current_usage": 25, + "request_max_limit": 20, + "request_current_usage": 5 + }); + limit.as_object_mut().unwrap().remove(usage_field); + let result = + result_from_usage(parse_usage(&json!({ "rate_limit": limit }), now()).unwrap()); + let named = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == id) + .unwrap_or_else(|| panic!("missing named rate window {id}")); + + assert_eq!(named.window.used_percent, 0.0, "{usage_field}"); + assert!(named.usage_known, "{usage_field}"); + assert!(named.window.usage_known(), "{usage_field}"); + } +} + +#[test] +fn validates_gateway_before_request_url_is_built() { + assert!( + quota_url_for_test("https://bifrost.example.com/base/") + .unwrap() + .as_str() + .starts_with("https://bifrost.example.com/base/api/governance/virtual-keys/quota") + ); + assert!(quota_url_for_test("http://10.1.2.3:8080").is_ok()); + assert!(quota_url_for_test("http://bifrost.example.com").is_err()); + assert!(quota_url_for_test("https://user:secret@bifrost.example.com").is_err()); + assert!(quota_url_for_test("ftp://10.1.2.3").is_err()); +} + +#[tokio::test] +async fn rejects_public_http_before_resolving_any_credential() { + let provider = BifrostProvider::new(); + let ctx = FetchContext { + gateway_url: Some("http://public.example.com".into()), + ..FetchContext::default() + }; + let error = provider.fetch_api(&ctx).await.unwrap_err(); + assert!(error.to_string().contains("must use HTTPS")); +} + +#[tokio::test] +async fn does_not_forward_virtual_key_through_gateway_redirects() { + let redirect_target = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let redirect_target_addr = redirect_target.local_addr().unwrap(); + let gateway = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let gateway_addr = gateway.local_addr().unwrap(); + let gateway_task = tokio::spawn(async move { + let (mut stream, _) = gateway.accept().await.unwrap(); + let mut request = vec![0; 4096]; + let read = stream.read(&mut request).await.unwrap(); + let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); + assert!(request.contains("x-bf-vk: test-virtual-key")); + let response = format!( + "HTTP/1.1 302 Found\r\nLocation: http://{redirect_target_addr}/redirected\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + let redirect_target_task = + tokio::spawn( + async move { timeout(Duration::from_millis(250), redirect_target.accept()).await }, + ); + + let provider = BifrostProvider::new(); + let ctx = FetchContext { + gateway_url: Some(format!("http://{gateway_addr}")), + api_key: Some("test-virtual-key".into()), + ..FetchContext::default() + }; + let error = provider.fetch_api(&ctx).await.unwrap_err(); + assert!(error.to_string().contains("HTTP 302")); + gateway_task.await.unwrap(); + assert!(redirect_target_task.await.unwrap().is_err()); +} + +#[test] +fn parses_fixed_reset_and_calendar_label_without_inventing_calendar_time() { + let fixed = reset_timing(Some("1h"), Some("2025-12-31T23:30:00Z"), now()); + assert_eq!(fixed.window_minutes, Some(60)); + assert_eq!( + fixed.resets_at.unwrap().to_rfc3339(), + "2026-01-01T00:30:00+00:00" + ); + let calendar = reset_timing(Some("1M"), Some("2025-12-01T00:00:00Z"), now()); + assert_eq!(calendar.label, Some("Monthly")); + assert_eq!(calendar.resets_at, None); + assert_eq!(calendar.window_minutes, None); + assert_eq!(parse_duration("1h30m"), Some(5_400.0)); +} + +#[test] +fn budgets_and_overrides_are_mapped_to_percent_and_spend() { + let parsed = parse_usage( + &json!({ + "virtual_key_name": "Build key", + "budgets": [{ "id": "b1", "max_limit": 10, "current_usage": 5, + "override_amount": 5, "override_mode": "forever", "source_name": "Team" }] + }), + now(), + ) + .unwrap(); + let result = result_from_usage(parsed); + assert!((result.usage.primary.used_percent - (100.0 / 3.0)).abs() < 0.001); + assert_eq!(result.cost.unwrap().limit, Some(15.0)); +} + +#[test] +fn shortest_root_budget_is_primary_and_cost_without_summing_budgets() { + let parsed = parse_usage( + &json!({ + "budgets": [ + { "id": "monthly", "max_limit": 1_000, "current_usage": 200, + "reset_duration": "1M" }, + { "id": "daily", "max_limit": 100, "current_usage": 10, + "reset_duration": "1d" } + ] + }), + now(), + ) + .unwrap(); + let result = result_from_usage(parsed); + + assert_eq!(result.usage.primary.used_percent, 10.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Daily · $10.00 / $100.00") + ); + let secondary = result.usage.secondary.as_ref().unwrap(); + assert_eq!(secondary.used_percent, 20.0); + assert_eq!( + secondary.reset_description.as_deref(), + Some("Monthly · $200.00 / $1000.00") + ); + + let cost = result.cost.unwrap(); + assert_eq!(cost.used, 10.0); + assert_eq!(cost.limit, Some(100.0)); + assert_eq!(cost.period, "Daily"); +} + +#[test] +fn malformed_optional_scope_collection_fails_closed() { + assert!(parse_usage(&json!({ "provider_configs": {} }), now()).is_err()); +} + +fn result_for(value: Value) -> ProviderFetchResult { + result_from_usage(parse_usage(&value, now()).unwrap()) +} + +fn detail_rows(result: &ProviderFetchResult) -> Vec<(String, String, Option)> { + result + .display_details() + .iter() + .filter(|row| row.id().starts_with("bifrost-model-")) + .map(|row| { + ( + row.title().to_owned(), + row.value().to_owned(), + row.secondary_value().map(str::to_owned), + ) + }) + .collect() +} + +#[test] +fn identity_uses_key_name_and_first_root_budget_source() { + let result = result_for(json!({ + "virtual_key_name": " Build key ", + "budgets": [ + { "id": "monthly", "max_limit": 1_000, "current_usage": 1, "reset_duration": "1M", + "source_name": "Monthly team" }, + { "id": "daily", "max_limit": 100, "current_usage": 1, "reset_duration": "1d", + "source_name": "Daily team" } + ], + "provider_configs": [ + { "provider": "openai", "budgets": [ + { "id": "p", "max_limit": 5, "current_usage": 1, "reset_duration": "1h", + "source_name": "Provider team" } + ] } + ] + })); + assert_eq!(result.usage.account_email.as_deref(), Some("Build key")); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Daily team") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("API")); +} + +#[test] +fn identity_is_omitted_when_key_name_and_source_are_empty() { + let result = result_for(json!({ + "virtual_key_name": " ", + "budgets": [{ "id": "b", "max_limit": 5, "current_usage": 1, "source_name": " " }] + })); + assert_eq!(result.usage.account_email, None); + assert_eq!(result.usage.account_organization, None); + assert!( + parse_usage(&json!({ "virtual_key_name": 7 }), now()).is_err(), + "a non-string key name is a malformed response" + ); +} + +#[test] +fn reset_only_rate_limit_dimension_is_an_unknown_named_window() { + let result = result_for(json!({ + "rate_limit": { + "token_max_limit": 0, "token_current_usage": 40, "token_reset_duration": "1h", + "request_max_limit": 0, "request_current_usage": 3, + "request_last_reset": "2025-12-31T23:30:00Z" + } + })); + let windows = &result.usage.extra_rate_windows; + assert_eq!( + windows.len(), + 1, + "unconfigured request dimension is skipped" + ); + let tokens = &windows[0]; + assert_eq!(tokens.id, "bifrost-tokens-0"); + assert_eq!(tokens.window.used_percent, 0.0); + assert!(!tokens.usage_known); + assert!(!tokens.window.usage_known()); + assert_eq!(tokens.window.window_minutes, Some(60)); +} + +#[test] +fn root_budget_extras_use_scope_and_budget_id() { + let result = result_for(json!({ + "budgets": [ + { "id": "a", "max_limit": 10, "current_usage": 1, "reset_duration": "1h" }, + { "id": "b", "max_limit": 10, "current_usage": 1, "reset_duration": "1d" }, + { "id": "c", "max_limit": 10, "current_usage": 1, "reset_duration": "1w" } + ], + "model_configs": [ + { "model_name": "m", "budgets": [{ "id": "d", "max_limit": 10, "current_usage": 1 }] } + ] + })); + let ids = result + .usage + .extra_rate_windows + .iter() + .map(|window| window.id.as_str()) + .collect::>(); + assert_eq!(ids, ["bifrost-budget-c", "bifrost-model-0-budget-d"]); +} + +#[test] +fn missing_current_usage_is_a_known_zero_budget_window() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10 }] + })); + assert_eq!(result.usage.primary.used_percent, 0.0); + assert!(result.usage.primary.usage_known()); +} + +#[test] +fn model_rows_normalize_names_and_format_token_counts() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "us.anthropic.claude-sonnet-4-20250514-v1:0", "provider": "bedrock", + "total_cost": 2.5, "total_tokens": 1_250_000 }, + { "model": "gpt-4o-2024-08-06", "provider": "bedrock", + "total_cost": 0.5, "total_tokens": 999 }, + { "model": "idle", "total_cost": 0, "total_tokens": 0 } + ] }] + })); + assert_eq!( + detail_rows(&result), + [ + ( + "claude-sonnet-4".to_owned(), + "$2.50".to_owned(), + Some("1.3M tokens".to_owned()) + ), + ( + "gpt-4o".to_owned(), + "$0.50".to_owned(), + Some("999 tokens".to_owned()) + ), + ] + ); +} + +#[test] +fn model_rows_prefix_provider_only_when_providers_are_mixed() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "eu.mistral.large-v1:0", "provider": "bedrock", "total_cost": 2 }, + { "model": "gpt-4o", "provider": "openai", "total_cost": 1 } + ] }] + })); + let titles = detail_rows(&result) + .into_iter() + .map(|row| row.0) + .collect::>(); + assert_eq!(titles, ["bedrock · large", "openai · gpt-4o"]); +} + +#[test] +fn model_rows_use_raw_name_when_short_names_collide() { + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": [ + { "model": "us.anthropic.claude-3-haiku-20240307-v1:0", "total_cost": 3 }, + { "model": "claude-3-haiku-20240307", "total_cost": 2 }, + { "model": "gpt-4o", "total_cost": 1 } + ] }] + })); + let titles = detail_rows(&result) + .into_iter() + .map(|row| row.0) + .collect::>(); + assert_eq!( + titles, + [ + "us.anthropic.claude-3-haiku-20240307-v1:0", + "claude-3-haiku-20240307", + "gpt-4o" + ] + ); +} + +#[test] +fn model_rows_cap_at_five_with_an_other_models_count() { + let models = (0..7) + .map(|index| json!({ "model": format!("m{index}"), "total_cost": 10 - index })) + .collect::>(); + let result = result_for(json!({ + "budgets": [{ "id": "b", "max_limit": 10, "current_usage": 3, "per_model_usage": models }] + })); + let rows = detail_rows(&result); + assert_eq!(rows.len(), 6); + assert_eq!(rows[5], ("Other models".to_owned(), "2".to_owned(), None)); +} From 7522ae1483177d2a84fed2141681dd50bf5e1b80 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:27:30 +0700 Subject: [PATCH 143/407] Port upstream 0.66.0: Atlas Cloud typed balance, API login method, icon, single dashboard URL --- .../icons/ProviderIcon-atlascloud.svg | 3 + .../providers/providerIcons.test.ts | 7 + .../src/components/providers/providerIcons.ts | 4 +- rust/src/providers/atlascloud/mod.rs | 46 +++-- rust/src/providers/atlascloud/tests.rs | 166 +++++++++++++++--- rust/src/settings/api_keys.rs | 2 +- 6 files changed, 192 insertions(+), 36 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-atlascloud.svg diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-atlascloud.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-atlascloud.svg new file mode 100644 index 0000000000..843383df48 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-atlascloud.svg @@ -0,0 +1,3 @@ + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts index a204e1155a..6b2b1eb5ab 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts @@ -9,6 +9,13 @@ describe("provider icon registry", () => { } }); + it("ships the upstream Atlas Cloud glyph tinted by the brand color", () => { + const svg = PROVIDER_ICON_REGISTRY.atlascloud.svgPath; + expect(svg).toContain(" { 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 67e51621fb..ee6cebc2ce 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -6,6 +6,7 @@ import abacus from "./icons/ProviderIcon-abacus.svg?raw"; import alibaba from "./icons/ProviderIcon-alibaba.svg?raw"; import amp from "./icons/ProviderIcon-amp.svg?raw"; import antigravity from "./icons/ProviderIcon-antigravity.svg?raw"; +import atlascloud from "./icons/ProviderIcon-atlascloud.svg?raw"; import augment from "./icons/ProviderIcon-augment.svg?raw"; import bedrock from "./icons/ProviderIcon-bedrock.svg?raw"; import claude from "./icons/ProviderIcon-claude.svg?raw"; @@ -93,6 +94,7 @@ const RAW: Record = { alibaba: tint(alibaba), amp: tint(amp), antigravity: tint(antigravity), + atlascloud: tint(atlascloud), augment: tint(augment), bedrock: tint(bedrock), claude: tint(claude), @@ -218,7 +220,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { nanogpt: { id: "nanogpt", brandColor: "#687fa1", fallbackLetter: "N" }, infini: { id: "infini", brandColor: "#687fa1", fallbackLetter: "I" }, abacus: { id: "abacus", brandColor: "#7c3aed", fallbackLetter: "A", svgPath: RAW.abacus }, - atlascloud: { id: "atlascloud", brandColor: "#5975F5", fallbackLetter: "A" }, + atlascloud: { id: "atlascloud", brandColor: "#5975F5", fallbackLetter: "A", svgPath: RAW.atlascloud }, manus: { id: "manus", brandColor: "#34322d", fallbackLetter: "M", svgPath: RAW.manus }, mimo: { id: "mimo", brandColor: "#ff6900", fallbackLetter: "M", svgPath: RAW.mimo }, doubao: { id: "doubao", brandColor: "#2563eb", fallbackLetter: "D", svgPath: RAW.doubao }, diff --git a/rust/src/providers/atlascloud/mod.rs b/rust/src/providers/atlascloud/mod.rs index 41191f11d1..68cd687199 100644 --- a/rust/src/providers/atlascloud/mod.rs +++ b/rust/src/providers/atlascloud/mod.rs @@ -6,14 +6,16 @@ use serde::Deserialize; use std::time::Duration; use crate::core::{ - FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; use crate::providers::{BoundedBodyError, read_bounded_response}; const BALANCE_URL: &str = "https://api.atlascloud.ai/public/v1/balance"; const CREDENTIAL_TARGET: &str = "codexbar-atlascloud"; const API_KEY_ENV: &str = "ATLASCLOUD_API_KEY"; +/// Single canonical console URL, shared with the API-key settings catalog. +pub const DASHBOARD_URL: &str = "https://www.atlascloud.ai/console"; const MAX_RESPONSE_BYTES: usize = 1024 * 1024; const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); @@ -44,7 +46,7 @@ impl AtlasCloudProvider { supports_credits: false, default_enabled: false, is_primary: false, - dashboard_url: Some("https://atlascloud.ai/dashboard"), + dashboard_url: Some(DASHBOARD_URL), status_page_url: None, tertiary_label_key: None, }, @@ -87,12 +89,34 @@ impl AtlasCloudProvider { ProviderError::Parse(format!("Invalid Atlas Cloud response: {error}")) })?; let balance = parse_balance(body)?; - let detail = ProviderDisplayDetail::new("atlascloud-available", "Available", balance); - Ok(ProviderFetchResult::new( - UsageSnapshot::new(RateWindow::informational("Account balance")), - "api", - ) - .with_display_detail(detail)) + Ok(balance_result(balance)) + } +} + +/// Typed USD balance for balance formatting, Usage & Spend and currency +/// conversion, plus a signed display row. `CostSnapshot::with_balance` clamps +/// negatives to zero, so the row is what keeps a deficit (`-$1.25`) visible. +fn balance_result(balance: f64) -> ProviderFetchResult { + let usage = + UsageSnapshot::new(RateWindow::informational("Account balance")).with_login_method("API"); + let cost = CostSnapshot::new(0.0, "USD", "Atlas Cloud balance").with_balance(balance); + let detail = ProviderDisplayDetail::new( + "atlascloud-available", + "Available balance", + format_usd(balance), + ); + ProviderFetchResult::new(usage, "api") + .with_cost(cost) + .with_display_detail(detail) +} + +/// `$95.50` / `-$1.25`; a negative that rounds to zero shows no sign. +fn format_usd(amount: f64) -> String { + let magnitude = format!("{:.2}", amount.abs()); + if amount < 0.0 && magnitude != "0.00" { + format!("-${magnitude}") + } else { + format!("${magnitude}") } } @@ -155,7 +179,7 @@ struct AvailableBalance { value: String, } -fn parse_balance(body: &str) -> Result { +fn parse_balance(body: &str) -> Result { let response: BalanceResponse = serde_json::from_str(body) .map_err(|error| ProviderError::Parse(format!("Invalid Atlas Cloud response: {error}")))?; if response.object != "balance" @@ -176,7 +200,7 @@ fn parse_balance(body: &str) -> Result { if !parsed.is_finite() { return Err(parse_failure("available.value is not a finite number")); } - Ok(amount.to_owned()) + Ok(parsed) } fn is_decimal(value: &str) -> bool { diff --git a/rust/src/providers/atlascloud/tests.rs b/rust/src/providers/atlascloud/tests.rs index fae77cfe49..85879c11ae 100644 --- a/rust/src/providers/atlascloud/tests.rs +++ b/rust/src/providers/atlascloud/tests.rs @@ -9,9 +9,27 @@ fn balance_payload(amount: &str) -> String { } #[test] -fn accepts_decimal_zero_and_negative_balances_without_normalizing_text() { - for amount in ["12.340", "0", "-0.25"] { - assert_eq!(parse_balance(&balance_payload(amount)).unwrap(), amount); +fn parses_zero_and_negative_balances_without_clamping() { + for (amount, expected) in [ + ("12.340", 12.34), + ("0", 0.0), + ("0.00", 0.0), + ("-1.250000", -1.25), + ] { + assert_eq!(parse_balance(&balance_payload(amount)).unwrap(), expected); + } +} + +#[test] +fn formats_signed_usd_like_upstream() { + for (amount, expected) in [ + (95.5, "$95.50"), + (0.0, "$0.00"), + (-1.25, "-$1.25"), + (-0.001, "$0.00"), + (1234.5678, "$1234.57"), + ] { + assert_eq!(format_usd(amount), expected); } } @@ -35,13 +53,28 @@ fn rejects_wrong_envelope_currency_and_non_string_amounts() { #[test] fn accepts_only_signed_decimal_strings_that_parse_to_finite_numbers() { - for amount in ["", "-", "+1", ".5", "1.", "1e3", "NaN", "Infinity", "1.2.3"] { + for amount in [ + "", + " ", + "-", + "+1", + ".5", + "1.", + "1e3", + "1e999", + "0x10", + "NaN", + "Infinity", + "1.2.3", + "private-response", + ] { + let error = parse_balance(&balance_payload(amount)).expect_err(amount); + let ProviderError::Parse(message) = error else { + panic!("unexpected error kind for {amount:?}"); + }; assert!( - matches!( - parse_balance(&balance_payload(amount)), - Err(ProviderError::Parse(_)) - ), - "unexpectedly accepted {amount:?}" + !message.contains("private-response"), + "error echoed response text: {message}" ); } let overflow = "9".repeat(400); @@ -69,14 +102,22 @@ fn maps_auth_permission_rate_limit_and_server_statuses() { status_error(StatusCode::BAD_REQUEST), ProviderError::Other(message) if message.contains("400") )); - assert!(matches!( - status_error(StatusCode::INTERNAL_SERVER_ERROR), - ProviderError::Other(message) if message.contains("unavailable") - )); + for status in [ + StatusCode::INTERNAL_SERVER_ERROR, + StatusCode::SERVICE_UNAVAILABLE, + ] { + assert!(matches!( + status_error(status), + ProviderError::Other(message) if message.contains("unavailable") + )); + } } -#[tokio::test] -async fn sends_bearer_request_and_exposes_balance_as_display_only_detail() { +/// Serve one canned HTTP response; the handle yields the lowercased request. +fn provider_serving( + status_line: &'static str, + body: String, +) -> (AtlasCloudProvider, std::thread::JoinHandle) { let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind local test server"); let address = listener.local_addr().expect("local server address"); let server = std::thread::spawn(move || { @@ -84,10 +125,9 @@ async fn sends_bearer_request_and_exposes_balance_as_display_only_detail() { let mut request = [0_u8; 4096]; let read = stream.read(&mut request).expect("read request"); let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); - let body = balance_payload("-3.25"); write!( stream, - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + "HTTP/1.1 {status_line}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body.len(), body ) @@ -100,21 +140,101 @@ async fn sends_bearer_request_and_exposes_balance_as_display_only_detail() { .expect("test HTTP client"); let provider = AtlasCloudProvider::with_client(format!("http://{address}/public/v1/balance"), client); - let context = FetchContext { + (provider, server) +} + +fn keyed_context() -> FetchContext { + FetchContext { api_key: Some("test-atlas-key".into()), ..FetchContext::default() - }; + } +} + +#[tokio::test] +async fn sends_bearer_request_and_exposes_typed_balance_with_signed_display_row() { + let (provider, server) = provider_serving("200 OK", balance_payload("-3.25")); - let result = provider.fetch_usage(&context).await.expect("balance fetch"); + let result = provider + .fetch_usage(&keyed_context()) + .await + .expect("balance fetch"); let request = server.join().expect("test server thread"); assert!(request.starts_with("get /public/v1/balance ")); assert!(request.contains("authorization: bearer test-atlas-key")); assert_eq!(result.display_details().len(), 1); - assert_eq!(result.display_details()[0].title(), "Available"); - assert_eq!(result.display_details()[0].value(), "-3.25"); + assert_eq!(result.display_details()[0].title(), "Available balance"); + assert_eq!(result.display_details()[0].value(), "-$3.25"); assert!(result.usage.primary.is_informational); assert!(result.usage.secondary.is_none()); - assert!(result.cost.is_none()); + assert_eq!(result.usage.login_method.as_deref(), Some("API")); + let cost = result.cost.expect("typed balance carrier"); + assert_eq!(cost.currency_code, "USD"); + // `with_balance` clamps a deficit to zero; the display row keeps the sign. + assert_eq!(cost.balance, Some(0.0)); +} + +#[tokio::test] +async fn positive_balance_is_typed_and_formatted() { + let (provider, server) = provider_serving("200 OK", balance_payload("95.5")); + + let result = provider + .fetch_usage(&keyed_context()) + .await + .expect("balance fetch"); + server.join().expect("test server thread"); + assert_eq!(result.display_details()[0].value(), "$95.50"); + let cost = result.cost.expect("typed balance carrier"); + assert_eq!(cost.balance, Some(95.5)); + assert_eq!(cost.format_balance().as_deref(), Some("$95.50")); +} + +#[tokio::test] +async fn non_ok_statuses_do_not_echo_the_response_body() { + for status_line in [ + "400 Bad Request", + "401 Unauthorized", + "403 Forbidden", + "429 Too Many Requests", + "503 Service Unavailable", + ] { + let (provider, server) = provider_serving(status_line, "private-response".into()); + let error = provider + .fetch_usage(&keyed_context()) + .await + .expect_err(status_line); + server.join().expect("test server thread"); + let message = error.to_string(); + assert!( + !message.contains("private-response"), + "{status_line}: {message}" + ); + } +} + +#[tokio::test] +async fn unparseable_success_body_does_not_echo_the_response_body() { + let (provider, server) = provider_serving("200 OK", "private-response".into()); + let error = provider + .fetch_usage(&keyed_context()) + .await + .expect_err("invalid body"); + server.join().expect("test server thread"); + assert!(matches!(error, ProviderError::Parse(_))); + assert!(!error.to_string().contains("private-response")); +} + +#[test] +fn metadata_uses_the_single_canonical_console_url() { + let provider = AtlasCloudProvider::new(); + assert_eq!( + provider.metadata().dashboard_url, + Some("https://www.atlascloud.ai/console") + ); + let configured = crate::settings::get_api_key_providers() + .into_iter() + .find(|info| info.id == ProviderId::AtlasCloud) + .expect("Atlas Cloud API-key catalog entry"); + assert_eq!(configured.dashboard_url, provider.metadata().dashboard_url); } #[test] diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index 18fe27ed64..c1b0b55c8b 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -404,7 +404,7 @@ pub fn get_api_key_providers() -> Vec { "Get an API key from Atlas Cloud and set it in Preferences or ATLASCLOUD_API_KEY.", ), config_file_path: None, - dashboard_url: Some("https://www.atlascloud.ai/console"), + dashboard_url: Some(crate::providers::atlascloud::DASHBOARD_URL), }, ProviderConfigInfo { id: ProviderId::ZenMux, From 47f55c7a726364897075f7f4ba59c2feff2014fe Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:28:55 +0700 Subject: [PATCH 144/407] Port upstream 0.66.0: Vercel AI Gateway credit balance provider --- README.md | 1 + .../providers/icons/ProviderIcon-vercel.svg | 3 + .../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/src/core/provider.rs | 17 +- rust/src/core/provider_factory.rs | 7 +- rust/src/core/token_accounts.rs | 3 +- rust/src/providers/mod.rs | 2 + rust/src/providers/vercel/mod.rs | 221 +++++++++++++ rust/src/providers/vercel/tests.rs | 295 ++++++++++++++++++ rust/src/settings/api_keys.rs | 11 + 13 files changed, 561 insertions(+), 6 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-vercel.svg create mode 100644 rust/src/providers/vercel/mod.rs create mode 100644 rust/src/providers/vercel/tests.rs diff --git a/README.md b/README.md index 2e052ad2f9..4269d1a9f5 100755 --- a/README.md +++ b/README.md @@ -81,6 +81,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Kiro | Cookies / CLI | Monthly Credits, Overage | | Vertex AI | gcloud OAuth | Cost | | v0 | API Key | Billing quota, API rate limits, on-demand balance | +| Vercel AI Gateway | API Key | Team credit balance, lifetime spend | | Augment | Cookies | Credits | | OpenCode | Local Config | Usage | | Kimi | Cookies | 5h Rate, Weekly | diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-vercel.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-vercel.svg new file mode 100644 index 0000000000..b8ff338636 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-vercel.svg @@ -0,0 +1,3 @@ + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 2fdc4380d7..a46cab5434 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -60,6 +60,7 @@ import stepfun from "./icons/ProviderIcon-stepfun.svg?raw"; import sub2api from "./icons/ProviderIcon-sub2api.svg?raw"; import t3chat from "./icons/ProviderIcon-t3chat.svg?raw"; import venice from "./icons/ProviderIcon-venice.svg?raw"; +import vercel from "./icons/ProviderIcon-vercel.svg?raw"; import vertexai from "./icons/ProviderIcon-vertexai.svg?raw"; import warp from "./icons/ProviderIcon-warp.svg?raw"; import windsurf from "./icons/ProviderIcon-windsurf.svg?raw"; @@ -148,6 +149,7 @@ const RAW: Record = { sub2api: tint(sub2api), t3chat: tint(t3chat), venice: tint(venice), + vercel: tint(vercel), vertexai: tint(vertexai), warp: tint(warp), windsurf: tint(windsurf), @@ -230,6 +232,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { stepfun: { id: "stepfun", brandColor: "#999999", fallbackLetter: "S", svgPath: RAW.stepfun }, sub2api: { id: "sub2api", brandColor: "#2dc6d8", fallbackLetter: "S", svgPath: RAW.sub2api }, venice: { id: "venice", brandColor: "#111827", fallbackLetter: "V", svgPath: RAW.venice }, + vercel: { id: "vercel", brandColor: "#737373", fallbackLetter: "V", svgPath: RAW.vercel }, openaiapi: { id: "openaiapi", brandColor: "#10a37f", fallbackLetter: "O" }, chutes: { id: "chutes", brandColor: "#ff5c35", fallbackLetter: "C" }, litellm: { id: "litellm", brandColor: "#0ea5e9", fallbackLetter: "L" }, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ccc2e8fe2f..b50eb5e823 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -36,7 +36,7 @@ const HAS_DASHBOARD = new Set([ "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", "opencode", "opencodego", "openrouter", "perplexity", "qoder", "codebuddy", "sakana", "stepfun", "t3chat", "venice", "vertexai", "warp", "windsurf", - "xai", "zai", "fireworks", "meta", "muse", "nous", + "xai", "zai", "fireworks", "meta", "muse", "nous", "vercel", ]); /** Provider IDs that have a status page URL in the backend */ const HAS_STATUS_PAGE = new Set([ diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx index 48c200779c..1713c8336f 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx @@ -253,6 +253,7 @@ function providerSourceHintShort( case "xai": case "fireworks": case "meta": + case "vercel": return t("ProviderSourceApiShort"); case "kiro": return t("ProviderSourceKiroEnvShort"); diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index dda4c24441..9458b3b9d3 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -51,6 +51,7 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["crof", "Crof"], ["stepfun", "StepFun"], ["venice", "Venice"], + ["vercel", "Vercel AI Gateway"], ["openaiapi", "OpenAI API"], ["grok", "Grok"], ["elevenlabs", "ElevenLabs"], diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..2b0a8bf158 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -93,6 +93,7 @@ pub enum ProviderId { Muse, Replicate, Nous, + Vercel, } impl ProviderId { @@ -179,6 +180,7 @@ impl ProviderId { ProviderId::Muse, ProviderId::Replicate, ProviderId::Nous, + ProviderId::Vercel, ] } @@ -226,6 +228,7 @@ impl ProviderId { ProviderId::Meta => "meta", ProviderId::Muse => "muse", ProviderId::Nous => "nous", + ProviderId::Vercel => "vercel", ProviderId::AiAnd => "aiand", ProviderId::Windsurf => "windsurf", ProviderId::Manus => "manus", @@ -313,6 +316,7 @@ impl ProviderId { ProviderId::Meta => "Meta", ProviderId::Muse => "Muse Code", ProviderId::Nous => "Nous Portal", + ProviderId::Vercel => "Vercel AI Gateway", ProviderId::AiAnd => "ai&", ProviderId::Windsurf => "Windsurf", ProviderId::Manus => "Manus", @@ -415,6 +419,7 @@ impl ProviderId { ProviderId::Meta => None, ProviderId::Muse => None, ProviderId::Nous => None, + ProviderId::Vercel => None, ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, @@ -495,6 +500,9 @@ impl ProviderId { "fireworks" | "fireworks-ai" | "fw" => Some(ProviderId::Fireworks), "muse" | "muse-code" | "muse code" => Some(ProviderId::Muse), "nous" | "nous-portal" | "nous portal" | "hermes" => Some(ProviderId::Nous), + "vercel" | "vercel-ai-gateway" | "vercel ai gateway" | "ai-gateway" | "ai gateway" => { + Some(ProviderId::Vercel) + } "meta" | "metaspark" | "meta-spark" | "muse-spark" | "musespark" | "muse spark" | "meta muse spark" => Some(ProviderId::Meta), "aiand" | "ai&" | "ai-and" | "ai and" => Some(ProviderId::AiAnd), @@ -958,6 +966,10 @@ pub fn cli_name_map() -> HashMap<&'static str, ProviderId> { map.insert("nous-portal", ProviderId::Nous); map.insert("nous portal", ProviderId::Nous); map.insert("hermes", ProviderId::Nous); + map.insert("vercel-ai-gateway", ProviderId::Vercel); + map.insert("vercel ai gateway", ProviderId::Vercel); + map.insert("ai-gateway", ProviderId::Vercel); + map.insert("ai gateway", ProviderId::Vercel); map.insert("metaspark", ProviderId::Meta); map.insert("meta-spark", ProviderId::Meta); map.insert("muse-spark", ProviderId::Meta); @@ -1106,6 +1118,8 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Muse => "#0668E1", ProviderId::Replicate => "#000000", ProviderId::Nous => "#D6A55C", + // Upstream uses white; a mid neutral keeps contrast on light and dark surfaces. + ProviderId::Vercel => "#737373", } } @@ -1120,7 +1134,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 80); + assert_eq!(all.len(), 81); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1181,6 +1195,7 @@ mod tests { assert!(all.contains(&ProviderId::Replicate)); assert!(all.contains(&ProviderId::Muse)); assert!(all.contains(&ProviderId::Nous)); + assert!(all.contains(&ProviderId::Vercel)); } #[test] diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 712c1d4fb9..3094ebb337 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -20,9 +20,9 @@ use crate::providers::{ NousProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PiProvider, PoeProvider, QoderProvider, QwenCloudProvider, ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, - T3ChatProvider, TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, - WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, - ZoomMateProvider, + T3ChatProvider, TypeSafeProvider, V0Provider, VeniceProvider, VercelProvider, VertexAIProvider, + WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, + ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -111,6 +111,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Meta => Box::new(MetaProvider::new()), ProviderId::Muse => Box::new(MuseProvider::new()), ProviderId::Nous => Box::new(NousProvider::new()), + ProviderId::Vercel => Box::new(VercelProvider::new()), } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 6dedc0fbde..62e56ea54b 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -377,7 +377,8 @@ impl TokenAccountSupport { | ProviderId::Fireworks | ProviderId::Meta | ProviderId::Nous - | ProviderId::Muse => None, + | ProviderId::Muse + | ProviderId::Vercel => None, } } diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index de9d265f6e..bab539330d 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -80,6 +80,7 @@ pub mod t3chat; pub mod typesafe; pub mod v0; pub mod venice; +pub mod vercel; pub mod vertexai; pub mod warp; pub mod wayfinder; @@ -162,6 +163,7 @@ pub use t3chat::T3ChatProvider; pub use typesafe::TypeSafeProvider; pub use v0::V0Provider; pub use venice::VeniceProvider; +pub use vercel::VercelProvider; pub use vertexai::VertexAIProvider; pub use warp::WarpProvider; pub use wayfinder::WayfinderProvider; diff --git a/rust/src/providers/vercel/mod.rs b/rust/src/providers/vercel/mod.rs new file mode 100644 index 0000000000..b509a8c72e --- /dev/null +++ b/rust/src/providers/vercel/mod.rs @@ -0,0 +1,221 @@ +//! Vercel AI Gateway credit balance provider. +//! +//! `GET https://ai-gateway.vercel.sh/v1/credits` with a bearer API key returns +//! the team's remaining balance and lifetime spend as decimal strings +//! (upstream CodexBar 0.66.0, `Resources/Plugins/vercel.js`). The endpoint +//! reports no quota or billing period, so the primary window is informational +//! and no percentage is invented. + +use async_trait::async_trait; +use reqwest::{Client, StatusCode, redirect::Policy}; +use serde::Deserialize; +use std::time::Duration; + +use crate::core::{ + CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +const CREDITS_URL: &str = "https://ai-gateway.vercel.sh/v1/credits"; +const CREDENTIAL_TARGET: &str = "codexbar-vercel"; +const ENV_KEYS: &[&str] = &["AI_GATEWAY_API_KEY"]; +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +const SECTION_LABEL: &str = "Team credits"; + +pub struct VercelProvider { + metadata: ProviderMetadata, + client: Client, + credits_url: String, +} + +impl VercelProvider { + pub fn new() -> Self { + let client = crate::core::credentialed_http_client_builder() + .redirect(Policy::none()) + .timeout(REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()); + Self::with_client(CREDITS_URL, client) + } + + fn with_client(credits_url: impl Into, client: Client) -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::Vercel, + display_name: "Vercel AI Gateway", + session_label: "Balance", + weekly_label: "Balance", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai-gateway"), + status_page_url: None, + tertiary_label_key: None, + }, + client, + credits_url: credits_url.into(), + } + } + + async fn fetch_credits( + &self, + ctx: &FetchContext, + ) -> Result { + let api_key = + crate::providers::resolve_api_key(ctx.api_key.as_deref(), CREDENTIAL_TARGET, ENV_KEYS)?; + let response = self + .client + .get(&self.credits_url) + .bearer_auth(api_key.trim()) + .header(reqwest::header::ACCEPT, "application/json") + .send() + .await?; + let status = response.status(); + if status != StatusCode::OK { + return Err(status_error(status)); + } + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::Read(error) => ProviderError::Network(error), + BoundedBodyError::TooLarge => parse_failure(), + })?; + Ok(build_result(parse_credits(&body)?)) + } +} + +impl Default for VercelProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for VercelProvider { + fn id(&self) -> ProviderId { + ProviderId::Vercel + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => self.fetch_credits(ctx).await, + SourceMode::Web | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::OAuth] + } +} + +/// Wire shape at upstream v0.66.0: both amounts are decimal strings. +#[derive(Debug, Deserialize)] +struct CreditsResponse { + balance: String, + total_used: String, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +struct Credits { + balance: f64, + total_used: f64, +} + +/// Status mapping keeps the code but never the response body. +fn status_error(status: StatusCode) -> ProviderError { + let code = status.as_u16(); + match status { + StatusCode::UNAUTHORIZED => ProviderError::AuthRequired, + StatusCode::FORBIDDEN => ProviderError::Other(format!( + "Vercel AI Gateway returned HTTP {code}; the API key lacks permission to read credits." + )), + StatusCode::TOO_MANY_REQUESTS => ProviderError::Other(format!( + "Vercel AI Gateway rate limit reached (HTTP {code})." + )), + status if status.is_server_error() => { + ProviderError::Other(format!("Vercel AI Gateway is unavailable (HTTP {code}).")) + } + _ => ProviderError::Other(format!("Vercel AI Gateway returned HTTP {code}.")), + } +} + +/// Fixed message: response data (and serde's value-echoing errors) never +/// reach the user-facing error. +fn parse_failure() -> ProviderError { + ProviderError::Parse("Vercel AI Gateway returned an unrecognized credit balance.".into()) +} + +fn parse_credits(body: &[u8]) -> Result { + let response: CreditsResponse = serde_json::from_slice(body).map_err(|_| parse_failure())?; + let balance = parse_amount(&response.balance).ok_or_else(parse_failure)?; + let total_used = parse_amount(&response.total_used).ok_or_else(parse_failure)?; + if total_used < 0.0 { + return Err(parse_failure()); + } + Ok(Credits { + balance, + total_used, + }) +} + +/// Accept only `^-?\d+(?:\.\d+)?$` that converts to a finite number. +fn parse_amount(value: &str) -> Option { + let digits = value.strip_prefix('-').unwrap_or(value); + let (integer, fraction) = match digits.split_once('.') { + Some((integer, fraction)) => (integer, Some(fraction)), + None => (digits, None), + }; + let all_digits = |part: &str| !part.is_empty() && part.bytes().all(|b| b.is_ascii_digit()); + if !all_digits(integer) || fraction.is_some_and(|part| !all_digits(part)) { + return None; + } + value + .parse::() + .ok() + .filter(|amount| amount.is_finite()) +} + +/// `$95.50`, `-$1.25`, `$0.00`; a negative that rounds to zero carries no sign. +fn format_usd(value: f64) -> String { + let cents = format!("{:.2}", value.abs()); + if value < 0.0 && cents != "0.00" { + format!("-${cents}") + } else { + format!("${cents}") + } +} + +fn build_result(credits: Credits) -> ProviderFetchResult { + let usage = + UsageSnapshot::new(RateWindow::informational(SECTION_LABEL)).with_login_method("API"); + // `with_balance` clamps to >= 0, so a negative balance is left untyped and + // stays visible only through the signed display row below. + let mut cost = CostSnapshot::new(credits.total_used, "USD", SECTION_LABEL); + if credits.balance >= 0.0 { + cost = cost.with_balance(credits.balance); + } + ProviderFetchResult::new(usage, "api") + .with_cost(cost) + .with_display_detail(ProviderDisplayDetail::new( + "vercel-balance", + "Available balance", + format_usd(credits.balance), + )) + .with_display_detail(ProviderDisplayDetail::new( + "vercel-lifetime-spend", + "Lifetime spend", + format_usd(credits.total_used), + )) +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/providers/vercel/tests.rs b/rust/src/providers/vercel/tests.rs new file mode 100644 index 0000000000..780aa0b138 --- /dev/null +++ b/rust/src/providers/vercel/tests.rs @@ -0,0 +1,295 @@ +use super::*; +use std::io::{Read, Write}; +use std::net::TcpListener; + +/// Upstream fixture shape (`APIBalancePluginTests.swift`). +fn credits_payload(balance: &str) -> String { + format!(r#"{{"balance":"{balance}","total_used":"4.50"}}"#) +} + +fn assert_parse_failure(result: Result) { + match result { + Err(ProviderError::Parse(message)) => { + assert!(!message.contains("private-response"), "{message}"); + } + other => panic!("expected parse failure, got {other:?}"), + } +} + +#[test] +fn parses_upstream_fixture_into_typed_amounts() { + let credits = parse_credits(credits_payload("95.50").as_bytes()).unwrap(); + assert_eq!(credits.balance, 95.5); + assert_eq!(credits.total_used, 4.5); +} + +#[test] +fn accepts_zero_negative_and_unpadded_balances() { + for (value, expected) in [("0.00", 0.0), ("-1.250000", -1.25), ("7", 7.0)] { + let credits = parse_credits(credits_payload(value).as_bytes()).unwrap(); + assert_eq!(credits.balance, expected, "{value}"); + } +} + +#[test] +fn rejects_malformed_amounts_without_echoing_response_data() { + for value in ["", " ", "NaN", "1e999", "0x10", "private-response"] { + assert_parse_failure(parse_credits(credits_payload(value).as_bytes())); + } + // Decimal strings that overflow f64 are not finite and must fail too. + assert_parse_failure(parse_credits(credits_payload(&"9".repeat(400)).as_bytes())); + for value in ["+1", ".5", "1.", "1e3", "Infinity", "1.2.3", "-", "--1"] { + assert_parse_failure(parse_credits(credits_payload(value).as_bytes())); + } +} + +#[test] +fn rejects_wrong_shapes_without_echoing_response_data() { + for body in [ + "private-response", + "{}", + "null", + "", + r#"{"balance":12,"total_used":"0"}"#, + r#"{"balance":"1.00","total_used":0}"#, + r#"{"balance":"1.00"}"#, + r#"{"total_used":"1.00"}"#, + r#"{"balance":null,"total_used":"1.00"}"#, + r#"["private-response"]"#, + ] { + assert_parse_failure(parse_credits(body.as_bytes())); + } +} + +#[test] +fn rejects_negative_lifetime_spend() { + assert_parse_failure(parse_credits(br#"{"balance":"1.00","total_used":"-1.00"}"#)); +} + +#[test] +fn ignores_unknown_fields_without_inventing_them() { + let credits = + parse_credits(br#"{"balance":"1.00","total_used":"2.00","extra":"private-response"}"#) + .unwrap(); + assert_eq!(credits.balance, 1.0); + assert_eq!(credits.total_used, 2.0); +} + +#[test] +fn formats_usd_with_sign_before_the_symbol() { + assert_eq!(format_usd(95.5), "$95.50"); + assert_eq!(format_usd(-1.25), "-$1.25"); + assert_eq!(format_usd(0.0), "$0.00"); + assert_eq!(format_usd(-0.001), "$0.00"); + assert_eq!(format_usd(4.5), "$4.50"); +} + +#[test] +fn maps_statuses_without_response_bodies() { + assert!(matches!( + status_error(StatusCode::UNAUTHORIZED), + ProviderError::AuthRequired + )); + for (status, needle) in [ + (StatusCode::FORBIDDEN, "HTTP 403"), + (StatusCode::TOO_MANY_REQUESTS, "rate limit"), + (StatusCode::SERVICE_UNAVAILABLE, "unavailable"), + (StatusCode::INTERNAL_SERVER_ERROR, "unavailable"), + (StatusCode::BAD_REQUEST, "HTTP 400"), + ] { + match status_error(status) { + ProviderError::Other(message) => { + assert!(message.contains(needle), "{status}: {message}"); + assert!(message.contains(&status.as_u16().to_string()), "{message}"); + } + other => panic!("{status}: unexpected {other:?}"), + } + } +} + +#[test] +fn balance_result_is_informational_with_typed_balance_and_rows() { + let result = build_result(Credits { + balance: 95.5, + total_used: 4.5, + }); + assert!(result.usage.primary.is_informational); + assert!(result.usage.secondary.is_none()); + assert_eq!(result.usage.login_method.as_deref(), Some("API")); + let cost = result.cost.as_ref().expect("typed cost"); + assert_eq!(cost.balance, Some(95.5)); + assert_eq!(cost.used, 4.5); + assert_eq!(cost.currency_code, "USD"); + assert!(!cost.always_visible, "lifetime spend is not a 30-day spend"); + let rows: Vec<_> = result + .display_details() + .iter() + .map(|row| (row.title(), row.value())) + .collect(); + assert_eq!( + rows, + [("Available balance", "$95.50"), ("Lifetime spend", "$4.50")] + ); +} + +#[test] +fn zero_balance_is_a_typed_zero_and_negative_stays_display_only() { + let zero = build_result(Credits { + balance: 0.0, + total_used: 4.5, + }); + assert_eq!(zero.cost.as_ref().unwrap().balance, Some(0.0)); + assert_eq!(zero.display_details()[0].value(), "$0.00"); + + let negative = build_result(Credits { + balance: -1.25, + total_used: 4.5, + }); + assert_eq!(negative.cost.as_ref().unwrap().balance, None); + assert_eq!(negative.display_details()[0].value(), "-$1.25"); +} + +#[test] +fn metadata_matches_upstream_descriptor() { + let provider = VercelProvider::new(); + let metadata = provider.metadata(); + assert_eq!(metadata.display_name, "Vercel AI Gateway"); + assert_eq!(metadata.session_label, "Balance"); + assert_eq!(metadata.weekly_label, "Balance"); + assert!(!metadata.default_enabled); + assert_eq!( + metadata.dashboard_url, + Some("https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai-gateway") + ); + assert_eq!(ProviderId::Vercel.cli_name(), "vercel"); + assert_eq!( + ProviderId::from_cli_name("vercel-ai-gateway"), + Some(ProviderId::Vercel) + ); + assert_eq!( + provider.available_sources(), + [SourceMode::Auto, SourceMode::OAuth] + ); +} + +#[tokio::test] +async fn unsupported_sources_fail_before_any_request() { + let provider = VercelProvider::with_client("http://127.0.0.1:1/v1/credits", Client::new()); + let context = FetchContext { + source_mode: SourceMode::Web, + api_key: Some("test-key".into()), + ..FetchContext::default() + }; + assert!(matches!( + provider.fetch_usage(&context).await, + Err(ProviderError::UnsupportedSource(SourceMode::Web)) + )); +} + +/// One-shot local server; returns the lower-cased request head it received. +fn serve_once( + status_line: &'static str, + body: String, +) -> (String, std::thread::JoinHandle) { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("bind local test server"); + let address = listener.local_addr().expect("local server address"); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept request"); + let mut request = [0_u8; 4096]; + let read = stream.read(&mut request).expect("read request"); + let request = String::from_utf8_lossy(&request[..read]).to_ascii_lowercase(); + write!( + stream, + "HTTP/1.1 {status_line}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len(), + ) + .expect("write response"); + request + }); + (format!("http://{address}/v1/credits"), server) +} + +fn test_provider(url: String) -> VercelProvider { + let client = Client::builder() + .redirect(Policy::none()) + .build() + .expect("test HTTP client"); + VercelProvider::with_client(url, client) +} + +fn keyed_context() -> FetchContext { + FetchContext { + api_key: Some(" test-gateway-key ".into()), + ..FetchContext::default() + } +} + +#[tokio::test] +async fn sends_trimmed_bearer_request_and_reports_balance() { + let (url, server) = serve_once("200 OK", credits_payload("-3.25")); + let result = test_provider(url) + .fetch_usage(&keyed_context()) + .await + .expect("credits fetch"); + let request = server.join().expect("test server thread"); + assert!(request.starts_with("get /v1/credits "), "{request}"); + assert!(request.contains("authorization: bearer test-gateway-key\r\n")); + assert_eq!(result.source_label, "api"); + assert_eq!(result.display_details()[0].value(), "-$3.25"); + assert_eq!(result.display_details()[1].value(), "$4.50"); + assert!(result.usage.primary.is_informational); +} + +#[tokio::test] +async fn http_failures_expose_status_but_never_the_body() { + for (status_line, needle) in [ + ("401 Unauthorized", "Authentication required"), + ("403 Forbidden", "HTTP 403"), + ("429 Too Many Requests", "rate limit"), + ("503 Service Unavailable", "unavailable"), + ("400 Bad Request", "HTTP 400"), + ] { + let (url, server) = serve_once(status_line, "private-response".into()); + let error = test_provider(url) + .fetch_usage(&keyed_context()) + .await + .expect_err(status_line); + server.join().expect("test server thread"); + let message = error.to_string(); + assert!(message.contains(needle), "{status_line}: {message}"); + assert!(!message.contains("private-response"), "{message}"); + } +} + +#[tokio::test] +async fn redirects_are_not_followed_with_the_api_key() { + let (url, server) = serve_once( + "302 Found\r\nLocation: http://127.0.0.1:9/steal", + String::new(), + ); + let error = test_provider(url) + .fetch_usage(&keyed_context()) + .await + .expect_err("redirect is not a success"); + server.join().expect("test server thread"); + assert!(error.to_string().contains("HTTP 302"), "{error}"); +} + +#[tokio::test] +async fn malformed_success_body_is_a_parse_error() { + let (url, server) = serve_once("200 OK", "private-response".into()); + let error = test_provider(url) + .fetch_usage(&keyed_context()) + .await + .expect_err("malformed body"); + server.join().expect("test server thread"); + assert!(matches!(error, ProviderError::Parse(_)), "{error}"); + assert!(!error.to_string().contains("private-response")); +} + +#[test] +fn missing_key_uses_the_shared_not_installed_error() { + let result = + crate::providers::resolve_api_key(None, "codexbar-vercel-test-without-credential", &[]); + assert!(matches!(result, Err(ProviderError::NotInstalled(_)))); +} diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index c2adf3565c..09a921c725 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -426,6 +426,17 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: Some("https://portal.neuralwatt.com/dashboard"), }, + ProviderConfigInfo { + id: ProviderId::Vercel, + name: "Vercel AI Gateway", + requires_api_key: true, + api_key_env_var: Some("AI_GATEWAY_API_KEY"), + api_key_help: Some( + "Add a Vercel AI Gateway API key to show the team's credit balance and lifetime spend.", + ), + config_file_path: None, + dashboard_url: Some("https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai-gateway"), + }, ProviderConfigInfo { id: ProviderId::Doubao, name: "Doubao / Volcengine Ark", From 2651567f58d22912213c3364ba97985815038051 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:29:17 +0700 Subject: [PATCH 145/407] Port upstream 0.67.0: retain inherited counter origin for direct fork chains A fork whose parent is itself a fork with no token snapshot before the child forked now resolves to the parent's inherited totals instead of leaving the child unresolved. Ancestor resolution is depth-guarded at 64 and the Codex cache schema is bumped to 5 so stale unresolved entries are rebuilt. --- CHANGELOG.md | 1 + rust/src/core/jsonl_scanner.rs | 7 + rust/src/core/jsonl_scanner/codex.rs | 19 +- rust/src/core/jsonl_scanner/codex/parser.rs | 8 + rust/src/cost_scanner/codex.rs | 1 + rust/src/cost_scanner/codex/logical_target.rs | 61 +++- rust/src/cost_scanner/tests.rs | 3 + rust/src/cost_scanner/tests/direct_fork.rs | 325 ++++++++++++++++++ 8 files changed, 411 insertions(+), 14 deletions(-) create mode 100644 rust/src/cost_scanner/tests/direct_fork.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index c054768f2c..55885d7346 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ menu-bar layout. ### Fixed - Claude: when Hide Personal Info is enabled, keep saved account rows distinguishable with stable localized `Account N` labels and matching redacted tooltips. +- Codex: a fork of a fork whose parent has no token snapshot before the child forked now inherits the whole chain's cumulative counter origin instead of staying unresolved; ancestor chains deeper than 64 stay unresolved, and the Codex cost cache is rebuilt once (cache schema 5). --- diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 13e78bcfe9..3e1285ee16 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -396,6 +396,11 @@ pub struct CodexForkAccountingState { pub history_base_thread_id: Option, pub fork_timestamp: Option, pub inherited_totals: Option, + /// Timestamp of this fork's first own token event. Until then its + /// cumulative counters still sit at `inherited_totals`, so a descendant + /// forked earlier resolves to that origin. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub first_token_timestamp: 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 @@ -449,6 +454,8 @@ pub struct CodexParseResult { pub token_timestamps_monotonic: Option, /// Last token timestamp observed by the parser. pub last_token_timestamp: Option, + /// First token timestamp observed by a from-zero parse. + pub first_token_timestamp: Option, /// Number of timestamp comparisons performed while validating this parse. pub token_timestamp_comparisons: u64, /// Newly consumed bytes in this parse pass. diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 108d4284d7..99f0467868 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -14,8 +14,10 @@ use parser::{CodexParseMode, CodexParserState}; /// 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. Version 4 reparses copied-prefix -/// subagents with locally inferred component baselines. Rebuild older artifacts. -pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 4; +/// subagents with locally inferred component baselines. Version 5 records a +/// fork's first token timestamp so direct-fork chains resolve through parents +/// with no token snapshot at the descendant's fork time. Rebuild older artifacts. +pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 5; /// 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 @@ -293,6 +295,18 @@ impl JsonlScanner { } } + /// Strict counterpart of [`Self::codex_timestamp_at_or_before`]: `earlier` + /// is before `later`. Malformed timestamps fail closed. + pub(crate) fn codex_timestamp_before(earlier: &str, later: &str) -> bool { + match ( + parse_rfc3339_timestamp(earlier), + parse_rfc3339_timestamp(later), + ) { + (Some(earlier), Some(later)) => earlier < later, + _ => false, + } + } + /// Parse a Codex JSONL file pub fn parse_codex_file( file_path: &Path, @@ -643,6 +657,7 @@ impl JsonlScanner { last_totals: parser.previous_totals, token_timestamps_monotonic: parser.token_timestamps_monotonic, last_token_timestamp: parser.previous_token_timestamp, + first_token_timestamp: parser.first_token_timestamp, token_timestamp_comparisons: parser.token_timestamp_comparisons, bytes_read, is_complete, diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index bdce9c68d8..fbfa19b187 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -14,6 +14,10 @@ pub(super) struct CodexParserState { saw_interleaved_totals: bool, pub(super) records: Vec<(CodexUsageRecord, i64)>, pub(super) previous_token_timestamp: Option, + /// First token timestamp seen by this parse. A fork uses it to prove that + /// its inherited origin is still the counter state at a descendant's + /// cutoff. Resumed non-fork parses do not observe the whole prefix. + pub(super) first_token_timestamp: Option, previous_token_timestamp_parsed: Option>, pub(super) token_timestamps_monotonic: Option, pub(super) token_timestamp_comparisons: u64, @@ -278,6 +282,7 @@ impl CodexParserState { saw_interleaved_totals: false, records: Vec::new(), previous_token_timestamp, + first_token_timestamp: None, previous_token_timestamp_parsed, // A parser always validates a fresh prefix. `None` is only an // input marker for the legacy-cache path, not an output state. @@ -846,6 +851,9 @@ impl CodexParserState { self.token_timestamps_monotonic = Some(false); } } + if self.first_token_timestamp.is_none() { + self.first_token_timestamp = Some(timestamp.to_string()); + } self.previous_token_timestamp = Some(timestamp.to_string()); self.previous_token_timestamp_parsed = current_parsed; } diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 6c0ebc5efa..6d747d7de2 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -746,6 +746,7 @@ impl CostScanner { history_base_thread_id: history_base_thread_id.clone(), fork_timestamp: codex_fork_timestamp.clone(), inherited_totals: parse_result.fork_baseline.clone(), + first_token_timestamp: parse_result.first_token_timestamp.clone(), remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), locally_resolved, }) diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index dd42b96a98..e5e94899c6 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -21,6 +21,9 @@ pub(super) enum CodexLineageDecision { Unsafe, } +/// Longest ancestor chain a parent baseline is resolved through. +const CODEX_MAX_LINEAGE_DEPTH: usize = 64; + struct CodexLineageNode { path: String, session_id: Option, @@ -446,6 +449,22 @@ impl CodexLineagePlanner { node_index: usize, child_fork_timestamp: Option<&str>, ) -> Option { + self.parent_owner_baseline_at_depth(cache, node_index, child_fork_timestamp, 0) + } + + /// Cumulative counters of one parent at a descendant's fork time. Ancestors + /// are revalidated on the way up, and a chain deeper than + /// [`CODEX_MAX_LINEAGE_DEPTH`] stays unresolved. + fn parent_owner_baseline_at_depth( + &self, + cache: &CostUsageCache, + node_index: usize, + child_fork_timestamp: Option<&str>, + depth: usize, + ) -> Option { + if depth >= CODEX_MAX_LINEAGE_DEPTH { + return None; + } let graph = self.graph()?; let node = graph.nodes.get(node_index)?; if graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { @@ -459,22 +478,23 @@ impl CodexLineagePlanner { return None; } - if super::codex_usage_uses_parent(usage) { - let inherited = usage - .codex_fork_accounting_state - .as_ref()? - .inherited_totals - .as_ref()?; + let fork_origin = if super::codex_usage_uses_parent(usage) { + let state = usage.codex_fork_accounting_state.as_ref()?; + let inherited = state.inherited_totals.as_ref()?; let parent_index = graph.parent_indices[node_index]?; - let baseline = self.parent_owner_baseline( + let baseline = self.parent_owner_baseline_at_depth( cache, parent_index, usage.codex_fork_timestamp.as_deref(), + depth + 1, )?; if &baseline != inherited { return None; } - } + Some((state, inherited)) + } else { + None + }; let metadata = fs::metadata(&node.path).ok()?; let expected_identity = usage.codex_file_identity.as_ref()?; @@ -490,11 +510,28 @@ impl CodexLineagePlanner { { 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) + if let Some(last_token_timestamp) = usage.codex_last_token_timestamp.as_deref() + && JsonlScanner::codex_timestamp_at_or_before( + last_token_timestamp, + child_fork_timestamp, + ) + { + return usage.last_totals.clone(); + } + + // No token snapshot of this fork exists at the descendant's cutoff, so + // its counters are still the origin it inherited when it forked. + let (state, inherited) = fork_origin?; + let forked_before_cutoff = usage.codex_fork_timestamp.as_deref().is_some_and(|forked| { + JsonlScanner::codex_timestamp_at_or_before(forked, child_fork_timestamp) + }); + let first_own_token_after_cutoff = usage.codex_last_token_timestamp.is_none() + || state.first_token_timestamp.as_deref().is_some_and(|first| { + JsonlScanner::codex_timestamp_before(child_fork_timestamp, first) + }); + (forked_before_cutoff && first_own_token_after_cutoff).then(|| inherited.clone()) } } diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 69f0c8dffa..abcc523e66 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -3345,6 +3345,9 @@ fn incomplete_or_buffered_empty_codex_fragment_is_not_marked_complete() { #[path = "tests/copied_prefix.rs"] mod copied_prefix; #[cfg(test)] +#[path = "tests/direct_fork.rs"] +mod direct_fork; +#[cfg(test)] #[path = "tests/lineage_cache.rs"] mod lineage_cache; #[cfg(test)] diff --git a/rust/src/cost_scanner/tests/direct_fork.rs b/rust/src/cost_scanner/tests/direct_fork.rs new file mode 100644 index 0000000000..1b6980ca00 --- /dev/null +++ b/rust/src/cost_scanner/tests/direct_fork.rs @@ -0,0 +1,325 @@ +//! Direct fork chains: `root -> parent (fork) -> child (fork)`. +//! +//! Port of upstream `CodexDirectForkBaselineTests`. The child of a fork must +//! inherit the cumulative-counter origin of the whole chain, including when +//! the intermediate parent has no token snapshot before the child forks. +use super::*; +use serde_json::{Value, json}; + +const CACHE_SCHEMA_BEFORE_DIRECT_FORK_BASELINES: u64 = 4; + +fn session_meta(id: &str, parent: Option<&str>, at: DateTime) -> Value { + let mut payload = json!({"id": id, "source": "vscode", "thread_source": "user"}); + if let Some(parent) = parent { + payload["forked_from_id"] = json!(parent); + } + json!({"type": "session_meta", "timestamp": at.to_rfc3339(), "payload": payload}) +} + +fn token_count(at: DateTime, total_input: i64, last_input: i64) -> Value { + json!({ + "type": "event_msg", "timestamp": at.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": "gpt-5.4", + "total_token_usage": {"input_tokens": total_input, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last_input, "output_tokens": 0}, + }} + }) +} + +fn write_rows(dir: &Path, name: &str, rows: &[Value]) -> PathBuf { + std::fs::create_dir_all(dir).unwrap(); + let path = dir.join(name); + let body = rows + .iter() + .map(|row| format!("{row}\n")) + .collect::(); + std::fs::write(&path, body).unwrap(); + path +} + +struct Chain { + _root: tempfile::TempDir, + sessions: PathBuf, + cache_root: PathBuf, + root_file: PathBuf, + parent_file: PathBuf, + child_file: PathBuf, + base: DateTime, +} + +impl Chain { + fn at(&self, seconds: i64) -> DateTime { + self.base + Duration::seconds(seconds) + } + + fn root_rows(&self, total_input: i64) -> Vec { + vec![ + session_meta("root", None, self.at(0)), + token_count(self.at(1), total_input, total_input), + ] + } +} + +/// `root(1000)` -> `parent` forked at t=2 (optional own token event at +/// `parent_event_time`, total 1040 / last 40) -> `child` forked at t=4. +fn write_chain(parent_event_time: i64) -> Chain { + 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 day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + let mut chain = Chain { + sessions, + cache_root, + root_file: PathBuf::new(), + parent_file: PathBuf::new(), + child_file: PathBuf::new(), + base, + _root: root, + }; + + chain.root_file = write_rows(&day_dir, "root.jsonl", &chain.root_rows(1_000)); + let mut parent_rows = vec![session_meta("parent", Some("root"), chain.at(2))]; + if parent_event_time > 0 { + parent_rows.push(token_count(chain.at(parent_event_time), 1_040, 40)); + } + chain.parent_file = write_rows(&day_dir, "parent.jsonl", &parent_rows); + + let (inherited, copied_last) = if parent_event_time == 3 { + (1_040, 40) + } else { + (1_000, 1_000) + }; + chain.child_file = write_rows( + &day_dir, + "child.jsonl", + &[ + session_meta("child", Some("parent"), chain.at(4)), + token_count(chain.at(5), inherited, copied_last), + token_count(chain.at(6), inherited + 20, 20), + token_count(chain.at(7), inherited + 20, 20), + ], + ); + // Discovery is oldest-first, so make parent-before-child order explicit + // rather than depending on filesystem timestamp resolution. + for (age, file) in [ + (30, &chain.root_file), + (20, &chain.parent_file), + (10, &chain.child_file), + ] { + let modified = std::time::SystemTime::now() - std::time::Duration::from_secs(age); + std::fs::File::options() + .write(true) + .open(file) + .unwrap() + .set_modified(modified) + .unwrap(); + } + chain +} + +fn scanner(sessions: &Path, cache_root: &Path, bounded: bool) -> CostScanner { + let mut options = CostScanOptions::app_driven(); + if bounded { + 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 scan_to_completion(scanner: &CostScanner) -> (CostSummary, CostUsageCache) { + for _ in 0..80 { + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + if !cache.codex_scan_incomplete { + return (summary, cache); + } + } + panic!("bounded Codex scan never completed"); +} + +fn billed_input(cache: &CostUsageCache, path: &Path) -> i64 { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!( + !usage.codex_unresolved_fork_parent, + "{path:?} is unresolved" + ); + usage + .days + .values() + .flat_map(HashMap::values) + .map(|row| row[0]) + .sum() +} + +fn downgrade_cache_schema(cache_root: &Path) { + let path = cache_root.join("cost-usage").join("codex-v1.json"); + let mut cache: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + cache["codex_cache_schema_version"] = json!(CACHE_SCHEMA_BEFORE_DIRECT_FORK_BASELINES); + std::fs::write(&path, serde_json::to_vec(&cache).unwrap()).unwrap(); +} + +fn assert_chain_billing( + cache: &CostUsageCache, + summary: &CostSummary, + chain: &Chain, + parent_event_time: i64, +) { + assert_eq!(billed_input(cache, &chain.child_file), 20); + let expected = 1_020 + if parent_event_time > 0 { 40 } else { 0 }; + assert_eq!(summary.input_tokens, expected); +} + +#[test] +fn direct_fork_chain_preserves_cumulative_inheritance() { + for parent_event_time in [0_i64, 3, 8] { + for bounded in [false, true] { + let chain = write_chain(parent_event_time); + let scanner = scanner(&chain.sessions, &chain.cache_root, bounded); + + // Cold scan, then an unchanged warm scan from the persisted cache. + for _ in 0..2 { + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + } + + // A cache written before this correction is rebuilt, not trusted. + downgrade_cache_schema(&chain.cache_root); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + + // A forced rescan is a cold scan against an empty cache. + let forced_cache_root = chain.cache_root.with_file_name("cache-forced"); + let forced = self::scanner(&chain.sessions, &forced_cache_root, bounded); + let (summary, cache) = scan_to_completion(&forced); + assert_chain_billing(&cache, &summary, &chain, parent_event_time); + } + } +} + +#[test] +fn changed_root_revalidates_empty_parent_descendants() { + let chain = write_chain(0); + let scanner = scanner(&chain.sessions, &chain.cache_root, false); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, 0); + let root_dir = chain.root_file.parent().unwrap(); + + // Root rewritten with the same counters: the grandchild is revalidated + // through the empty parent and keeps its own 20 tokens. + write_rows(root_dir, "root.jsonl", &chain.root_rows(1_000)); + let (summary, cache) = scan_to_completion(&scanner); + assert_chain_billing(&cache, &summary, &chain, 0); + + // Root history rewritten to a counter above what the child replays. The + // child must not keep the baseline it inherited through the empty + // parent; it fails closed instead of billing against a stale origin. + write_rows(root_dir, "root.jsonl", &chain.root_rows(1_010)); + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let child = &cache.files[&chain.child_file.to_string_lossy().to_string()]; + assert!(child.codex_unresolved_fork_parent); + assert!(child.days.is_empty()); + assert_eq!(summary.input_tokens, 1_010); +} + +#[test] +fn cyclic_empty_fork_ancestry_stays_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 day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + let files = [("a", "b"), ("b", "a")].map(|(id, parent)| { + write_rows( + &day_dir, + &format!("{id}.jsonl"), + &[session_meta(id, Some(parent), base)], + ) + }); + let scanner = scanner(&sessions, &cache_root, false); + + for _ in 0..2 { + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 0); + for file in &files { + let usage = &cache.files[&file.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + } + } +} + +/// `root(1000)` -> `depth` empty forks -> `leaf` (20 own tokens). Returns the +/// scan result for the leaf. +fn scan_empty_fork_ladder(depth: usize) -> (CostSummary, CostUsageCache, PathBuf) { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let base = Utc::now() - Duration::hours(1); + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + write_rows( + &day_dir, + "s000.jsonl", + &[ + session_meta("s000", None, base), + token_count(base + Duration::seconds(1), 1_000, 1_000), + ], + ); + for level in 1..=depth { + write_rows( + &day_dir, + &format!("s{level:03}.jsonl"), + &[session_meta( + &format!("s{level:03}"), + Some(&format!("s{:03}", level - 1)), + base + Duration::seconds(2), + )], + ); + } + let leaf = write_rows( + &day_dir, + "leaf.jsonl", + &[ + session_meta( + "leaf", + Some(&format!("s{depth:03}")), + base + Duration::seconds(3), + ), + token_count(base + Duration::seconds(4), 1_000, 1_000), + token_count(base + Duration::seconds(5), 1_020, 20), + ], + ); + let scanner = scanner(&sessions, &root.path().join("cache"), false); + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + (summary, cache, leaf) +} + +#[test] +fn empty_fork_ladder_resolves_within_depth_guard() { + let (summary, cache, leaf) = scan_empty_fork_ladder(8); + assert_eq!(billed_input(&cache, &leaf), 20); + assert_eq!(summary.input_tokens, 1_020); +} + +#[test] +fn empty_fork_ladder_beyond_depth_guard_stays_unresolved() { + let (summary, cache, leaf) = scan_empty_fork_ladder(70); + let usage = &cache.files[&leaf.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + assert_eq!(summary.input_tokens, 1_000); +} From 5d092e5fc56709eeb6b879b11c9ca7bfac31f2e4 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:31:45 +0700 Subject: [PATCH 146/407] Port upstream 0.67.0: add twelve preferred currencies from one shared catalog --- .../src-tauri/src/commands/currency.rs | 16 +- apps/desktop-tauri/src/lib/currency.test.ts | 34 ++++ apps/desktop-tauri/src/lib/currency.ts | 40 ++-- .../src/lib/currencyCatalog.generated.ts | 32 +++ .../settings/tabs/GeneralTab.test.tsx | 9 +- .../src/surfaces/settings/tabs/GeneralTab.tsx | 7 +- rust/src/currency.rs | 191 +++++++++++++++--- 7 files changed, 255 insertions(+), 74 deletions(-) create mode 100644 apps/desktop-tauri/src/lib/currencyCatalog.generated.ts diff --git a/apps/desktop-tauri/src-tauri/src/commands/currency.rs b/apps/desktop-tauri/src-tauri/src/commands/currency.rs index 5eae3cd246..753a544bac 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/currency.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/currency.rs @@ -1,8 +1,7 @@ use std::{collections::HashMap, path::PathBuf}; use codexbar::currency::{ - SUPPORTED_CURRENCY_CODES, convert_amount, fallback_rates, fetch_exchange_rates, - normalize_preferred_currency, + CURRENCIES, convert_amount, fallback_rates, fetch_exchange_rates, normalize_preferred_currency, }; use serde::{Deserialize, Serialize}; use tauri::State; @@ -102,12 +101,12 @@ fn cache_path() -> Option { 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() + for currency in CURRENCIES { + if let Some(rate) = rates.get(currency.code).copied() && rate.is_finite() && rate > 0.0 { - clean.insert((*code).to_string(), rate); + clean.insert(currency.code.to_string(), rate); } } if clean @@ -154,16 +153,15 @@ fn persist_rates(cached: &PersistedRates) { #[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_eq!(rates.len(), CURRENCIES.len()); + for currency in CURRENCIES { assert!( rates - .get(*code) + .get(currency.code) .is_some_and(|rate| rate.is_finite() && *rate > 0.0) ); } diff --git a/apps/desktop-tauri/src/lib/currency.test.ts b/apps/desktop-tauri/src/lib/currency.test.ts index 96a05d5879..3a6e8de056 100644 --- a/apps/desktop-tauri/src/lib/currency.test.ts +++ b/apps/desktop-tauri/src/lib/currency.test.ts @@ -1,6 +1,8 @@ import { describe, expect, it } from "vitest"; import { + CURRENCY_PICKER_OPTIONS, FALLBACK_CURRENCY_RATES, + SUPPORTED_CURRENCIES, convertCurrencyAmount, formatDisplayCurrency, mergeValidCurrencyRates, @@ -23,6 +25,38 @@ describe("preferred currency display", () => { expect(display).toMatch(/485/); }); + it("lists the twelve added currencies after TRY with upstream fallback rates", () => { + const added = ["NZD", "SEK", "NOK", "DKK", "PLN", "BRL", "MXN", "ZAR", "THB", "IDR", "VND", "UAH"]; + const tryIndex = SUPPORTED_CURRENCIES.indexOf("TRY"); + expect(SUPPORTED_CURRENCIES.slice(tryIndex + 1)).toEqual(added); + expect(CURRENCY_PICKER_OPTIONS.slice(tryIndex + 1).map((option) => option.label)).toEqual([ + "NZD ($)", "SEK (kr)", "NOK (kr)", "DKK (kr)", "PLN (zł)", "BRL (R$)", + "MXN ($)", "ZAR (R)", "THB (฿)", "IDR (Rp)", "VND (₫)", "UAH (₴)", + ]); + expect(FALLBACK_CURRENCY_RATES).toMatchObject({ + NZD: 1.761, SEK: 9.908, NOK: 9.48, DKK: 6.554, PLN: 3.838, BRL: 5.117, + MXN: 17.47, ZAR: 16.36, THB: 33.37, IDR: 17836, VND: 25962, UAH: 44.86, + }); + expect(Object.keys(FALLBACK_CURRENCY_RATES)).toEqual([...SUPPORTED_CURRENCIES]); + for (const code of added) expect(normalizePreferredCurrency(code.toLowerCase())).toBe(code); + }); + + it("converts each added currency through the USD pivot", () => { + for (const code of ["NZD", "SEK", "NOK", "DKK", "PLN", "BRL", "MXN", "ZAR", "THB", "IDR", "VND", "UAH"]) { + const rate = FALLBACK_CURRENCY_RATES[code]; + expect(convertCurrencyAmount(10, "USD", code, FALLBACK_CURRENCY_RATES)).toBeCloseTo(10 * rate, 8); + expect(convertCurrencyAmount(10 * rate, code, "USD", FALLBACK_CURRENCY_RATES)).toBeCloseTo(10, 8); + } + expect(convertCurrencyAmount(10, "EUR", "SEK", FALLBACK_CURRENCY_RATES)).toBeCloseTo((10 / 0.92) * 9.908, 8); + }); + + it("formats zero-decimal currencies without fractional units", () => { + const vnd = formatDisplayCurrency(0.01, "USD", "VND", FALLBACK_CURRENCY_RATES); + expect(vnd).toMatch(/260/); + expect(vnd).not.toMatch(/[.,]\d{1,2}\D*$/); + expect(formatDisplayCurrency(1, "USD", "JPY", FALLBACK_CURRENCY_RATES)).toMatch(/154(?![.,]\d)/); + }); + 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"); diff --git a/apps/desktop-tauri/src/lib/currency.ts b/apps/desktop-tauri/src/lib/currency.ts index fb83245eed..110744de74 100644 --- a/apps/desktop-tauri/src/lib/currency.ts +++ b/apps/desktop-tauri/src/lib/currency.ts @@ -1,32 +1,18 @@ -export const SUPPORTED_CURRENCIES = [ - "USD", "GBP", "EUR", "CZK", "CNY", "JPY", "KRW", "CAD", "AUD", "HKD", "TWD", "SGD", - "INR", "CHF", "AED", "TRY", -] as const; +import { CURRENCY_CATALOG } from "./currencyCatalog.generated"; -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, -}; +// The catalog (order, symbols, offline rates) is generated from rust/src/currency.rs. +export const SUPPORTED_CURRENCIES: readonly string[] = CURRENCY_CATALOG.map((entry) => entry.code); + +export const FALLBACK_CURRENCY_RATES: Record = Object.fromEntries( + CURRENCY_CATALOG.map((entry) => [entry.code, entry.fallbackRate]), +); + +export const CURRENCY_PICKER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = + CURRENCY_CATALOG.map((entry) => ({ value: entry.code, label: `${entry.code} (${entry.symbol})` })); 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"; + return code === "AUTO" || SUPPORTED_CURRENCIES.includes(code) ? code : "AUTO"; } export function convertCurrencyAmount( @@ -38,8 +24,7 @@ export function convertCurrencyAmount( 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 (!SUPPORTED_CURRENCIES.includes(source) || !SUPPORTED_CURRENCIES.includes(target)) return null; if (source === target) return amount; const sourceRate = source === "USD" ? 1 : rates[source]; const targetRate = target === "USD" ? 1 : rates[target]; @@ -77,7 +62,6 @@ export function formatDisplayCurrency( return new Intl.NumberFormat(undefined, { style: "currency", currency: preferred, - maximumFractionDigits: 2, }).format(converted); } catch { return `${converted.toFixed(2)} ${preferred}`; diff --git a/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts b/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts new file mode 100644 index 0000000000..e1ec5b1f2a --- /dev/null +++ b/apps/desktop-tauri/src/lib/currencyCatalog.generated.ts @@ -0,0 +1,32 @@ +// Generated from rust/src/currency.rs (CURRENCIES). Do not edit by hand. +// Regenerate: 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 }, + { code: "EUR", symbol: "€", fallbackRate: 0.92 }, + { code: "CZK", symbol: "Kč", fallbackRate: 21 }, + { code: "CNY", symbol: "¥", fallbackRate: 7.27 }, + { code: "JPY", symbol: "¥", fallbackRate: 154 }, + { code: "KRW", symbol: "₩", fallbackRate: 1428.9 }, + { code: "CAD", symbol: "$", fallbackRate: 1.38 }, + { code: "AUD", symbol: "$", fallbackRate: 1.55 }, + { code: "HKD", symbol: "$", fallbackRate: 7.8 }, + { code: "TWD", symbol: "NT$", fallbackRate: 32.3 }, + { code: "SGD", symbol: "$", fallbackRate: 1.34 }, + { code: "INR", symbol: "₹", fallbackRate: 84.5 }, + { code: "CHF", symbol: "Fr.", fallbackRate: 0.8 }, + { code: "AED", symbol: "د.إ", fallbackRate: 3.6725 }, + { code: "TRY", symbol: "₺", fallbackRate: 48.5 }, + { code: "NZD", symbol: "$", fallbackRate: 1.761 }, + { code: "SEK", symbol: "kr", fallbackRate: 9.908 }, + { code: "NOK", symbol: "kr", fallbackRate: 9.48 }, + { code: "DKK", symbol: "kr", fallbackRate: 6.554 }, + { code: "PLN", symbol: "zł", fallbackRate: 3.838 }, + { code: "BRL", symbol: "R$", fallbackRate: 5.117 }, + { code: "MXN", symbol: "$", fallbackRate: 17.47 }, + { code: "ZAR", symbol: "R", fallbackRate: 16.36 }, + { code: "THB", symbol: "฿", fallbackRate: 33.37 }, + { code: "IDR", symbol: "Rp", fallbackRate: 17836 }, + { code: "VND", symbol: "₫", fallbackRate: 25962 }, + { code: "UAH", symbol: "₴", fallbackRate: 44.86 }, +] as const; diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx index d17c651de3..37b77c7d1b 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx @@ -309,11 +309,16 @@ describe("GeneralTab language picker", () => { }); describe("GeneralTab preferred currency picker", () => { - it("offers AUTO and every supported currency with the TRY symbol", () => { + it("offers AUTO then every catalog currency as CODE (symbol) in catalog order", () => { render(); const select = screen.getByLabelText("PreferredCurrencyLabel"); const options = Array.from(select.querySelectorAll("option"), (option) => option.textContent); - expect(options).toEqual(["AUTO", "USD", "GBP", "EUR", "CZK", "CNY", "JPY", "KRW", "CAD", "AUD", "HKD", "TWD", "SGD", "INR", "CHF", "AED", "TRY (₺)"]); + expect(options).toEqual([ + "AUTO", "USD ($)", "GBP (£)", "EUR (€)", "CZK (Kč)", "CNY (¥)", "JPY (¥)", "KRW (₩)", + "CAD ($)", "AUD ($)", "HKD ($)", "TWD (NT$)", "SGD ($)", "INR (₹)", "CHF (Fr.)", + "AED (د.إ)", "TRY (₺)", "NZD ($)", "SEK (kr)", "NOK (kr)", "DKK (kr)", "PLN (zł)", + "BRL (R$)", "MXN ($)", "ZAR (R)", "THB (฿)", "IDR (Rp)", "VND (₫)", "UAH (₴)", + ]); }); it("persists an explicitly selected preferred currency", () => { diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx index 765c9934dc..80391eac54 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.tsx @@ -14,7 +14,7 @@ import type { UsageThresholdOverride, } from "../../../types/bridge"; import type { LocaleKey } from "../../../i18n/keys"; -import { SUPPORTED_CURRENCIES } from "../../../lib/currency"; +import { CURRENCY_PICKER_OPTIONS } from "../../../lib/currency"; import type { TabProps } from "../settingsTabs"; const FALLBACK_LANGUAGE_OPTIONS: LanguageOption[] = [ @@ -274,10 +274,7 @@ export default function GeneralTab({ ariaLabel={t("PreferredCurrencyLabel")} options={[ { value: "AUTO", label: "AUTO" }, - ...SUPPORTED_CURRENCIES.map((code) => ({ - value: code, - label: code === "TRY" ? "TRY (₺)" : code, - })), + ...CURRENCY_PICKER_OPTIONS, ]} onChange={(value) => set({ preferredCurrencyCode: value })} /> diff --git a/rust/src/currency.rs b/rust/src/currency.rs index 8040d9c92b..608dc9978a 100644 --- a/rust/src/currency.rs +++ b/rust/src/currency.rs @@ -3,33 +3,91 @@ use std::collections::HashMap; use std::time::Duration; -pub const SUPPORTED_CURRENCY_CODES: &[&str] = &[ - "USD", "GBP", "EUR", "CZK", "CNY", "JPY", "KRW", "CAD", "AUD", "HKD", "TWD", "SGD", "INR", - "CHF", "AED", "TRY", -]; +/// One entry of the preferred-display-currency catalog. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct CurrencyInfo { + pub code: &'static str, + pub symbol: &'static str, + /// Offline USD-pivot rate used until live rates are available. + pub fallback_rate: f64, +} + +const fn currency(code: &'static str, symbol: &'static str, fallback_rate: f64) -> CurrencyInfo { + CurrencyInfo { + code, + symbol, + fallback_rate, + } +} -pub const FALLBACK_RATES: &[(&str, f64)] = &[ - ("USD", 1.0), - ("GBP", 0.79), - ("EUR", 0.92), - ("CZK", 21.0), - ("CNY", 7.27), - ("JPY", 154.0), - ("KRW", 1428.90), - ("CAD", 1.38), - ("AUD", 1.55), - ("HKD", 7.80), - ("TWD", 32.30), - ("SGD", 1.34), - ("INR", 84.50), - ("CHF", 0.80), - ("AED", 3.6725), - ("TRY", 48.5), +/// 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`. +pub const CURRENCIES: &[CurrencyInfo] = &[ + currency("USD", "$", 1.0), + currency("GBP", "£", 0.79), + currency("EUR", "€", 0.92), + currency("CZK", "Kč", 21.0), + currency("CNY", "¥", 7.27), + currency("JPY", "¥", 154.0), + currency("KRW", "₩", 1428.90), + currency("CAD", "$", 1.38), + currency("AUD", "$", 1.55), + currency("HKD", "$", 7.80), + currency("TWD", "NT$", 32.30), + currency("SGD", "$", 1.34), + currency("INR", "₹", 84.50), + currency("CHF", "Fr.", 0.80), + currency("AED", "د.إ", 3.6725), + currency("TRY", "₺", 48.5), + // Rates below are from open.er-api.com on 2026-09-24. + currency("NZD", "$", 1.761), + currency("SEK", "kr", 9.908), + currency("NOK", "kr", 9.480), + currency("DKK", "kr", 6.554), + currency("PLN", "zł", 3.838), + currency("BRL", "R$", 5.117), + currency("MXN", "$", 17.47), + currency("ZAR", "R", 16.36), + currency("THB", "฿", 33.37), + currency("IDR", "Rp", 17836.0), + currency("VND", "₫", 25962.0), + currency("UAH", "₴", 44.86), ]; +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 { + 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\ + export const CURRENCY_CATALOG = [\n", + ); + for currency in CURRENCIES { + out.push_str(&format!( + " {{ code: \"{}\", symbol: \"{}\", fallbackRate: {} }},\n", + currency.code, currency.symbol, currency.fallback_rate + )); + } + out.push_str("] as const;\n"); + out +} + pub fn normalize_preferred_currency(value: &str) -> String { let code = value.trim().to_ascii_uppercase(); - if code == "AUTO" || SUPPORTED_CURRENCY_CODES.contains(&code.as_str()) { + if code == "AUTO" || is_supported_currency(&code) { code } else { "AUTO".to_string() @@ -37,9 +95,9 @@ pub fn normalize_preferred_currency(value: &str) -> String { } pub fn fallback_rates() -> HashMap { - FALLBACK_RATES + CURRENCIES .iter() - .map(|(code, rate)| ((*code).to_string(), *rate)) + .map(|currency| (currency.code.to_string(), currency.fallback_rate)) .collect() } @@ -52,12 +110,12 @@ pub fn parse_exchange_rates(payload: &[u8]) -> Option> { } let rates = value.get("rates")?.as_object()?; let mut parsed = HashMap::new(); - for code in SUPPORTED_CURRENCY_CODES { - let Some(rate) = rates.get(*code).and_then(serde_json::Value::as_f64) else { + for currency in CURRENCIES { + let Some(rate) = rates.get(currency.code).and_then(serde_json::Value::as_f64) else { continue; }; if rate.is_finite() && rate > 0.0 { - parsed.insert((*code).to_string(), rate); + parsed.insert(currency.code.to_string(), rate); } } if parsed @@ -82,9 +140,7 @@ pub fn convert_amount( } let source = source_code.trim().to_ascii_uppercase(); let target = target_code.trim().to_ascii_uppercase(); - if !SUPPORTED_CURRENCY_CODES.contains(&source.as_str()) - || !SUPPORTED_CURRENCY_CODES.contains(&target.as_str()) - { + if !is_supported_currency(&source) || !is_supported_currency(&target) { return None; } if source == target { @@ -184,4 +240,79 @@ mod tests { malformed.insert("TRY".into(), f64::NAN); assert_eq!(convert_amount(12.0, "USD", "TRY", &malformed), None); } + + const UPSTREAM_ADDED: &[(&str, &str, f64)] = &[ + ("NZD", "$", 1.761), + ("SEK", "kr", 9.908), + ("NOK", "kr", 9.480), + ("DKK", "kr", 6.554), + ("PLN", "zł", 3.838), + ("BRL", "R$", 5.117), + ("MXN", "$", 17.47), + ("ZAR", "R", 16.36), + ("THB", "฿", 33.37), + ("IDR", "Rp", 17836.0), + ("VND", "₫", 25962.0), + ("UAH", "₴", 44.86), + ]; + + #[test] + fn added_currencies_follow_try_in_upstream_order_with_upstream_rates() { + let codes: Vec<&str> = CURRENCIES.iter().map(|currency| currency.code).collect(); + let try_index = codes.iter().position(|code| *code == "TRY").unwrap(); + assert_eq!(codes.len(), try_index + 1 + UPSTREAM_ADDED.len()); + for (offset, (code, symbol, rate)) in UPSTREAM_ADDED.iter().enumerate() { + let entry = &CURRENCIES[try_index + 1 + offset]; + assert_eq!( + (entry.code, entry.symbol, entry.fallback_rate), + (*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] + fn added_currencies_convert_through_the_usd_pivot_and_normalize() { + let rates = fallback_rates(); + assert_eq!(rates.len(), CURRENCIES.len()); + for (code, _, rate) in UPSTREAM_ADDED { + assert_eq!(normalize_preferred_currency(&code.to_lowercase()), *code); + let from_usd = convert_amount(10.0, "USD", code, &rates).unwrap(); + assert!((from_usd - 10.0 * rate).abs() < 1e-9, "{code}"); + let back = convert_amount(from_usd, code, "USD", &rates).unwrap(); + assert!((back - 10.0).abs() < 1e-9, "{code}"); + } + let cross = convert_amount(10.0, "EUR", "SEK", &rates).unwrap(); + assert!((cross - (10.0 / 0.92) * 9.908).abs() < 1e-9); + } + + #[test] + fn live_rate_parsing_keeps_added_currencies() { + let rates = parse_exchange_rates( + br#"{"result":"success","base_code":"USD","rates":{"USD":1,"VND":26000,"IDR":17900,"UAH":45,"XXX":3}}"#, + ) + .unwrap(); + assert_eq!(rates.get("VND"), Some(&26000.0)); + assert_eq!(rates.get("IDR"), Some(&17900.0)); + assert_eq!(rates.get("UAH"), Some(&45.0)); + assert!(!rates.contains_key("XXX")); + } + + #[test] + fn currency_catalog_typescript_matches_the_checked_in_module() { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../apps/desktop-tauri/src/lib/currencyCatalog.generated.ts"); + let expected = render_typescript_catalog(); + if std::env::var_os("UPDATE_CURRENCY_CATALOG").is_some() { + std::fs::write(&path, &expected).unwrap(); + } + let actual = std::fs::read_to_string(&path).unwrap_or_default(); + assert_eq!( + actual.replace("\r\n", "\n"), + expected, + "regenerate with UPDATE_CURRENCY_CATALOG=1 cargo test -p codexbar currency_catalog" + ); + } } From 67ae0c63acdced100727906969371cefe6120575 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:34:02 +0700 Subject: [PATCH 147/407] Port upstream 0.68.0: add dashboard usage display control --- docs/CLI.md | 2 + rust/src/cli/serve/dashboard/dashboard.html | 58 ++++++++++++-- rust/src/cli/serve/dashboard/html.rs | 87 +++++++++++++++++++++ 3 files changed, 142 insertions(+), 5 deletions(-) diff --git a/docs/CLI.md b/docs/CLI.md index a9c25bcdcc..063a48c2c1 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -91,6 +91,8 @@ codexbar serve --port 8080 Typical endpoints: `/health`, `/usage`, `/cost`, and `/dashboard/v1/snapshot`. Loopback default keeps local use simple; treat non-loopback as a threat-model choice because the token for protected requests crosses the network over HTTP. +The dashboard page at `/` has a "Usage display" control (Follow server, Used, Remaining). Follow server uses `host.usageBarsShowUsed` from the snapshot; the other two override it for this browser only, stored in `localStorage` under `codexbar.dashboard.usageDisplay`. Changing it re-renders the cached snapshot without refetching. Bar colour always follows consumption (70% used warns, 90% used is full), whichever value is shown. + Pass `--metrics` to enable the Prometheus text endpoint at `/metrics`; it returns `404` when the flag is absent. The endpoint uses the same Host allowlist, Bearer token, snapshot cache, and single-flight collection as the dashboard snapshot. A scrape never waits for provider I/O: it returns the last successful snapshot while an expired value refreshes in the background, or `codexbar_up 0` until the first collection succeeds. ```powershell diff --git a/rust/src/cli/serve/dashboard/dashboard.html b/rust/src/cli/serve/dashboard/dashboard.html index 73e117b756..16360f7e59 100644 --- a/rust/src/cli/serve/dashboard/dashboard.html +++ b/rust/src/cli/serve/dashboard/dashboard.html @@ -9,10 +9,12 @@