diff --git a/src/credentials/mod.rs b/src/credentials/mod.rs index 6caddc5..fc48896 100644 --- a/src/credentials/mod.rs +++ b/src/credentials/mod.rs @@ -265,24 +265,55 @@ pub fn load_gemini() -> Option { }) } -/// The Grok / xAI token from the OS keyring (cairn-code device OAuth / API key) -/// or the Grok CLI's `~/.grok/auth.json`. +/// The Grok / xAI token from the Grok CLI's `~/.grok/auth.json` or the OS +/// keyring (cairn-code device OAuth / API key), whichever still holds a live +/// session. +/// +/// Order alone is not enough here. Only the Grok CLI refreshes its own file, +/// and the refresh probe is what drives it; nothing rewrites the `cairn-code` +/// keyring copy. A stale keyring entry that shadowed the file would therefore +/// stay stale forever: every call would see an expired token, run the probe, +/// refresh the file it then ignored, and hand back the same dead token. pub fn load_grok() -> Option { + pick_live(grok_candidates()) +} + +/// The first still-live credential, or the canonical one when none are. +/// +/// Falling back to the head of the list rather than `None` keeps an expired +/// token in play so the caller's refresh probe still has something to revive. +fn pick_live(candidates: Vec) -> Option { + match candidates.iter().find(|token| token.is_fresh()) { + Some(fresh) => Some(fresh.clone()), + None => candidates.into_iter().next(), + } +} + +/// Every Grok credential this machine holds, canonical store first. +fn grok_candidates() -> Vec { + let mut found = Vec::new(); + + if let Some(root) = read_json(&home_dir().join(".grok").join("auth.json")) + && let Some(token) = parse_grok(&root) + { + found.push(token); + } + for (service, account) in [ ("cairn-code", "oauth:xai"), ("cairn-code", "xai"), ("grok", "auth"), ("grok", "oauth:grok"), ] { - if let Some(data) = keyring::read(service, account) { - if let Ok(value) = serde_json::from_slice::(&data) { - if let Some(token) = parse_grok(&value) { - return Some(token); - } - } else { + let Some(data) = keyring::read(service, account) else { + continue; + }; + match serde_json::from_slice::(&data) { + Ok(value) => found.extend(parse_grok(&value)), + Err(_) => { let key = String::from_utf8_lossy(&data).trim().to_string(); if !key.is_empty() && !key.starts_with('{') { - return Some(Token { + found.push(Token { access_token: key, auth_method: "api_key".into(), ..Default::default() @@ -292,8 +323,7 @@ pub fn load_grok() -> Option { } } - let root = read_json(&home_dir().join(".grok").join("auth.json"))?; - parse_grok(&root) + found } pub(crate) fn parse_grok(root: &Value) -> Option { @@ -686,6 +716,48 @@ mod tests { assert_eq!(token.auth_method, "oidc"); } + #[test] + fn a_stale_keyring_entry_never_shadows_a_live_credential() { + let live = now_unix() + 3_600; + let stale = Token { + access_token: "expired-keyring-copy".into(), + expires_at: Some(now_unix() - 3_600), + ..Default::default() + }; + let fresh = Token { + access_token: "refreshed-cli-token".into(), + expires_at: Some(live), + ..Default::default() + }; + + // Canonical store first but expired, keyring second and live. + let picked = pick_live(vec![stale.clone(), fresh.clone()]).unwrap(); + assert_eq!(picked.access_token, "refreshed-cli-token"); + + // Live canonical store wins over anything behind it. + let picked = pick_live(vec![fresh.clone(), stale.clone()]).unwrap(); + assert_eq!(picked.access_token, "refreshed-cli-token"); + } + + #[test] + fn all_expired_falls_back_to_the_canonical_store_for_the_probe() { + let older = Token { + access_token: "canonical".into(), + expires_at: Some(now_unix() - 7_200), + ..Default::default() + }; + let newer = Token { + access_token: "keyring".into(), + expires_at: Some(now_unix() - 60), + ..Default::default() + }; + assert_eq!( + pick_live(vec![older, newer]).unwrap().access_token, + "canonical" + ); + assert!(pick_live(Vec::new()).is_none()); + } + #[test] fn parse_claude_oauth_direct_and_nested() { let nested = json!({ diff --git a/src/fetch.rs b/src/fetch.rs index 2bdb345..2e9bff7 100644 --- a/src/fetch.rs +++ b/src/fetch.rs @@ -587,8 +587,9 @@ impl<'a> Fetcher<'a> { }; // Billing is where the percentage lives; the session expiry is only a - // stand-in until it answers. - let mut used_percent = 0.0; + // stand-in until it answers. `None` means it never answered, which is + // not the same as a window that has gone unspent. + let mut used_percent = None; let mut reset = match token.as_ref().and_then(|t| t.expires_at) { Some(expiry) => crate::time::countdown_between(expiry, now_unix()), None => "Active".to_string(), @@ -598,9 +599,7 @@ impl<'a> Fetcher<'a> { "https://cli-chat-proxy.grok.com/v1/billing?format=credits", )) && let Some(config) = billing.get("config") { - if let Some(percent) = config.get("creditUsagePercent").and_then(Value::as_f64) { - used_percent = percent; - } + used_percent = Some(grok_used_percent(config)); if let Some(end) = config .get("currentPeriod") .and_then(|period| period.get("end")) @@ -612,14 +611,15 @@ impl<'a> Fetcher<'a> { } let account = if email.is_empty() { "Active" } else { &email }; + let window = match used_percent { + Some(percent) => { + UsageWindow::new("Weekly", percent).text(format!("{percent:.0}% used")) + } + None => UsageWindow::new("Weekly", 0.0).text("usage unavailable"), + }; ProviderUsage::healthy( Provider::Grok, - vec![ - UsageWindow::new("Weekly", used_percent) - .reset(reset) - .seconds(WEEK) - .text(format!("{used_percent:.0}% used")), - ], + vec![window.reset(reset).seconds(WEEK)], format!("Grok CLI ({account})"), ) } @@ -792,6 +792,19 @@ pub fn fetch_all(http: &dyn HttpClient, configs: &[ProviderConfig]) -> Vec f64 { + config + .get("creditUsagePercent") + .and_then(Value::as_f64) + .unwrap_or(0.0) +} + #[cfg(test)] mod tests { use super::*; @@ -862,6 +875,51 @@ mod tests { "monthly":{"status":"ok","percent":50,"resetsAt":"2099-09-12T22:42:28.112Z"} }}"#; + const GROK_USER: &str = r#"{"email":"grokuser@example.com"}"#; + + #[test] + fn an_omitted_credit_percent_is_a_real_zero_not_missing_data() { + // Protobuf JSON drops a field sitting at its default, so a billing + // period with no spend answers without `creditUsagePercent` at all. + let unspent = serde_json::json!({ + "currentPeriod": {"type": "USAGE_PERIOD_TYPE_WEEKLY"}, + "onDemandUsed": {"val": 0} + }); + assert_eq!(grok_used_percent(&unspent), 0.0); + + let spent = serde_json::json!({"creditUsagePercent": 88.0}); + assert_eq!(grok_used_percent(&spent), 88.0); + } + + #[test] + fn grok_billing_that_never_answers_reports_unavailable_rather_than_zero() { + // Only the user call is routed; billing gets no route and so errors. + let http = FakeHttp::new(vec![("v1/user", 200, GROK_USER)]); + let usage = Fetcher::new(&http).fetch(&keyed("grok", "xai-test-key")); + + assert_eq!(usage.status, crate::model::Status::Healthy); + let window = &usage.windows[0]; + assert_eq!(window.label, "Weekly"); + assert_eq!( + window.percent_text(), + "usage unavailable", + "a billing call that never landed must not read as 0% used" + ); + } + + #[test] + fn grok_reports_the_billing_percentage_when_it_answers() { + let billing = r#"{"config":{"creditUsagePercent":88.0}}"#; + let http = FakeHttp::new(vec![ + ("v1/billing", 200, billing), + ("v1/user", 200, GROK_USER), + ]); + let usage = Fetcher::new(&http).fetch(&keyed("grok", "xai-test-key")); + + assert_eq!(usage.windows[0].percent_text(), "88% used"); + assert_eq!(usage.windows[0].used_percent, 88.0); + } + #[test] fn opencode_reports_all_three_windows_and_names_the_spent_one() { let http = FakeHttp::new(vec![("zen/go/v1/usage", 200, OPENCODE_USAGE)]);