diff --git a/rust/src/providers/claude/accounts/login.rs b/rust/src/providers/claude/accounts/login.rs index a3f64d00f4..9a7625cc45 100644 --- a/rust/src/providers/claude/accounts/login.rs +++ b/rust/src/providers/claude/accounts/login.rs @@ -416,10 +416,19 @@ mod tests { assert!(!removable.exists()); } + #[cfg(windows)] + fn make_junction(link: &Path, target: &Path, failure: &str) { + use std::os::windows::process::CommandExt; + let output = Command::new("powershell.exe") + .args(["-NoProfile", "-NonInteractive", "-Command", "New-Item -ItemType Junction -Path $env:CODEXBAR_TEST_LINK -Target $env:CODEXBAR_TEST_TARGET | Out-Null"]) + .env("CODEXBAR_TEST_LINK", link).env("CODEXBAR_TEST_TARGET", target) + .creation_flags(0x0800_0000).output().unwrap(); + assert!(output.status.success(), "{failure}"); + } + #[cfg(windows)] #[test] fn startup_cleanup_does_not_follow_junctions() { - use std::os::windows::process::CommandExt; let dir = tempfile::tempdir().unwrap(); let root = dir.path().join("logins"); let outside = dir.path().join("outside"); @@ -427,13 +436,10 @@ mod tests { std::fs::create_dir_all(&outside).unwrap(); std::fs::write(outside.join("keep.txt"), "preserved").unwrap(); let link = root.join(uuid::Uuid::new_v4().to_string()); - let output = Command::new("powershell.exe") - .args(["-NoProfile", "-NonInteractive", "-Command", "New-Item -ItemType Junction -Path $env:CODEXBAR_TEST_LINK -Target $env:CODEXBAR_TEST_TARGET | Out-Null"]) - .env("CODEXBAR_TEST_LINK", &link).env("CODEXBAR_TEST_TARGET", &outside) - .creation_flags(0x0800_0000).output().unwrap(); - assert!( - output.status.success(), - "Failed to create the isolated test junction." + make_junction( + &link, + &outside, + "Failed to create the isolated test junction.", ); cleanup_login_root(&root).unwrap(); assert!(outside.join("keep.txt").exists()); @@ -730,18 +736,14 @@ mod tests { #[cfg(windows)] #[test] fn ordinary_junction_is_not_treated_as_wsl_backed() { - use std::os::windows::process::CommandExt; let dir = tempfile::tempdir().unwrap(); let target = dir.path().join("target"); let link = dir.path().join("link"); std::fs::create_dir_all(&target).unwrap(); - let output = Command::new("powershell.exe") - .args(["-NoProfile", "-NonInteractive", "-Command", "New-Item -ItemType Junction -Path $env:CODEXBAR_TEST_LINK -Target $env:CODEXBAR_TEST_TARGET | Out-Null"]) - .env("CODEXBAR_TEST_LINK", &link).env("CODEXBAR_TEST_TARGET", &target) - .creation_flags(0x0800_0000).output().unwrap(); - assert!( - output.status.success(), - "Failed to create the local junction fixture." + make_junction( + &link, + &target, + "Failed to create the local junction fixture.", ); assert!(!path_is_wsl_backed(&link)); assert!(!path_is_wsl_backed(&target)); diff --git a/rust/src/providers/claude/admin_api.rs b/rust/src/providers/claude/admin_api.rs index c16b2748fc..d5ebb52d27 100644 --- a/rust/src/providers/claude/admin_api.rs +++ b/rust/src/providers/claude/admin_api.rs @@ -209,6 +209,18 @@ struct CacheCreation { total_input_tokens: Option, } +impl MessageResult { + fn input_tokens(&self) -> i64 { + self.uncached_input_tokens.unwrap_or(0) + + self + .cache_creation + .as_ref() + .and_then(|c| c.total_input_tokens) + .unwrap_or(0) + + self.cache_read_input_tokens.unwrap_or(0) + } +} + fn result_from_admin_usage( costs: &CostReportResponse, messages: &MessagesUsageResponse, @@ -225,14 +237,7 @@ fn result_from_admin_usage( .data .iter() .flat_map(|bucket| &bucket.results) - .map(|r| { - r.uncached_input_tokens.unwrap_or(0) - + r.cache_creation - .as_ref() - .and_then(|c| c.total_input_tokens) - .unwrap_or(0) - + r.cache_read_input_tokens.unwrap_or(0) - }) + .map(MessageResult::input_tokens) .sum(); let output_tokens: i64 = messages .data @@ -278,26 +283,13 @@ fn result_from_admin_usage( .with_login_method("Admin API"); usage.updated_at = now; - let mut model_tokens: HashMap = HashMap::new(); - for result in messages.data.iter().flat_map(|bucket| &bucket.results) { - let name = result - .model - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - .unwrap_or("Claude API"); - let tokens = result.uncached_input_tokens.unwrap_or(0) - + result - .cache_creation - .as_ref() - .and_then(|c| c.total_input_tokens) - .unwrap_or(0) - + result.cache_read_input_tokens.unwrap_or(0) - + result.output_tokens.unwrap_or(0); - *model_tokens.entry(name.to_string()).or_default() += tokens; - } - let mut top_models: Vec<_> = model_tokens.into_iter().collect(); - top_models.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + let top_models = ranked_totals( + messages.data.iter().flat_map(|bucket| &bucket.results), + |result| result.model.as_deref(), + "Claude API", + |result| result.input_tokens() + result.output_tokens.unwrap_or(0), + |a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)), + ); for (idx, (model, tokens)) in top_models.into_iter().take(3).enumerate() { usage = usage.with_extra_rate_window( format!("model-{idx}"), @@ -306,19 +298,18 @@ fn result_from_admin_usage( ); } - let mut cost_items: HashMap = HashMap::new(); - for result in costs.data.iter().flat_map(|bucket| &bucket.results) { - let name = result - .description - .as_deref() - .or(result.cost_type.as_deref()) - .map(str::trim) - .filter(|s| !s.is_empty()) - .unwrap_or("Claude API"); - *cost_items.entry(name.to_string()).or_default() += usd_from_lowest_unit(&result.amount); - } - let mut top_items: Vec<_> = cost_items.into_iter().collect(); - top_items.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + let top_items = ranked_totals( + costs.data.iter().flat_map(|bucket| &bucket.results), + |result| { + result + .description + .as_deref() + .or(result.cost_type.as_deref()) + }, + "Claude API", + |result| usd_from_lowest_unit(&result.amount), + |a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal), + ); for (idx, (item, cost)) in top_items.into_iter().take(3).enumerate() { usage = usage.with_extra_rate_window( format!("cost-{idx}"), @@ -342,21 +333,16 @@ fn result_from_admin_usage( /// 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 { + let ranked = ranked_totals( + costs.data.iter().flat_map(|bucket| &bucket.results), + |result| result.workspace_id.as_deref(), + "Default", + |result| usd_from_lowest_unit(&result.amount), + |a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)), + ); + if ranked.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) @@ -372,6 +358,27 @@ fn workspace_spend_rows(costs: &CostReportResponse) -> Vec( + rows: impl Iterator, + name: impl Fn(&'a R) -> Option<&'a str>, + fallback: &str, + value: impl Fn(&R) -> V, + order: impl FnMut(&(String, V), &(String, V)) -> std::cmp::Ordering, +) -> Vec<(String, V)> { + let mut totals: HashMap = HashMap::new(); + for row in rows { + let key = name(row) + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or(fallback); + *totals.entry(key.to_string()).or_default() += value(row); + } + let mut ranked: Vec<_> = totals.into_iter().collect(); + ranked.sort_by(order); + ranked +} + fn usd_from_lowest_unit(raw: &str) -> f64 { raw.parse::().unwrap_or(0.0) / 100.0 } diff --git a/rust/src/providers/claude/admin_api_tests.rs b/rust/src/providers/claude/admin_api_tests.rs index 3a5391c2fb..737cdc6b51 100644 --- a/rust/src/providers/claude/admin_api_tests.rs +++ b/rust/src/providers/claude/admin_api_tests.rs @@ -152,3 +152,70 @@ fn workspace_option_groups_the_existing_cost_request_only() { cost_report_query(&start, &end, false).len() + 1 ); } + +#[test] +fn messages_and_costs_fold_into_ranked_windows() { + let messages: MessagesUsageResponse = serde_json::from_str( + r#"{"data":[ + {"starting_at":"2026-09-23T00:00:00Z","ending_at":"2026-09-24T00:00:00Z","results":[ + {"model":"claude-b","uncached_input_tokens":10,"cache_read_input_tokens":5,"output_tokens":5}, + {"model":" ","uncached_input_tokens":1,"cache_creation":{"total_input_tokens":2}}, + {"model":"claude-a","uncached_input_tokens":20}]}, + {"starting_at":"2026-09-24T00:00:00Z","ending_at":"2026-09-25T00:00:00Z","results":[ + {"model":"claude-c","output_tokens":7}, + {"uncached_input_tokens":4}]}]}"#, + ) + .unwrap(); + let report = costs( + r#"{"data":[ + {"starting_at":"2026-09-22T00:00:00Z","ending_at":"2026-09-23T00:00:00Z","results":[ + {"amount":"300","description":" Web search "}, + {"amount":"100","cost_type":"tokens"}, + {"amount":"200","description":"","cost_type":"session"}, + {"amount":"50"}]}]}"#, + ); + let result = result_from_admin_usage(&report, &messages, now(), false); + let usage = &result.usage; + let windows: Vec<_> = usage + .extra_rate_windows + .iter() + .map(|window| { + ( + window.id.as_str(), + window.title.as_str(), + window.window.reset_description.as_deref().unwrap(), + ) + }) + .collect(); + assert_eq!( + windows, + [ + ("input-tokens", "Input tokens", "42"), + ("output-tokens", "Output tokens", "12"), + ("model-0", "Model: claude-a", "20 tokens"), + ("model-1", "Model: claude-b", "20 tokens"), + ("model-2", "Model: Claude API", "7 tokens"), + ("cost-0", "Cost: Web search", "$3.00"), + ("cost-1", "Cost: Claude API", "$2.50"), + ("cost-2", "Cost: tokens", "$1.00"), + ] + ); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("$6.50 over last 30 days") + ); + assert_eq!( + usage + .secondary + .as_ref() + .unwrap() + .reset_description + .as_deref(), + Some("54 tokens") + ); + assert_eq!( + usage.primary.resets_at.map(|at| at.to_rfc3339()).as_deref(), + Some("2026-09-22T00:00:00+00:00") + ); + assert!(result.display_details().is_empty()); +} diff --git a/rust/src/providers/claude/claude_swap.rs b/rust/src/providers/claude/claude_swap.rs index 4fc13b9e63..ecfe3734fb 100644 --- a/rust/src/providers/claude/claude_swap.rs +++ b/rust/src/providers/claude/claude_swap.rs @@ -24,17 +24,13 @@ mod sanitize; use chrono::{DateTime, Utc}; use serde::Serialize; -pub use parser::{parse_account_list, parse_switch_result, validate_switch_target}; +pub use parser::parse_account_list; pub use projection::{ ClaudeSwapAccount, ClaudeSwapAccountAction, ClaudeSwapHistoricalUsageDto, ClaudeSwapScopedWindowDto, ClaudeSwapSpendWindowDto, ClaudeSwapUsageWindowDto, - HISTORICAL_USAGE_PROVENANCE, action_for_account, project_accounts, + action_for_account, project_accounts, }; -pub use runner::{ - DEFAULT_TIMEOUT, MAX_OUTPUT_BYTES, SWITCH_TIMEOUT, list_arguments, read_account_list, - read_account_list_with_timeout, resolve_executable_path, switch_account, switch_arguments, -}; -pub use sanitize::{MAX_DIAGNOSTIC_CHARS, MAX_LABEL_CHARS, sanitize_display}; +pub use runner::{read_account_list, switch_account}; #[derive(Debug, thiserror::Error)] pub enum ClaudeSwapError { diff --git a/rust/src/providers/claude/claude_swap/parser.rs b/rust/src/providers/claude/claude_swap/parser.rs index 5ff813871c..45ba70ca69 100644 --- a/rust/src/providers/claude/claude_swap/parser.rs +++ b/rust/src/providers/claude/claude_swap/parser.rs @@ -211,10 +211,12 @@ fn parse_row( }) } -/// Strictly parse the schema-v1 `cswap --list --json` envelope. -pub fn parse_account_list(raw: &str) -> Result { - let value: Value = serde_json::from_str(raw).map_err(|_| ClaudeSwapError::NotJsonObject)?; - let object = value.as_object().ok_or(ClaudeSwapError::NotJsonObject)?; +/// Parse a schema-v1 object and surface a reported `error` envelope before +/// any command-specific field is read. +fn parse_envelope(raw: &str) -> Result, ClaudeSwapError> { + let Ok(Value::Object(object)) = serde_json::from_str::(raw) else { + return Err(ClaudeSwapError::NotJsonObject); + }; let schema_version = object .get("schemaVersion") @@ -248,6 +250,12 @@ pub fn parse_account_list(raw: &str) -> Result Result { + let object = parse_envelope(raw)?; let raw_accounts = object .get("accounts") @@ -318,41 +326,7 @@ pub fn parse_account_list(raw: &str) -> Result --json` envelope. pub fn parse_switch_result(raw: &str) -> Result { - let value: Value = serde_json::from_str(raw).map_err(|_| ClaudeSwapError::NotJsonObject)?; - let object = value.as_object().ok_or(ClaudeSwapError::NotJsonObject)?; - - let schema_version = object - .get("schemaVersion") - .and_then(Value::as_i64) - .ok_or(ClaudeSwapError::MissingSchemaVersion)?; - if schema_version != 1 { - return Err(ClaudeSwapError::UnsupportedSchemaVersion(schema_version)); - } - if let Some(error) = object.get("error").and_then(Value::as_object) { - let kind = sanitize_display( - error.get("type").and_then(Value::as_str).unwrap_or("Error"), - MAX_LABEL_CHARS, - ); - let message = sanitize_display( - error - .get("message") - .and_then(Value::as_str) - .unwrap_or("unknown error"), - MAX_DIAGNOSTIC_CHARS, - ); - return Err(ClaudeSwapError::ReportedError { - kind: if kind.is_empty() { - "Error".to_string() - } else { - kind - }, - message: if message.is_empty() { - "unknown error".to_string() - } else { - message - }, - }); - } + let object = parse_envelope(raw)?; let switched = object .get("switched") @@ -428,376 +402,5 @@ pub fn validate_switch_target( } #[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - fn list_fixture() -> Value { - json!({ - "schemaVersion": 1, - "activeAccountNumber": 2, - "accounts": [ - { - "number": 1, - "email": "same@example.com", - "organizationName": "Work", - "active": false, - "usageStatus": "ok", - "usage": { - "fiveHour": { "pct": 120.0, "resetsAt": "2026-09-12T01:00:00Z" }, - "sevenDay": { "pct": 18.0 }, - "scoped": [{ "name": "Fable only", "pct": 4.0 }] - }, - "usageFetchedAt": "2026-09-12T00:30:00.000Z" - }, - { - "number": 2, - "email": "same@example.com", - "organizationName": "Personal", - "active": true, - "usageStatus": "ok", - "usage": { "fiveHour": { "pct": 81.0 }, "sevenDay": { "pct": 18.0 } } - }, - { - "number": 3, - "email": "expired@example.com", - "organizationName": "", - "alias": "Backup", - "active": false, - "usageStatus": "token_expired" - } - ] - }) - } - - #[test] - fn parses_schema_v1_and_normalizes_windows() { - let parsed = parse_account_list(&list_fixture().to_string()).unwrap(); - assert_eq!(parsed.active_account_number, Some(2)); - assert_eq!(parsed.accounts.len(), 3); - let first = &parsed.accounts[0]; - assert_eq!(first.number, 1); - assert_eq!(first.organization_name, "Work"); - // Out-of-range percentages are clamped like upstream. - assert_eq!(first.usage.five_hour.as_ref().unwrap().used_percent, 100.0); - assert!(first.usage.five_hour.as_ref().unwrap().resets_at.is_some()); - assert_eq!(first.usage.scoped[0].name, "Fable only"); - assert_eq!(first.usage_status, ClaudeSwapUsageStatus::Ok); - } - - #[test] - fn switching_capability_defaults_to_true_and_accepts_booleans() { - let parsed = parse_account_list(&list_fixture().to_string()).unwrap(); - assert!(parsed.supports_account_switching); - for supported in [true, false] { - let mut fixture = list_fixture(); - fixture["supportsAccountSwitching"] = json!(supported); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - assert_eq!(parsed.supports_account_switching, supported); - } - } - - #[test] - fn rejects_non_boolean_switching_capability() { - for value in [ - json!(null), - json!(0), - json!(1), - json!("false"), - json!([]), - json!({}), - ] { - let mut fixture = list_fixture(); - fixture["supportsAccountSwitching"] = value.clone(); - assert!( - matches!( - parse_account_list(&fixture.to_string()), - Err(ClaudeSwapError::MalformedShape(ref message)) - if message == "supportsAccountSwitching is not a boolean" - ), - "{value} must be rejected" - ); - } - } - - #[test] - fn rejects_unknown_and_missing_schema_versions() { - let mut unknown = list_fixture(); - unknown["schemaVersion"] = json!(2); - assert!(matches!( - parse_account_list(&unknown.to_string()), - Err(ClaudeSwapError::UnsupportedSchemaVersion(2)) - )); - - let mut missing = list_fixture(); - missing.as_object_mut().unwrap().remove("schemaVersion"); - assert!(matches!( - parse_account_list(&missing.to_string()), - Err(ClaudeSwapError::MissingSchemaVersion) - )); - - assert!(matches!( - parse_account_list("not json"), - Err(ClaudeSwapError::NotJsonObject) - )); - assert!(matches!( - parse_account_list("[]"), - Err(ClaudeSwapError::NotJsonObject) - )); - } - - #[test] - fn unknown_status_is_not_echoed_to_the_row() { - let mut fixture = list_fixture(); - fixture["accounts"][2]["usageStatus"] = json!("super_secret_token\u{1b}]0;leak\u{07}"); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let row = parsed.accounts.iter().find(|a| a.number == 3).unwrap(); - assert_eq!(row.usage_status, ClaudeSwapUsageStatus::Unknown); - assert!(!row.usage_status.as_label().contains("super_secret_token")); - } - - #[test] - fn external_labels_strip_escapes_and_respect_length_bounds() { - let hostile = format!("\u{1b}[31mEvil\u{1b}[0m\n{}", "x".repeat(400)); - let mut fixture = list_fixture(); - fixture["accounts"][2]["alias"] = json!(hostile); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let alias = parsed - .accounts - .iter() - .find(|a| a.number == 3) - .unwrap() - .alias - .as_deref() - .unwrap(); - assert!(!alias.contains('\u{1b}')); - assert!(alias.contains("Evil")); - assert!(alias.chars().count() <= MAX_LABEL_CHARS); - } - - #[test] - fn reported_error_envelope_is_sanitized_and_bounded() { - let raw = json!({ - "schemaVersion": 1, - "error": { - "type": "\u{1b}[31mBad\u{07}", - "message": format!("\u{1b}]0;leak\u{07}{}", "y".repeat(900)) - } - }); - match parse_account_list(&raw.to_string()) { - Err(ClaudeSwapError::ReportedError { kind, message }) => { - assert!(!kind.contains('\u{1b}')); - assert!(!message.contains('\u{1b}')); - assert!(message.chars().count() <= MAX_DIAGNOSTIC_CHARS); - } - other => panic!("expected reported error, got {other:?}"), - } - } - - #[test] - fn surfaces_error_envelope_instead_of_partial_accounts() { - let raw = json!({ - "schemaVersion": 1, - "error": { "type": "LockHeld", "message": "another cswap is running" } - }); - match parse_account_list(&raw.to_string()) { - Err(ClaudeSwapError::ReportedError { kind, message }) => { - assert_eq!(kind, "LockHeld"); - assert!(message.contains("another cswap")); - } - other => panic!("expected reported error, got {other:?}"), - } - } - - #[test] - fn rejects_disagreeing_active_fields_and_duplicate_slots() { - let mut disagree = list_fixture(); - disagree["activeAccountNumber"] = json!(1); - assert!(matches!( - parse_account_list(&disagree.to_string()), - Err(ClaudeSwapError::MalformedShape(_)) - )); - - let mut duplicate = list_fixture(); - { - let accounts = duplicate["accounts"].as_array_mut().unwrap(); - accounts[1]["number"] = json!(1); - accounts[0]["active"] = json!(false); - } - duplicate["activeAccountNumber"] = json!(null); - assert!(matches!( - parse_account_list(&duplicate.to_string()), - Err(ClaudeSwapError::MalformedShape(_)) - )); - } - - #[test] - fn ignores_malformed_scoped_rows_without_dropping_valid_windows() { - let mut fixture = list_fixture(); - fixture["accounts"][0]["usage"]["scoped"] = json!([ - { "name": "Fable only", "pct": 4.0 }, - { "name": "", "pct": 9.0 }, - { "pct": 3.0 }, - "nonsense", - { "name": "Broken reset", "pct": 2.0, "resetsAt": "not-a-date" } - ]); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let scoped = &parsed.accounts[0].usage.scoped; - assert_eq!(scoped.len(), 1); - assert_eq!(scoped[0].name, "Fable only"); - assert!(parsed.accounts[0].usage.five_hour.is_some()); - } - - #[test] - fn parses_spend_disabled_and_last_good_usage_as_additive_data() { - let mut fixture = list_fixture(); - fixture["accounts"][0]["disabled"] = json!(true); - fixture["accounts"][0]["usage"]["spend"] = json!({ - "used": 12.5, - "limit": 50.0, - "pct": 25.0, - "currency": " USD ", - "resetsAt": "2026-09-13T00:00:00Z" - }); - fixture["accounts"][0]["lastGoodUsage"] = json!({ - "fiveHour": { "pct": 44.0 }, - "sevenDay": { "pct": 19.0 }, - "spend": { "used": 8.0, "limit": 40.0, "pct": 20.0, "currency": "EUR" } - }); - fixture["accounts"][0]["lastGoodFetchedAt"] = json!("2026-09-12T00:45:00Z"); - fixture["accounts"][2]["usageStatus"] = json!("foreign_credential"); - - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let first = &parsed.accounts[0]; - assert!(first.is_disabled); - assert_eq!( - first.usage.spend.as_ref().unwrap().currency_code.as_deref(), - Some("USD") - ); - assert_eq!(first.usage.spend.as_ref().unwrap().used, 12.5); - let history = first.historical_usage.as_ref().unwrap(); - assert_eq!( - history.measurement.five_hour.as_ref().unwrap().used_percent, - 44.0 - ); - assert_eq!( - history - .measurement - .spend - .as_ref() - .unwrap() - .currency_code - .as_deref(), - Some("EUR") - ); - assert_eq!(history.fetched_at.to_rfc3339(), "2026-09-12T00:45:00+00:00"); - assert_eq!( - parsed.accounts[2].usage_status, - ClaudeSwapUsageStatus::ForeignCredential - ); - } - - #[test] - fn drops_invalid_additive_history_and_spend_without_dropping_live_usage() { - let mut fixture = list_fixture(); - fixture["accounts"][0]["usage"]["spend"] = json!({ - "used": -1.0, - "limit": 50.0, - "pct": 25.0 - }); - fixture["accounts"][0]["lastGoodUsage"] = json!({ - "fiveHour": { "pct": "not-a-number" }, - "scoped": [{ "name": "valid scope", "pct": 3.0 }] - }); - fixture["accounts"][0]["lastGoodFetchedAt"] = json!("not-a-date"); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let first = &parsed.accounts[0]; - assert!(first.usage.five_hour.is_some()); - assert!(first.usage.spend.is_none()); - assert!(first.historical_usage.is_none()); - - fixture["accounts"][0]["lastGoodFetchedAt"] = json!("2026-09-12T00:45:00Z"); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - let history = parsed.accounts[0].historical_usage.as_ref().unwrap(); - assert!(history.measurement.five_hour.is_none()); - assert_eq!(history.measurement.scoped.len(), 1); - } - - #[test] - fn missing_spend_currency_remains_unknown() { - let mut fixture = list_fixture(); - fixture["accounts"][0]["usage"]["spend"] = json!({ - "used": 2.0, - "limit": 20.0, - "pct": 10.0 - }); - let parsed = parse_account_list(&fixture.to_string()).unwrap(); - assert_eq!( - parsed.accounts[0] - .usage - .spend - .as_ref() - .unwrap() - .currency_code, - None - ); - } - - #[test] - fn switch_result_requires_matching_target_slot() { - let raw = json!({ - "schemaVersion": 1, - "switched": true, - "from": { "number": 2 }, - "to": { "number": 3 }, - "reason": "switched" - }); - let parsed = parse_switch_result(&raw.to_string()).unwrap(); - assert!(parsed.switched); - assert_eq!(parsed.from_account_number, Some(2)); - assert_eq!(parsed.to_account_number, 3); - assert!(validate_switch_target(3, &parsed).is_ok()); - - let wrong = json!({ - "schemaVersion": 1, - "switched": true, - "from": { "number": 1 }, - "to": { "number": 2 }, - "reason": "switched" - }); - let wrong = parse_switch_result(&wrong.to_string()).unwrap(); - assert!(matches!( - validate_switch_target(3, &wrong), - Err(ClaudeSwapError::MismatchedTarget { - expected: 3, - actual: 2 - }) - )); - } - - #[test] - fn switch_result_rejects_missing_reason_and_bad_schema() { - let missing_reason = json!({ - "schemaVersion": 1, - "switched": true, - "from": { "number": 1 }, - "to": { "number": 2 } - }); - assert!(matches!( - parse_switch_result(&missing_reason.to_string()), - Err(ClaudeSwapError::MalformedShape(_)) - )); - - let bad_schema = json!({ - "schemaVersion": 9, - "switched": true, - "from": { "number": 1 }, - "to": { "number": 2 }, - "reason": "switched" - }); - assert!(matches!( - parse_switch_result(&bad_schema.to_string()), - Err(ClaudeSwapError::UnsupportedSchemaVersion(9)) - )); - } -} +#[path = "parser_tests.rs"] +mod tests; diff --git a/rust/src/providers/claude/claude_swap/parser_tests.rs b/rust/src/providers/claude/claude_swap/parser_tests.rs new file mode 100644 index 0000000000..f031bb2651 --- /dev/null +++ b/rust/src/providers/claude/claude_swap/parser_tests.rs @@ -0,0 +1,416 @@ +use super::*; +use serde_json::json; + +fn list_fixture() -> Value { + json!({ + "schemaVersion": 1, + "activeAccountNumber": 2, + "accounts": [ + { + "number": 1, + "email": "same@example.com", + "organizationName": "Work", + "active": false, + "usageStatus": "ok", + "usage": { + "fiveHour": { "pct": 120.0, "resetsAt": "2026-09-12T01:00:00Z" }, + "sevenDay": { "pct": 18.0 }, + "scoped": [{ "name": "Fable only", "pct": 4.0 }] + }, + "usageFetchedAt": "2026-09-12T00:30:00.000Z" + }, + { + "number": 2, + "email": "same@example.com", + "organizationName": "Personal", + "active": true, + "usageStatus": "ok", + "usage": { "fiveHour": { "pct": 81.0 }, "sevenDay": { "pct": 18.0 } } + }, + { + "number": 3, + "email": "expired@example.com", + "organizationName": "", + "alias": "Backup", + "active": false, + "usageStatus": "token_expired" + } + ] + }) +} + +#[test] +fn parses_schema_v1_and_normalizes_windows() { + let parsed = parse_account_list(&list_fixture().to_string()).unwrap(); + assert_eq!(parsed.active_account_number, Some(2)); + assert_eq!(parsed.accounts.len(), 3); + let first = &parsed.accounts[0]; + assert_eq!(first.number, 1); + assert_eq!(first.organization_name, "Work"); + // Out-of-range percentages are clamped like upstream. + assert_eq!(first.usage.five_hour.as_ref().unwrap().used_percent, 100.0); + assert!(first.usage.five_hour.as_ref().unwrap().resets_at.is_some()); + assert_eq!(first.usage.scoped[0].name, "Fable only"); + assert_eq!(first.usage_status, ClaudeSwapUsageStatus::Ok); +} + +#[test] +fn switching_capability_defaults_to_true_and_accepts_booleans() { + let parsed = parse_account_list(&list_fixture().to_string()).unwrap(); + assert!(parsed.supports_account_switching); + for supported in [true, false] { + let mut fixture = list_fixture(); + fixture["supportsAccountSwitching"] = json!(supported); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + assert_eq!(parsed.supports_account_switching, supported); + } +} + +#[test] +fn rejects_non_boolean_switching_capability() { + for value in [ + json!(null), + json!(0), + json!(1), + json!("false"), + json!([]), + json!({}), + ] { + let mut fixture = list_fixture(); + fixture["supportsAccountSwitching"] = value.clone(); + assert!( + matches!( + parse_account_list(&fixture.to_string()), + Err(ClaudeSwapError::MalformedShape(ref message)) + if message == "supportsAccountSwitching is not a boolean" + ), + "{value} must be rejected" + ); + } +} + +#[test] +fn rejects_unknown_and_missing_schema_versions() { + let mut unknown = list_fixture(); + unknown["schemaVersion"] = json!(2); + assert!(matches!( + parse_account_list(&unknown.to_string()), + Err(ClaudeSwapError::UnsupportedSchemaVersion(2)) + )); + + let mut missing = list_fixture(); + missing.as_object_mut().unwrap().remove("schemaVersion"); + assert!(matches!( + parse_account_list(&missing.to_string()), + Err(ClaudeSwapError::MissingSchemaVersion) + )); + + assert!(matches!( + parse_account_list("not json"), + Err(ClaudeSwapError::NotJsonObject) + )); + assert!(matches!( + parse_account_list("[]"), + Err(ClaudeSwapError::NotJsonObject) + )); +} + +#[test] +fn unknown_status_is_not_echoed_to_the_row() { + let mut fixture = list_fixture(); + fixture["accounts"][2]["usageStatus"] = json!("super_secret_token\u{1b}]0;leak\u{07}"); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let row = parsed.accounts.iter().find(|a| a.number == 3).unwrap(); + assert_eq!(row.usage_status, ClaudeSwapUsageStatus::Unknown); + assert!(!row.usage_status.as_label().contains("super_secret_token")); +} + +#[test] +fn external_labels_strip_escapes_and_respect_length_bounds() { + let hostile = format!("\u{1b}[31mEvil\u{1b}[0m\n{}", "x".repeat(400)); + let mut fixture = list_fixture(); + fixture["accounts"][2]["alias"] = json!(hostile); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let alias = parsed + .accounts + .iter() + .find(|a| a.number == 3) + .unwrap() + .alias + .as_deref() + .unwrap(); + assert!(!alias.contains('\u{1b}')); + assert!(alias.contains("Evil")); + assert!(alias.chars().count() <= MAX_LABEL_CHARS); +} + +#[test] +fn reported_error_envelope_is_sanitized_and_bounded() { + let raw = json!({ + "schemaVersion": 1, + "error": { + "type": "\u{1b}[31mBad\u{07}", + "message": format!("\u{1b}]0;leak\u{07}{}", "y".repeat(900)) + } + }); + match parse_account_list(&raw.to_string()) { + Err(ClaudeSwapError::ReportedError { kind, message }) => { + assert!(!kind.contains('\u{1b}')); + assert!(!message.contains('\u{1b}')); + assert!(message.chars().count() <= MAX_DIAGNOSTIC_CHARS); + } + other => panic!("expected reported error, got {other:?}"), + } +} + +#[test] +fn surfaces_error_envelope_instead_of_partial_accounts() { + let raw = json!({ + "schemaVersion": 1, + "error": { "type": "LockHeld", "message": "another cswap is running" } + }); + match parse_account_list(&raw.to_string()) { + Err(ClaudeSwapError::ReportedError { kind, message }) => { + assert_eq!(kind, "LockHeld"); + assert!(message.contains("another cswap")); + } + other => panic!("expected reported error, got {other:?}"), + } +} + +#[test] +fn rejects_disagreeing_active_fields_and_duplicate_slots() { + let mut disagree = list_fixture(); + disagree["activeAccountNumber"] = json!(1); + assert!(matches!( + parse_account_list(&disagree.to_string()), + Err(ClaudeSwapError::MalformedShape(_)) + )); + + let mut duplicate = list_fixture(); + { + let accounts = duplicate["accounts"].as_array_mut().unwrap(); + accounts[1]["number"] = json!(1); + accounts[0]["active"] = json!(false); + } + duplicate["activeAccountNumber"] = json!(null); + assert!(matches!( + parse_account_list(&duplicate.to_string()), + Err(ClaudeSwapError::MalformedShape(_)) + )); +} + +#[test] +fn ignores_malformed_scoped_rows_without_dropping_valid_windows() { + let mut fixture = list_fixture(); + fixture["accounts"][0]["usage"]["scoped"] = json!([ + { "name": "Fable only", "pct": 4.0 }, + { "name": "", "pct": 9.0 }, + { "pct": 3.0 }, + "nonsense", + { "name": "Broken reset", "pct": 2.0, "resetsAt": "not-a-date" } + ]); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let scoped = &parsed.accounts[0].usage.scoped; + assert_eq!(scoped.len(), 1); + assert_eq!(scoped[0].name, "Fable only"); + assert!(parsed.accounts[0].usage.five_hour.is_some()); +} + +#[test] +fn parses_spend_disabled_and_last_good_usage_as_additive_data() { + let mut fixture = list_fixture(); + fixture["accounts"][0]["disabled"] = json!(true); + fixture["accounts"][0]["usage"]["spend"] = json!({ + "used": 12.5, + "limit": 50.0, + "pct": 25.0, + "currency": " USD ", + "resetsAt": "2026-09-13T00:00:00Z" + }); + fixture["accounts"][0]["lastGoodUsage"] = json!({ + "fiveHour": { "pct": 44.0 }, + "sevenDay": { "pct": 19.0 }, + "spend": { "used": 8.0, "limit": 40.0, "pct": 20.0, "currency": "EUR" } + }); + fixture["accounts"][0]["lastGoodFetchedAt"] = json!("2026-09-12T00:45:00Z"); + fixture["accounts"][2]["usageStatus"] = json!("foreign_credential"); + + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let first = &parsed.accounts[0]; + assert!(first.is_disabled); + assert_eq!( + first.usage.spend.as_ref().unwrap().currency_code.as_deref(), + Some("USD") + ); + assert_eq!(first.usage.spend.as_ref().unwrap().used, 12.5); + let history = first.historical_usage.as_ref().unwrap(); + assert_eq!( + history.measurement.five_hour.as_ref().unwrap().used_percent, + 44.0 + ); + assert_eq!( + history + .measurement + .spend + .as_ref() + .unwrap() + .currency_code + .as_deref(), + Some("EUR") + ); + assert_eq!(history.fetched_at.to_rfc3339(), "2026-09-12T00:45:00+00:00"); + assert_eq!( + parsed.accounts[2].usage_status, + ClaudeSwapUsageStatus::ForeignCredential + ); +} + +#[test] +fn drops_invalid_additive_history_and_spend_without_dropping_live_usage() { + let mut fixture = list_fixture(); + fixture["accounts"][0]["usage"]["spend"] = json!({ + "used": -1.0, + "limit": 50.0, + "pct": 25.0 + }); + fixture["accounts"][0]["lastGoodUsage"] = json!({ + "fiveHour": { "pct": "not-a-number" }, + "scoped": [{ "name": "valid scope", "pct": 3.0 }] + }); + fixture["accounts"][0]["lastGoodFetchedAt"] = json!("not-a-date"); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let first = &parsed.accounts[0]; + assert!(first.usage.five_hour.is_some()); + assert!(first.usage.spend.is_none()); + assert!(first.historical_usage.is_none()); + + fixture["accounts"][0]["lastGoodFetchedAt"] = json!("2026-09-12T00:45:00Z"); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + let history = parsed.accounts[0].historical_usage.as_ref().unwrap(); + assert!(history.measurement.five_hour.is_none()); + assert_eq!(history.measurement.scoped.len(), 1); +} + +#[test] +fn missing_spend_currency_remains_unknown() { + let mut fixture = list_fixture(); + fixture["accounts"][0]["usage"]["spend"] = json!({ + "used": 2.0, + "limit": 20.0, + "pct": 10.0 + }); + let parsed = parse_account_list(&fixture.to_string()).unwrap(); + assert_eq!( + parsed.accounts[0] + .usage + .spend + .as_ref() + .unwrap() + .currency_code, + None + ); +} + +#[test] +fn list_and_switch_share_the_envelope_checks() { + const NOT_OBJECT: &str = "claude-swap returned output that is not a JSON object."; + const NO_SCHEMA: &str = "claude-swap output has no schemaVersion field."; + const SCHEMA_2: &str = + "claude-swap output uses unsupported schema version 2; CodexBar supports version 1."; + const UNKNOWN: &str = "claude-swap reported Error: unknown error"; + let rows: [(&str, &str, &str); 12] = [ + ("not json", NOT_OBJECT, NOT_OBJECT), + ("[]", NOT_OBJECT, NOT_OBJECT), + (r#""text""#, NOT_OBJECT, NOT_OBJECT), + ("{}", NO_SCHEMA, NO_SCHEMA), + (r#"{"schemaVersion": "1"}"#, NO_SCHEMA, NO_SCHEMA), + (r#"{"schemaVersion": 1.5}"#, NO_SCHEMA, NO_SCHEMA), + (r#"{"schemaVersion": 2}"#, SCHEMA_2, SCHEMA_2), + ( + r#"{"schemaVersion": 2, "error": {"type": "LockHeld"}}"#, + SCHEMA_2, + SCHEMA_2, + ), + (r#"{"schemaVersion": 1, "error": {}}"#, UNKNOWN, UNKNOWN), + ( + r#"{"schemaVersion": 1, "error": {"type": "", "message": ""}}"#, + UNKNOWN, + UNKNOWN, + ), + ( + r#"{"schemaVersion": 1, "error": {"type": "LockHeld", "message": "busy"}}"#, + "claude-swap reported LockHeld: busy", + "claude-swap reported LockHeld: busy", + ), + ( + r#"{"schemaVersion": 1, "error": "text"}"#, + "claude-swap output is malformed: missing accounts array", + "claude-swap output is malformed: missing switched flag", + ), + ]; + for (raw, list_expected, switch_expected) in rows { + let list = parse_account_list(raw).unwrap_err().to_string(); + let switch = parse_switch_result(raw).unwrap_err().to_string(); + assert_eq!(list, list_expected, "{raw}"); + assert_eq!(switch, switch_expected, "{raw}"); + } +} + +#[test] +fn switch_result_requires_matching_target_slot() { + let raw = json!({ + "schemaVersion": 1, + "switched": true, + "from": { "number": 2 }, + "to": { "number": 3 }, + "reason": "switched" + }); + let parsed = parse_switch_result(&raw.to_string()).unwrap(); + assert!(parsed.switched); + assert_eq!(parsed.from_account_number, Some(2)); + assert_eq!(parsed.to_account_number, 3); + assert!(validate_switch_target(3, &parsed).is_ok()); + + let wrong = json!({ + "schemaVersion": 1, + "switched": true, + "from": { "number": 1 }, + "to": { "number": 2 }, + "reason": "switched" + }); + let wrong = parse_switch_result(&wrong.to_string()).unwrap(); + assert!(matches!( + validate_switch_target(3, &wrong), + Err(ClaudeSwapError::MismatchedTarget { + expected: 3, + actual: 2 + }) + )); +} + +#[test] +fn switch_result_rejects_missing_reason_and_bad_schema() { + let missing_reason = json!({ + "schemaVersion": 1, + "switched": true, + "from": { "number": 1 }, + "to": { "number": 2 } + }); + assert!(matches!( + parse_switch_result(&missing_reason.to_string()), + Err(ClaudeSwapError::MalformedShape(_)) + )); + + let bad_schema = json!({ + "schemaVersion": 9, + "switched": true, + "from": { "number": 1 }, + "to": { "number": 2 }, + "reason": "switched" + }); + assert!(matches!( + parse_switch_result(&bad_schema.to_string()), + Err(ClaudeSwapError::UnsupportedSchemaVersion(9)) + )); +} diff --git a/rust/src/providers/claude/claude_swap/projection.rs b/rust/src/providers/claude/claude_swap/projection.rs index f9382451ff..eaf92779ee 100644 --- a/rust/src/providers/claude/claude_swap/projection.rs +++ b/rust/src/providers/claude/claude_swap/projection.rs @@ -346,6 +346,17 @@ mod tests { parse_account_list(&raw.to_string()).unwrap() } + /// Parse and project a list holding only `account`. + fn project_one(account: serde_json::Value, active: Option) -> ClaudeSwapAccount { + let raw = json!({ + "schemaVersion": 1, + "activeAccountNumber": active, + "accounts": [account] + }); + let parsed = parse_account_list(&raw.to_string()).unwrap(); + project_accounts(&parsed, false).remove(0) + } + #[test] fn same_email_accounts_get_distinct_stable_ids_and_labels() { let projected = project_accounts(&list_fixture(), false); @@ -415,20 +426,16 @@ mod tests { #[test] fn unknown_status_is_neither_echoed_nor_actionable() { - let raw = json!({ - "schemaVersion": 1, - "activeAccountNumber": null, - "accounts": [{ + let account = project_one( + json!({ "number": 1, "email": "x@example.com", "active": false, "usageStatus": "super_secret_token\u{1b}]0;leak\u{07}", "usage": { "fiveHour": { "pct": 1.0 } } - }] - }); - let parsed = parse_account_list(&raw.to_string()).unwrap(); - let projected = project_accounts(&parsed, false); - let account = &projected[0]; + }), + None, + ); assert_eq!(account.status, "unknown"); assert!(account.action.is_none()); let error = account.error.as_deref().unwrap(); @@ -438,18 +445,15 @@ mod tests { #[test] fn foreign_credentials_expose_explicit_reauthentication_action() { - let raw = json!({ - "schemaVersion": 1, - "activeAccountNumber": 1, - "accounts": [{ + let account = project_one( + json!({ "number": 1, "email": "x@example.com", "active": true, "usageStatus": "foreign_credential" - }] - }); - let parsed = parse_account_list(&raw.to_string()).unwrap(); - let account = &project_accounts(&parsed, false)[0]; + }), + Some(1), + ); assert_eq!( account.action, Some(ClaudeSwapAccountAction::Reauthenticate) @@ -465,10 +469,8 @@ mod tests { #[test] fn historical_usage_is_typed_and_marked_as_source_reported() { - let raw = json!({ - "schemaVersion": 1, - "activeAccountNumber": null, - "accounts": [{ + let account = project_one( + json!({ "number": 1, "email": "x@example.com", "active": false, @@ -479,10 +481,9 @@ mod tests { "spend": { "used": 2.0, "limit": 20.0, "pct": 10.0, "currency": "USD" } }, "lastGoodFetchedAt": "2026-09-12T00:45:00Z" - }] - }); - let parsed = parse_account_list(&raw.to_string()).unwrap(); - let account = &project_accounts(&parsed, false)[0]; + }), + None, + ); assert!(account.is_disabled); assert_eq!( account.historical_usage.as_ref().unwrap().provenance, @@ -503,10 +504,8 @@ mod tests { #[test] fn spend_only_ok_usage_is_not_reported_as_empty() { - let raw = json!({ - "schemaVersion": 1, - "activeAccountNumber": null, - "accounts": [{ + let account = project_one( + json!({ "number": 1, "email": "spend@example.com", "active": false, @@ -514,10 +513,9 @@ mod tests { "usage": { "spend": { "used": 2.0, "limit": 20.0, "pct": 10.0 } } - }] - }); - let parsed = parse_account_list(&raw.to_string()).unwrap(); - let account = &project_accounts(&parsed, false)[0]; + }), + None, + ); assert!(account.spend.is_some()); assert!(account.error.is_none()); } diff --git a/rust/src/providers/claude/claude_swap/runner.rs b/rust/src/providers/claude/claude_swap/runner.rs index aa3e2cb519..defa919227 100644 --- a/rust/src/providers/claude/claude_swap/runner.rs +++ b/rust/src/providers/claude/claude_swap/runner.rs @@ -301,8 +301,12 @@ fn wait_bounded_for_child(child: &mut Child) { } } +/// No job containment off Windows; the uninhabited guard keeps one call shape. #[cfg(not(windows))] -fn terminate_child_tree(child: &mut Child) { +type ProcessTreeGuard = std::convert::Infallible; + +#[cfg(not(windows))] +fn terminate_child_tree(child: &mut Child, _tree: Option) { if let Err(error) = child.kill() { tracing::warn!(%error, "failed to kill the direct claude-swap child"); } @@ -356,13 +360,12 @@ fn run_bounded( return Err(ClaudeSwapError::Process(error)); } }; + #[cfg(not(windows))] + let mut process_tree: Option = None; let stdout = match child.stdout.take() { Some(stdout) => stdout, None => { - #[cfg(windows)] terminate_child_tree(&mut child, process_tree.take()); - #[cfg(not(windows))] - terminate_child_tree(&mut child); return Err(ClaudeSwapError::Process( "Failed to capture stdout.".to_string(), )); @@ -371,10 +374,7 @@ fn run_bounded( let stderr = match child.stderr.take() { Some(stderr) => stderr, None => { - #[cfg(windows)] terminate_child_tree(&mut child, process_tree.take()); - #[cfg(not(windows))] - terminate_child_tree(&mut child); return Err(ClaudeSwapError::Process( "Failed to capture stderr.".to_string(), )); @@ -418,10 +418,7 @@ fn run_bounded( Ok(Some(_status)) => child_exited = true, Ok(None) => {} Err(e) => { - #[cfg(windows)] terminate_child_tree(&mut child, process_tree.take()); - #[cfg(not(windows))] - terminate_child_tree(&mut child); return Err(ClaudeSwapError::Process(e.to_string())); } } @@ -432,19 +429,13 @@ fn run_bounded( // A disconnected channel only means failure when a reader ended // without delivering its outcome; after both readers send it is normal. if disconnected && !readers_done { - #[cfg(windows)] terminate_child_tree(&mut child, process_tree.take()); - #[cfg(not(windows))] - terminate_child_tree(&mut child); return Err(ClaudeSwapError::Process( "claude-swap output streams closed unexpectedly.".to_string(), )); } if Instant::now() >= deadline { - #[cfg(windows)] terminate_child_tree(&mut child, process_tree.take()); - #[cfg(not(windows))] - terminate_child_tree(&mut child); return Err(ClaudeSwapError::TimedOut(timeout.as_secs())); } std::thread::sleep(POLL_INTERVAL); diff --git a/rust/src/providers/claude/claude_swap/sanitize.rs b/rust/src/providers/claude/claude_swap/sanitize.rs index 44c4dd6298..6c75d6a91b 100644 --- a/rust/src/providers/claude/claude_swap/sanitize.rs +++ b/rust/src/providers/claude/claude_swap/sanitize.rs @@ -106,81 +106,46 @@ mod tests { use super::*; #[test] - fn strips_terminal_escapes_and_bounds_length() { - let csi = "\u{1b}[31mred\u{1b}[0m"; - assert_eq!(sanitize_display(csi, MAX_LABEL_CHARS), "red"); - - let osc = "a\u{1b}]0;ignored\u{07}b"; - assert_eq!(sanitize_display(osc, MAX_LABEL_CHARS), "ab"); + fn sanitizes_display_text() { + let rows = [ + // Terminal escapes (CSI, OSC). + ("\u{1b}[31mred\u{1b}[0m", "red"), + ("a\u{1b}]0;ignored\u{07}b", "ab"), + // Line breaks and repeated spaces collapse. + ("one\r\ntwo\u{2028}three", "one two three"), + (" a b ", "a b"), + // Bidi controls: RLO ... PDF around reversed text would otherwise render + // as "txet"; then LRM / RLM / ALM / isolates. + ("safe\u{202E}txet\u{202C}", "safetxet"), + ("left\u{200F}right\u{200E}\u{061C}", "leftright"), + ("a\u{2066}b\u{2069}c", "abc"), + ("x\u{202A}y\u{202B}z\u{202D}", "xyz"), + // Default-ignorable code points. + ("co\u{00AD}de", "code"), + ("a\u{200B}\u{200C}\u{200D}b", "ab"), + ("word\u{2060}joiner", "wordjoiner"), + ("\u{FEFF}bom", "bom"), + ("e\u{FE0F}motion", "emotion"), + ("tag\u{E0061}\u{E007F}end", "tagend"), + ("filler\u{3164}text", "fillertext"), + // U+0301 is a combining acute accent (Mn), not default-ignorable. + ("e\u{0301}", "e\u{0301}"), + ]; + for (input, expected) in rows { + assert_eq!( + sanitize_display(input, MAX_LABEL_CHARS), + expected, + "{input:?}" + ); + } + } + #[test] + fn bounds_display_length() { let long = "x".repeat(MAX_LABEL_CHARS + 50); assert_eq!( sanitize_display(&long, MAX_LABEL_CHARS).chars().count(), MAX_LABEL_CHARS ); } - - #[test] - fn collapses_line_breaks_and_repeated_spaces() { - let multiline = "one\r\ntwo\u{2028}three"; - assert_eq!( - sanitize_display(multiline, MAX_LABEL_CHARS), - "one two three" - ); - assert_eq!(sanitize_display(" a b ", MAX_LABEL_CHARS), "a b"); - } - - #[test] - fn strips_bidi_controls_that_could_reorder_text() { - // RLO ... PDF around reversed text would otherwise render as "txet". - assert_eq!( - sanitize_display("safe\u{202E}txet\u{202C}", MAX_LABEL_CHARS), - "safetxet" - ); - // LRM / RLM / ALM / isolates. - assert_eq!( - sanitize_display("left\u{200F}right\u{200E}\u{061C}", MAX_LABEL_CHARS), - "leftright" - ); - assert_eq!( - sanitize_display("a\u{2066}b\u{2069}c", MAX_LABEL_CHARS), - "abc" - ); - assert_eq!( - sanitize_display("x\u{202A}y\u{202B}z\u{202D}", MAX_LABEL_CHARS), - "xyz" - ); - } - - #[test] - fn strips_default_ignorable_code_points() { - assert_eq!(sanitize_display("co\u{00AD}de", MAX_LABEL_CHARS), "code"); - assert_eq!( - sanitize_display("a\u{200B}\u{200C}\u{200D}b", MAX_LABEL_CHARS), - "ab" - ); - assert_eq!( - sanitize_display("word\u{2060}joiner", MAX_LABEL_CHARS), - "wordjoiner" - ); - assert_eq!(sanitize_display("\u{FEFF}bom", MAX_LABEL_CHARS), "bom"); - assert_eq!( - sanitize_display("e\u{FE0F}motion", MAX_LABEL_CHARS), - "emotion" - ); - assert_eq!( - sanitize_display("tag\u{E0061}\u{E007F}end", MAX_LABEL_CHARS), - "tagend" - ); - assert_eq!( - sanitize_display("filler\u{3164}text", MAX_LABEL_CHARS), - "fillertext" - ); - } - - #[test] - fn preserves_ordinary_combining_marks() { - // U+0301 is a combining acute accent (Mn), not default-ignorable. - assert_eq!(sanitize_display("e\u{0301}", MAX_LABEL_CHARS), "e\u{0301}"); - } } diff --git a/rust/src/providers/claude/cli_binary.rs b/rust/src/providers/claude/cli_binary.rs new file mode 100644 index 0000000000..3c5a05a732 --- /dev/null +++ b/rust/src/providers/claude/cli_binary.rs @@ -0,0 +1,116 @@ +#[cfg(windows)] +use std::os::windows::process::CommandExt; +#[cfg(windows)] +use std::process::{Command as StdCommand, Stdio}; + +/// Locate the Claude CLI for shell integrations that need to reopen a session. +pub fn locate_claude_binary() -> Option { + if let Some(path) = std::env::var_os("CLAUDE_BINARY") + .map(std::path::PathBuf::from) + .filter(|path| path.is_file()) + { + return Some(path); + } + + #[cfg(windows)] + { + let candidates = [ + // Direct install + dirs::data_local_dir().map(|p| p.join("Programs").join("claude").join("claude.exe")), + // npm global (AppData\Roaming\npm) + dirs::data_local_dir().map(|p| p.join("npm").join("claude.cmd")), + dirs::home_dir().map(|h| { + h.join("AppData") + .join("Roaming") + .join("npm") + .join("claude.cmd") + }), + // npm global alternate (~\.npm-global) + dirs::home_dir().map(|h| h.join(".npm-global").join("claude.cmd")), + // Volta managed + dirs::data_local_dir().map(|p| { + p.join("Volta") + .join("tools") + .join("image") + .join("packages") + .join("@anthropic-ai") + .join("claude-code") + .join("bin") + .join("claude.cmd") + }), + // fnm managed (via shim) + dirs::data_local_dir().map(|p| p.join("fnm_multishells").join("claude.cmd")), + // PATH lookup + find_windows_claude_in_path(), + ]; + + candidates.into_iter().flatten().find(|p| p.exists()) + } + + #[cfg(not(windows))] + { + which::which("claude").ok() + } +} + +#[cfg(windows)] +fn find_windows_claude_in_path() -> Option { + const CREATE_NO_WINDOW: u32 = 0x08000000; + + let mut command = StdCommand::new("where"); + command + .arg("claude") + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .creation_flags(CREATE_NO_WINDOW); + let output = command.output().ok()?; + + if !output.status.success() { + return None; + } + + let mut matches: Vec<_> = String::from_utf8_lossy(&output.stdout) + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(std::path::PathBuf::from) + .collect(); + + matches.sort_by_key(|path| { + match path + .extension() + .and_then(|ext| ext.to_str()) + .map(|ext| ext.to_ascii_lowercase()) + .as_deref() + { + Some("cmd") => 0, + Some("bat") => 1, + Some("exe") => 2, + _ => 3, + } + }); + + matches.into_iter().find(|path| path.exists()) +} + +/// Detect the version of the claude CLI +pub(super) fn detect_claude_version() -> Option { + let claude_path = locate_claude_binary()?; + + #[cfg(windows)] + const CREATE_NO_WINDOW: u32 = 0x08000000; + + let mut cmd = std::process::Command::new(claude_path); + cmd.args(["--version"]); + #[cfg(windows)] + cmd.creation_flags(CREATE_NO_WINDOW); + + let output = cmd.output().ok()?; + + if output.status.success() { + let version_str = String::from_utf8_lossy(&output.stdout); + super::super::extract_semver(&version_str) + } else { + None + } +} diff --git a/rust/src/providers/claude/cli_probe.rs b/rust/src/providers/claude/cli_probe.rs new file mode 100644 index 0000000000..359f5ec803 --- /dev/null +++ b/rust/src/providers/claude/cli_probe.rs @@ -0,0 +1,568 @@ +use super::cli_binary::locate_claude_binary; +use super::cli_text::is_workspace_trust_prompt; +use super::{ClaudeProvider, accounts, auto_precision, cli_screen, trust_dialog, unix_now_secs}; +use crate::cli::tty_runner::{TtyCommandOptions, TtyCommandRunner}; +use crate::core::ProviderError; +use std::time::{Duration, Instant}; + +const CLAUDE_PROBE_SESSION_ID_FILE: &str = ".codexbar-session-id"; +const CLAUDE_PROBE_LOCK_FILE: &str = ".codexbar-probe.lock"; +pub(super) const CLAUDE_PROBE_CACHE_FILE: &str = ".codexbar-usage-cache.json"; +/// How long a second codexbar process waits for a running probe to finish. +const CLAUDE_PROBE_LOCK_WAIT: Duration = Duration::from_secs(30); +/// Every codexbar process (the `serve` daemon, one-off `usage` calls from the +/// companion) launches its own Claude CLI for a probe. The interactive +/// `/usage` screen costs 6-10 s of CPU each time, so a recent successful +/// probe output is shared across processes for this long. +pub(super) const CLAUDE_PROBE_CACHE_TTL: Duration = Duration::from_secs(45); + +#[derive(serde::Serialize, serde::Deserialize)] +pub(super) struct ClaudeProbeCache { + pub(super) captured_at_unix: u64, + /// `claude_login_fingerprint` of the login the screen belongs to. + #[serde(default)] + pub(super) login: String, + pub(super) output: String, +} + +pub(super) fn load_cached_probe_output(probe_dir: &std::path::Path, login: &str) -> Option { + let raw = std::fs::read_to_string(probe_dir.join(CLAUDE_PROBE_CACHE_FILE)).ok()?; + let cache: ClaudeProbeCache = serde_json::from_str(&raw).ok()?; + let age = unix_now_secs().saturating_sub(cache.captured_at_unix); + if auto_precision::probe_screen_is_superseded(cache.captured_at_unix) + || login.is_empty() + || cache.login != login + || age > CLAUDE_PROBE_CACHE_TTL.as_secs() + || cache.output.trim().is_empty() + { + return None; + } + tracing::debug!(age_secs = age, "Reusing recent Claude CLI probe output"); + Some(cache.output) +} + +pub(super) fn store_cached_probe_output(probe_dir: &std::path::Path, login: &str, output: &str) { + let cache = ClaudeProbeCache { + captured_at_unix: unix_now_secs(), + login: login.to_string(), + output: output.to_string(), + }; + // Atomic, because other processes read the cache without the probe lock. + let stored = serde_json::to_vec(&cache) + .map_err(anyhow::Error::from) + .and_then(|json| { + crate::atomic_file::write_atomic(&probe_dir.join(CLAUDE_PROBE_CACHE_FILE), &json) + }); + if let Err(err) = stored { + tracing::debug!(error = %err, "failed to persist Claude probe cache"); + } +} + +/// Identifies the Claude login a probe runs under without reading any +/// credential: the location, size and modification time of Claude Code's +/// `.credentials.json`, which every login, token refresh and account switch +/// rewrites. Without a credentials file a probe screen is never shared. +fn claude_login_fingerprint() -> Option { + let credentials = accounts::config_dir().ok()?.join(".credentials.json"); + login_fingerprint_at(&credentials) +} + +pub(super) fn login_fingerprint_at(credentials: &std::path::Path) -> Option { + let metadata = std::fs::metadata(credentials).ok()?; + let modified = metadata + .modified() + .ok()? + .duration_since(std::time::UNIX_EPOCH) + .ok()?; + let identity = format!( + "{}|{}|{}", + credentials.display(), + metadata.len(), + modified.as_nanos() + ); + Some(crate::core::sha256_hex(identity.as_bytes())) +} + +/// The probe screen as it may appear in a log: secrets and email addresses +/// (the screen can show the signed-in account) are masked. +pub(super) fn redacted_probe_screen(visible: &str) -> String { + let redacted = crate::core::SecretRedactor::redact(visible); + crate::core::PersonalInfoRedactor::redact_emails_in_text(Some(&redacted), true) + .unwrap_or(redacted) +} + +/// Only a parseable usage screen is worth sharing; errors are retried live. +fn claude_cli_output_is_shareable(output: &str) -> bool { + claude_cli_error_from_output(output).is_none() + && ClaudeProvider::new().parse_cli_output(output).is_ok() +} + +/// Cross-process guard around the Claude PTY probe. Claude Code refuses to +/// start a session whose `--session-id` is already running ("Session ID … +/// is already in use"), so two codexbar processes (for example the +/// `serve` daemon and a one-off `usage` call) must not probe concurrently. +pub(super) struct ClaudeProbeLock(std::fs::File); + +impl ClaudeProbeLock { + /// Wait for the probe lock. `Ok(None)` means locking is unsupported here + /// and the probe runs unlocked. A probe still running elsewhere after the + /// wait is an error: probing alongside it would reuse its session id. + fn acquire(probe_dir: &std::path::Path) -> Result, ProviderError> { + Self::acquire_within(probe_dir, CLAUDE_PROBE_LOCK_WAIT) + } + + pub(super) fn acquire_within( + probe_dir: &std::path::Path, + wait: Duration, + ) -> Result, ProviderError> { + let path = probe_dir.join(CLAUDE_PROBE_LOCK_FILE); + let file = match std::fs::OpenOptions::new() + .create(true) + .read(true) + .write(true) + .truncate(false) + .open(&path) + { + Ok(file) => file, + Err(err) => { + tracing::debug!(error = %err, "Claude probe lock file unavailable; continuing unlocked"); + return Ok(None); + } + }; + let deadline = Instant::now() + wait; + loop { + match file.try_lock() { + Ok(()) => return Ok(Some(Self(file))), + Err(std::fs::TryLockError::WouldBlock) => {} + Err(std::fs::TryLockError::Error(err)) => { + tracing::debug!(error = %err, "Claude probe lock unavailable; continuing unlocked"); + return Ok(None); + } + } + if Instant::now() >= deadline { + return Err(ProviderError::Other( + "Timed out waiting for another CodexBar process to finish its Claude CLI \ + usage probe." + .to_string(), + )); + } + std::thread::sleep(Duration::from_millis(200)); + } + } +} + +impl Drop for ClaudeProbeLock { + fn drop(&mut self) { + // Best-effort unlock; closing the handle releases it anyway. + let _unlocked = self.0.unlock(); + } +} + +fn claude_usage_probe_dir() -> Result { + let base = dirs::data_local_dir() + .or_else(dirs::home_dir) + .ok_or_else(|| { + ProviderError::Other("Could not resolve a local data directory".to_string()) + })?; + let dir = base.join("CodexBar").join("claude-usage-probe"); + std::fs::create_dir_all(&dir).map_err(|e| { + ProviderError::Other(format!( + "Failed to prepare Claude CLI probe directory: {}", + e + )) + })?; + Ok(dir) +} + +/// Persist and reuse one probe session id so repeated `/usage` PTY launches do +/// not register a fresh empty Claude account session each refresh (upstream #2263). +pub(super) fn load_or_create_probe_session_id(probe_dir: &std::path::Path) -> String { + let path = probe_dir.join(CLAUDE_PROBE_SESSION_ID_FILE); + if let Ok(raw) = std::fs::read_to_string(&path) { + let trimmed = raw.trim(); + if uuid::Uuid::parse_str(trimmed).is_ok() { + return trimmed.to_ascii_lowercase(); + } + } + let id = uuid::Uuid::new_v4().to_string().to_ascii_lowercase(); + if let Err(err) = std::fs::write(&path, &id) { + tracing::debug!(error = %err, "failed to persist Claude probe session id"); + } + id +} + +/// Claude treats `--session-id` as create-only when a local transcript JSONL +/// already exists for that id. Clear probe-dir jsonl leftovers before reuse. +pub(super) fn cleanup_probe_session_jsonl(probe_dir: &std::path::Path) { + let Ok(entries) = std::fs::read_dir(probe_dir) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().and_then(|e| e.to_str()) == Some("jsonl") { + // Best-effort cleanup: a locked or missing probe file just stays. + let _removed = std::fs::remove_file(&path); + } + } + cleanup_probe_transcript(probe_dir); +} + +/// Claude stores the transcript for a working directory under +/// `/projects//.jsonl`, where the +/// config dir is `CLAUDE_CONFIG_DIR` or `~/.claude`. Remove the probe session +/// transcripts there, otherwise the fixed `--session-id` fails with "already +/// in use" on the next run. +fn cleanup_probe_transcript(probe_dir: &std::path::Path) { + let Ok(config_dir) = accounts::config_dir() else { + return; + }; + cleanup_probe_transcripts_in(&config_dir.join("projects"), probe_dir); +} + +pub(super) fn cleanup_probe_transcripts_in( + projects_root: &std::path::Path, + probe_dir: &std::path::Path, +) { + let project_dir = projects_root.join(claude_project_dir_name(probe_dir)); + let Ok(entries) = std::fs::read_dir(&project_dir) else { + return; + }; + for entry in entries.flatten() { + let path = entry.path(); + let is_file = entry.file_type().is_ok_and(|kind| kind.is_file()); + if is_file && path.extension().and_then(|e| e.to_str()) == Some("jsonl") { + // Best-effort cleanup: a locked transcript just stays. + let _removed = std::fs::remove_file(&path); + } + } + // Succeeds only when nothing else is left in the probe's project dir. + let _removed = std::fs::remove_dir(&project_dir); +} + +/// Longest project directory name Claude Code writes before it truncates the +/// name and appends a hash of the full path. +const CLAUDE_PROJECT_DIR_NAME_MAX: usize = 200; + +/// Claude Code's project directory name for a working directory: every UTF-16 +/// code unit that is not an ASCII letter or digit becomes `-` +/// (`C:\Users\x` -> `C--Users-x`), and long names are cut to 200 characters +/// plus `-`. (Claude Code also NFC-normalizes the +/// path first; Windows paths are normally NFC already.) +pub(super) fn claude_project_dir_name(dir: &std::path::Path) -> String { + let path = dir.to_string_lossy(); + let sanitized: String = path + .encode_utf16() + .map(|unit| match u8::try_from(unit) { + Ok(byte) if byte.is_ascii_alphanumeric() => char::from(byte), + _ => '-', + }) + .collect(); + if sanitized.len() <= CLAUDE_PROJECT_DIR_NAME_MAX { + return sanitized; + } + format!( + "{}-{}", + &sanitized[..CLAUDE_PROJECT_DIR_NAME_MAX], + javascript_hash_base36(&path) + ) +} + +/// `Math.abs(hash).toString(36)` of the JavaScript string hash +/// `hash = (hash << 5) - hash + charCode`, kept in 32 bits. +pub(super) fn javascript_hash_base36(text: &str) -> String { + let hash = text.encode_utf16().fold(0i32, |hash, unit| { + hash.wrapping_mul(31).wrapping_add(i32::from(unit)) + }); + let mut magnitude = i64::from(hash).unsigned_abs(); + let mut digits = Vec::new(); + loop { + digits.push(char::from_digit((magnitude % 36) as u32, 36).unwrap_or('0')); + magnitude /= 36; + if magnitude == 0 { + break; + } + } + digits.iter().rev().collect() +} + +/// Arguments shared by every Claude CLI `/usage` probe. +/// +/// The remote-control startup hook can otherwise change the interactive +/// session before the usage command is collected. Keep this override in one +/// helper so future CLI probe paths cannot silently omit it. +pub(super) fn claude_usage_settings_args() -> [String; 2] { + [ + "--settings".to_string(), + // Issue #778 (Claude Code 2.1.293): `tui: default` stops the "Try the new + // fullscreen renderer?" offer. Older builds ignore unknown settings keys. + r#"{"remoteControlAtStartup":false,"tui":"default"}"#.to_string(), + ] +} + +pub(super) fn claude_probe_launch_args(session_id: &str) -> Vec { + let mut args = vec![ + "--setting-sources".to_string(), + "user".to_string(), + "--allowed-tools".to_string(), + String::new(), + ]; + args.extend(claude_usage_settings_args()); + args.extend(["--session-id".to_string(), session_id.to_string()]); + args +} + +struct ClaudePtyProbeOptions { + script: &'static str, + timeout_secs: f64, + idle_timeout_secs: Option, + initial_delay_secs: f64, + script_char_delay_secs: f64, + script_line_delay_secs: f64, + screen_responder: Option, + /// Re-type the script at these offsets while no done marker is visible. + script_retry_delays_secs: &'static [f64], + script_done_substrings: &'static [&'static str], + script_echo_substrings: &'static [&'static str], + /// Idle window after the done marker appeared (trailing output only). + idle_timeout_after_done_secs: Option, + /// Reuse a recent usage screen another process stored, and share this + /// one (the `/usage` probe only, never the trust preflight). + share_output: bool, +} + +/// Offsets (seconds after launch) at which `/usage` is re-sent when Claude's +/// input widget was not ready for the first attempt. Claude Code needs roughly +/// 1-5 s to mount its prompt on Windows, and keystrokes before that are lost. +const CLAUDE_USAGE_RETRY_DELAYS_SECS: &[f64] = &[6.0, 9.5, 14.0]; +/// Output markers that prove `/usage` opened (limits view or activity stats). +const CLAUDE_USAGE_DONE_MARKERS: &[&str] = &[ + "current session", + "current week", + "total duration", + "favorite model:", + "total tokens:", +]; +/// The typed command as Claude echoes it into its prompt line. While this is +/// visible the first attempt is still being processed, so do not type again. +const CLAUDE_USAGE_ECHO_MARKERS: &[&str] = &["❯ /usage", "> /usage", "/usage show session cost"]; + +async fn run_claude_usage_pty_probe( + claude_path: std::path::PathBuf, + working_directory: std::path::PathBuf, +) -> Result { + run_claude_pty_probe( + claude_path, + working_directory, + ClaudePtyProbeOptions { + script: "/usage", + timeout_secs: 24.0, + idle_timeout_secs: Some(6.0), + initial_delay_secs: 3.0, + script_char_delay_secs: 0.04, + script_line_delay_secs: 0.0, + screen_responder: None, + script_retry_delays_secs: CLAUDE_USAGE_RETRY_DELAYS_SECS, + script_done_substrings: CLAUDE_USAGE_DONE_MARKERS, + script_echo_substrings: CLAUDE_USAGE_ECHO_MARKERS, + idle_timeout_after_done_secs: Some(1.5), + share_output: true, + }, + ) + .await +} + +async fn run_claude_trust_preflight( + claude_path: std::path::PathBuf, + working_directory: std::path::PathBuf, +) -> Result { + run_claude_pty_probe( + claude_path, + working_directory, + ClaudePtyProbeOptions { + script: "", + timeout_secs: 15.0, + idle_timeout_secs: Some(4.0), + initial_delay_secs: 0.6, + script_char_delay_secs: 0.0, + script_line_delay_secs: 0.0, + screen_responder: Some(trust_dialog::TRUST_RESPONDER), + script_retry_delays_secs: &[], + script_done_substrings: &[], + script_echo_substrings: &[], + idle_timeout_after_done_secs: None, + share_output: false, + }, + ) + .await +} + +pub(super) fn resolve_claude_cli_path() -> Result { + locate_claude_binary().ok_or_else(|| { + ProviderError::NotInstalled( + "Claude CLI not found. Install from https://docs.claude.ai/claude-code".to_string(), + ) + }) +} + +pub(super) async fn fetch_claude_cli_usage_text( + claude_path: std::path::PathBuf, +) -> Result { + let probe_dir = claude_usage_probe_dir()?; + if let Some(login) = claude_login_fingerprint() + && let Some(cached) = load_cached_probe_output(&probe_dir, &login) + { + return Ok(cached); + } + let combined = run_claude_usage_pty_probe(claude_path.clone(), probe_dir.clone()).await?; + if !is_workspace_trust_prompt(&cli_screen::render(&combined, true).to_lowercase()) { + return Ok(combined); + } + + run_claude_trust_preflight(claude_path.clone(), probe_dir.clone()).await?; + run_claude_usage_pty_probe(claude_path, probe_dir).await +} + +/// Windows launch failures the CLI prints instead of a usage screen, as +/// (lowercase marker, user-facing message). +const CLAUDE_CLI_ENVIRONMENT_ERRORS: &[(&str, &str)] = &[ + ( + "requires git-bash", + "Claude CLI requires Git Bash on Windows. Install Git for Windows or set \ + CLAUDE_CODE_GIT_BASH_PATH to your bash.exe path.", + ), + ( + "running scripts is disabled", + "Claude CLI could not start because PowerShell script execution is disabled. \ + Use claude.cmd or adjust the execution policy.", + ), + ( + "cannot run a document in the middle of a pipeline", + "Claude CLI resolved to a Unix shell script on Windows. Reinstall Claude Code or \ + ensure claude.cmd is first on PATH.", + ), +]; + +/// Auth markers are checked before the environment markers. +pub(super) fn claude_cli_error_from_output(output: &str) -> Option { + let lowered = output.to_lowercase(); + if lowered.contains("not logged in") || lowered.contains("login required") { + return Some(ProviderError::AuthRequired); + } + if lowered.contains("token expired") || lowered.contains("token_expired") { + return Some(ProviderError::OAuthExpired( + "Token expired. Run `claude login` to refresh.".to_string(), + )); + } + if lowered.contains("authentication_error") { + return Some(ProviderError::OAuth( + "Authentication error. Run `claude login`.".to_string(), + )); + } + CLAUDE_CLI_ENVIRONMENT_ERRORS + .iter() + .find(|(marker, _)| lowered.contains(marker)) + .map(|(_, message)| ProviderError::Other((*message).to_string())) +} + +/// Environment overrides for passive Claude CLI PTY probes. +pub(super) fn claude_passive_probe_env( + mut base: std::collections::HashMap, +) -> std::collections::HashMap { + // Passive status/usage probes must not mutate or update the user's Claude CLI installation. + base.insert("NO_COLOR".to_string(), "1".to_string()); + base.insert("DISABLE_AUTOUPDATER".to_string(), "1".to_string()); + // Issue #778 (Claude Code 2.1.293): stops the "Claude in Chrome extension + // detected" offer, which would swallow `/usage`. Claude Code reads this + // variable; older builds ignore it. Probe launch only; config is untouched. + base.insert("CLAUDE_CODE_ENABLE_CFC".to_string(), "0".to_string()); + base +} + +async fn run_claude_pty_probe( + claude_path: std::path::PathBuf, + working_directory: std::path::PathBuf, + probe: ClaudePtyProbeOptions, +) -> Result { + tokio::task::spawn_blocking(move || { + // Keep ownership in the worker: cancelling the async refresh does not + // stop spawn_blocking or its CLI process from rotating credentials. + let _account_operation = accounts::CREDENTIAL_OPERATION.blocking_lock(); + let login: Option<&dyn Fn() -> Option> = + probe.share_output.then_some(&claude_login_fingerprint); + run_locked_probe(&working_directory, login, || { + cleanup_probe_session_jsonl(&working_directory); + let session_id = load_or_create_probe_session_id(&working_directory); + let env = claude_passive_probe_env(TtyCommandRunner::enriched_environment()); + + let mut options = TtyCommandOptions::new() + .with_timeout(probe.timeout_secs) + .with_initial_delay(probe.initial_delay_secs) + .with_script_char_delay(probe.script_char_delay_secs) + .with_script_line_delay(probe.script_line_delay_secs) + .with_working_directory(working_directory.clone()) + .with_extra_args(claude_probe_launch_args(&session_id)); + if let Some(idle) = probe.idle_timeout_secs { + options = options.with_idle_timeout(idle); + } + if let Some(idle) = probe.idle_timeout_after_done_secs { + options = options.with_idle_timeout_after_done(idle); + } + if let Some(responder) = probe.screen_responder { + options = options.with_screen_responder(responder); + } + if !probe.script_retry_delays_secs.is_empty() { + options = options.with_script_retries( + probe.script_retry_delays_secs.to_vec(), + probe + .script_done_substrings + .iter() + .map(|marker| (*marker).to_string()) + .collect(), + probe + .script_echo_substrings + .iter() + .map(|marker| (*marker).to_string()) + .collect(), + ); + } + options.env = env.into(); + + TtyCommandRunner::new() + .run(&claude_path.to_string_lossy(), probe.script, options) + .map(|result| result.text) + .map_err(|error| match error { + crate::cli::tty_runner::TtyCommandError::TimedOut => ProviderError::Timeout, + other => ProviderError::Other(format!("Claude CLI failed: {}", other)), + }) + }) + }) + .await + .map_err(|e| ProviderError::Other(format!("Claude CLI probe failed: {}", e)))? +} + +/// Run one probe under the cross-process probe lock. `login` is set when the +/// screen may be shared and identifies the Claude login it belongs to: a +/// fresh screen another process stored for that login while this one waited +/// is reused, and a parseable screen is stored for the others unless the +/// login changed while the probe ran. +pub(super) fn run_locked_probe( + probe_dir: &std::path::Path, + login: Option<&dyn Fn() -> Option>, + probe: impl FnOnce() -> Result, +) -> Result { + let _probe_lock = ClaudeProbeLock::acquire(probe_dir)?; + let before = login.and_then(|login| login()); + if let Some(before) = &before + && let Some(cached) = load_cached_probe_output(probe_dir, before) + { + return Ok(cached); + } + let output = probe()?; + if let (Some(before), Some(login)) = (&before, login) + && login().as_ref() == Some(before) + && claude_cli_output_is_shareable(&output) + { + store_cached_probe_output(probe_dir, before, &output); + } + Ok(output) +} diff --git a/rust/src/providers/claude/cli_reset.rs b/rust/src/providers/claude/cli_reset.rs index 590f1d82d9..89cbe058bc 100644 --- a/rust/src/providers/claude/cli_reset.rs +++ b/rust/src/providers/claude/cli_reset.rs @@ -11,18 +11,25 @@ fn regex(cell: &'static OnceLock, pattern: &'static str) -> &'static Rege cell.get_or_init(|| Regex::new(pattern).expect("valid Claude CLI regex")) } -pub(super) fn parse_percent_line(line: &str) -> Option { +/// Every "N% used/left/..." in order, as used percent clamped to 0..=100. +pub(super) fn percent_matches(text: &str) -> impl Iterator + '_ { static PERCENT: OnceLock = OnceLock::new(); - let captures = regex( + regex( &PERCENT, r"(?i)(\d{1,3}(?:\.\d+)?)\s*%\s*(used|spent|consumed|left|remaining|available)", ) - .captures(line)?; - let value: f64 = captures.get(1)?.as_str().parse().ok()?; - match captures.get(2)?.as_str().to_ascii_lowercase().as_str() { - "left" | "remaining" | "available" => Some((100.0 - value).max(0.0)), - _ => Some(value.min(100.0)), - } + .captures_iter(text) + .filter_map(|captures| { + let value: f64 = captures.get(1)?.as_str().parse().ok()?; + match captures.get(2)?.as_str().to_ascii_lowercase().as_str() { + "left" | "remaining" | "available" => Some((100.0 - value).max(0.0)), + _ => Some(value.min(100.0)), + } + }) +} + +pub(super) fn parse_percent_line(line: &str) -> Option { + percent_matches(line).next() } pub(super) fn normalized_for_label_search(text: &str) -> String { @@ -37,6 +44,25 @@ pub(super) fn starts_next_usage_section(line: &str, current_label: &str) -> bool normalized.starts_with("current") && !normalized.contains(current_label) } +/// Up to `max_lines` lines from the label line at `idx`, ending before the +/// next "Current ..." heading that is not `label_normalized`. +pub(super) fn label_section<'a>( + lines: &'a [&'a str], + idx: usize, + label_normalized: &'a str, + max_lines: usize, +) -> impl Iterator { + lines + .iter() + .skip(idx) + .take(max_lines) + .enumerate() + .take_while(move |(offset, line)| { + *offset == 0 || !starts_next_usage_section(line, label_normalized) + }) + .map(|(_, line)| *line) +} + pub(super) fn extract_cli_scoped_weekly_limits( text: &str, now: DateTime, @@ -69,10 +95,7 @@ pub(super) fn extract_cli_scoped_weekly_limits( let mut used_percent = None; let mut reset_description = None; let current_label = normalized_for_label_search(line); - for (offset, section_line) in lines.iter().skip(idx).take(14).enumerate() { - if offset > 0 && starts_next_usage_section(section_line, ¤t_label) { - break; - } + for section_line in label_section(&lines, idx, ¤t_label, 14) { used_percent = used_percent.or_else(|| parse_percent_line(section_line)); reset_description = reset_description.or_else(|| { reset_re @@ -142,42 +165,33 @@ pub(super) fn parse_claude_reset_date_in_system_zone( let (raw, timezone) = normalize_claude_reset_text(text, system_timezone)?; let components = parse_claude_reset_components(&raw)?; let now_local = now.with_timezone(&timezone); - let candidates = match (components.year, components.month, components.day) { - (Some(year), Some(month), Some(day)) => local_reset_occurrences( - timezone, - year, - month, - day, - components.hour, - components.minute, - ), + // Without a year, try nearby years; without a date, yesterday to tomorrow. + let dates: Vec<(i32, u32, u32)> = match (components.year, components.month, components.day) { + (Some(year), Some(month), Some(day)) => vec![(year, month, day)], (None, Some(month), Some(day)) => (now_local.year() - 8..=now_local.year() + 8) - .flat_map(|year| { - local_reset_occurrences( - timezone, - year, - month, - day, - components.hour, - components.minute, - ) - }) + .map(|year| (year, month, day)) .collect(), (None, None, None) => (-1..=1) - .flat_map(|offset| { + .map(|offset| { let date = now_local.date_naive() + Duration::days(offset); - local_reset_occurrences( - timezone, - date.year(), - date.month(), - date.day(), - components.hour, - components.minute, - ) + (date.year(), date.month(), date.day()) }) .collect(), _ => return None, }; + let candidates = dates + .into_iter() + .flat_map(|(year, month, day)| { + local_reset_occurrences( + timezone, + year, + month, + day, + components.hour, + components.minute, + ) + }) + .collect(); resolve_claude_reset_occurrence(candidates, now, expected_window_minutes) } @@ -290,21 +304,14 @@ fn parse_claude_hour( } fn claude_month(month: &str) -> Option { - match month.to_ascii_lowercase().as_str() { - "jan" => Some(1), - "feb" => Some(2), - "mar" => Some(3), - "apr" => Some(4), - "may" => Some(5), - "jun" => Some(6), - "jul" => Some(7), - "aug" => Some(8), - "sep" => Some(9), - "oct" => Some(10), - "nov" => Some(11), - "dec" => Some(12), - _ => None, - } + const MONTHS: [&str; 12] = [ + "jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec", + ]; + let month = month.to_ascii_lowercase(); + (1..) + .zip(MONTHS) + .find(|(_, name)| *name == month) + .map(|(number, _)| number) } fn local_reset_occurrences( diff --git a/rust/src/providers/claude/cli_screen/tests.rs b/rust/src/providers/claude/cli_screen/tests.rs index 1d3fbd0db1..af0bd2f008 100644 --- a/rust/src/providers/claude/cli_screen/tests.rs +++ b/rust/src/providers/claude/cli_screen/tests.rs @@ -186,7 +186,7 @@ fn a_cursor_sequence_after_plain_text_switches_to_replay() { #[test] fn differential_redraw_fixture_keeps_the_scoped_weekly_quota() { // Plain escape stripping fuses the redraw fragments, which is the bug. - let stripped = crate::providers::claude::strip_ansi(USAGE_FIXTURE); + let stripped = strip_ansi(USAGE_FIXTURE); assert!(stripped.contains("51%usd"), "{stripped}"); let provider = ClaudeProvider::new(); @@ -227,3 +227,59 @@ fn redrawn_identity_uses_the_final_frame_with_cursor_positioned_spaces() { assert_eq!(usage.account_email.as_deref(), Some("fixture@example.com")); assert_eq!(usage.login_method.as_deref(), Some("Claude Max Account")); } + +#[test] +fn rendered_text_never_carries_an_escape() { + let inputs = [ + USAGE_FIXTURE, + STATUS_FIXTURE, + "plain \u{1b}[1mLogin method: Claude Max\u{1b}[22m", + "osc \u{1b}]0;title\u{7}tail \u{1b}]8;;x\u{1b}\\link", + "lone \u{1b}\n\u{1b}Xtrail \u{1b}", + "abc\u{1b}[2DZ \u{1b}(B done", + ]; + for input in inputs { + for preserve in [true, false] { + let rendered = render(input, preserve); + assert!(!rendered.contains('\u{1b}'), "{rendered:?}"); + assert_eq!(strip_ansi(&rendered), rendered); + } + } +} + +/// Plain escape stripping, the pre-replay behavior these tests compare against. +fn strip_ansi(text: &str) -> String { + let mut result = String::with_capacity(text.len()); + let mut chars = text.chars().peekable(); + + while let Some(c) = chars.next() { + if c == '\x1B' { + // Skip CSI sequences: ESC[...letter + if chars.peek() == Some(&'[') { + chars.next(); + let mut final_char = None; + while let Some(&next) = chars.peek() { + chars.next(); + if next.is_ascii_alphabetic() { + final_char = Some(next); + break; + } + } + if final_char == Some('C') { + result.push(' '); + } + // Skip OSC sequences: ESC]...BEL + } else if chars.peek() == Some(&']') { + for next in chars.by_ref() { + if next == '\x07' || next == '\\' { + break; + } + } + } + } else { + result.push(c); + } + } + + result +} diff --git a/rust/src/providers/claude/cli_text.rs b/rust/src/providers/claude/cli_text.rs new file mode 100644 index 0000000000..c2883079c5 --- /dev/null +++ b/rust/src/providers/claude/cli_text.rs @@ -0,0 +1,135 @@ +use super::cli_reset::{label_section, normalized_for_label_search, parse_percent_line}; +use regex_lite::Regex; + +pub(super) fn is_non_interactive_slash_command_response(text: &str) -> bool { + let mentions_usage_and_exit = text.contains("/usage") && text.contains("/exit"); + let says_entered_commands = + text.contains("i see you've entered") || text.contains("you've entered two slash commands"); + let says_no_slash_command = text.contains("available custom slash commands") + && text.contains("don't see these commands"); + let says_usage_is_cli_only = text + .contains("token usage and statistics are typically displayed by the cli interface") + || text.contains("i don't have direct access to those metrics"); + + mentions_usage_and_exit + && (says_entered_commands || says_no_slash_command || says_usage_is_cli_only) +} + +pub(super) fn is_workspace_trust_prompt(text: &str) -> bool { + text.contains("quick safety check") + && text.contains("trust this folder") + && text.contains("yes, i trust this folder") +} + +/// Current `/usage` panels show a local session summary above the plan limits, +/// so the presence of a limit section outranks the activity-stats markers. +pub(super) fn has_plan_limit_section(text: &str) -> bool { + text.contains("current session") || text.contains("current week") +} + +pub(super) fn is_cli_activity_stats_response(text: &str) -> bool { + let has_activity_overview = text.contains("favorite model:") || text.contains("total tokens:"); + let has_session_cost_summary = + text.contains("total duration") && text.contains("usage:") && text.contains("cache read"); + + has_activity_overview || has_session_cost_summary +} + +/// The weekly heading Claude prints, newest wording first. +pub(super) const WEEKLY_LABELS: [&str; 2] = ["current week (all models)", "current week"]; + +/// First match of `find` in any section headed by `label`; each section is +/// scanned for at most `max_lines` lines, label line included. +fn find_near_label( + text: &str, + label: &str, + max_lines: usize, + mut find: impl FnMut(&str) -> Option, +) -> Option { + let label_normalized = normalized_for_label_search(label); + let lines: Vec<&str> = text.lines().collect(); + lines + .iter() + .enumerate() + .filter(|(_, line)| normalized_for_label_search(line).contains(&label_normalized)) + .find_map(|(idx, _)| { + label_section(&lines, idx, &label_normalized, max_lines).find_map(&mut find) + }) +} + +/// Percentage near a label (e.g. "Current session"), as "used". +pub(super) fn extract_percent_near_label(text: &str, label: &str) -> Option { + find_near_label(text, label, 12, parse_percent_line) +} + +pub(super) fn is_exhausted_short_form(clean_lower: &str) -> bool { + clean_lower.contains("out of extra usage") || clean_lower.contains("hit your limit") +} + +/// Extract email address from text +pub(super) fn extract_email(text: &str) -> Option { + // Try explicit patterns first + let patterns = [ + r"Account:\s*([^\s@]+@[^\s@]+\.[^\s]+)", + r"Email:\s*([^\s@]+@[^\s@]+\.[^\s]+)", + r"([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,})", + ]; + + for pattern in patterns { + if let Ok(re) = Regex::new(pattern) + && let Some(caps) = re.captures(text) + && let Some(m) = caps.get(1) + { + return Some(m.as_str().trim().to_string()); + } + } + + None +} + +/// Extract login method / plan name from text +pub(super) fn extract_login_method(text: &str) -> Option { + // Look for explicit "Login method:" line + if let Ok(re) = Regex::new(r"(?i)login\s+method:\s*(.+)") + && let Some(caps) = re.captures(text) + && let Some(m) = caps.get(1) + { + let method = m.as_str().trim(); + if !method.is_empty() { + return Some(clean_plan_name(method)); + } + } + + // Look for "Claude " patterns + if let Ok(re) = Regex::new(r"(?i)(claude\s+(?:max|pro|ultra|team|free)[a-z0-9\s._-]*)") + && let Some(caps) = re.captures(text) + && let Some(m) = caps.get(1) + { + let plan = m.as_str().trim(); + if !plan.to_lowercase().contains("code") { + return Some(clean_plan_name(plan)); + } + } + + None +} + +/// Reset text near a label, from "resets" to the end of its line. +pub(super) fn extract_reset_description(text: &str, label: &str) -> Option { + find_near_label(text, label, 14, extract_inline_reset_description) +} + +/// Extract a "resets ..." suffix from a short single-line status. +pub(super) fn extract_inline_reset_description(text: &str) -> Option { + let lower = text.to_lowercase(); + let pos = lower.find("resets")?; + Some(text[pos..].trim().to_string()) +} + +/// Clean up a plan name from rendered (escape-free) text: drop bracketed +/// codes like `[22m` and trim. +fn clean_plan_name(text: &str) -> String { + let re = Regex::new(r"\[\d+m").unwrap_or_else(|_| Regex::new(".^").unwrap()); + let result = re.replace_all(text, ""); + result.trim().to_string() +} diff --git a/rust/src/providers/claude/mod.rs b/rust/src/providers/claude/mod.rs index b5fa7abe3d..734928868d 100755 --- a/rust/src/providers/claude/mod.rs +++ b/rust/src/providers/claude/mod.rs @@ -4,8 +4,11 @@ pub mod accounts; mod admin_api; mod auto_precision; pub mod claude_swap; +mod cli_binary; +mod cli_probe; mod cli_reset; mod cli_screen; +mod cli_text; mod oauth; pub mod quota_history; mod reset_credits; @@ -16,27 +19,30 @@ mod web_api; use async_trait::async_trait; use chrono::Utc; -use regex_lite::Regex; -#[cfg(windows)] -use std::os::windows::process::CommandExt; -#[cfg(windows)] -use std::process::{Command as StdCommand, Stdio}; use std::sync::LazyLock; use std::sync::Mutex; use std::time::{Duration, Instant}; -use crate::cli::tty_runner::{TtyCommandOptions, TtyCommandRunner}; use crate::core::{ FetchContext, LastGoodFailurePolicy, Provider, ProviderError, ProviderFetchResult, ProviderId, RateWindow, SourceMode, UsageSnapshot, }; use admin_api::ClaudeAdminApiFetcher; +use cli_binary::detect_claude_version; +pub use cli_binary::locate_claude_binary; +use cli_probe::{ + claude_cli_error_from_output, fetch_claude_cli_usage_text, redacted_probe_screen, + resolve_claude_cli_path, +}; #[cfg(test)] use cli_reset::parse_claude_reset_date_in_system_zone; -use cli_reset::{ - extract_cli_scoped_weekly_limits, normalized_for_label_search, parse_claude_reset_date, - parse_percent_line, starts_next_usage_section, +use cli_reset::{extract_cli_scoped_weekly_limits, parse_claude_reset_date, percent_matches}; +use cli_text::{ + WEEKLY_LABELS, extract_email, extract_inline_reset_description, extract_login_method, + extract_percent_near_label, extract_reset_description, has_plan_limit_section, + is_cli_activity_stats_response, is_exhausted_short_form, + is_non_interactive_slash_command_response, }; // ── Upstream 0.50.1 #2516: CLI usage-result cache ──────────────────────────── @@ -179,26 +185,6 @@ fn claude_plan_label(tier: &str) -> String { } } -const CLAUDE_PROBE_SESSION_ID_FILE: &str = ".codexbar-session-id"; -const CLAUDE_PROBE_LOCK_FILE: &str = ".codexbar-probe.lock"; -const CLAUDE_PROBE_CACHE_FILE: &str = ".codexbar-usage-cache.json"; -/// How long a second codexbar process waits for a running probe to finish. -const CLAUDE_PROBE_LOCK_WAIT: Duration = Duration::from_secs(30); -/// Every codexbar process (the `serve` daemon, one-off `usage` calls from the -/// companion) launches its own Claude CLI for a probe. The interactive -/// `/usage` screen costs 6-10 s of CPU each time, so a recent successful -/// probe output is shared across processes for this long. -const CLAUDE_PROBE_CACHE_TTL: Duration = Duration::from_secs(45); - -#[derive(serde::Serialize, serde::Deserialize)] -struct ClaudeProbeCache { - captured_at_unix: u64, - /// `claude_login_fingerprint` of the login the screen belongs to. - #[serde(default)] - login: String, - output: String, -} - fn unix_now_secs() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -206,569 +192,6 @@ fn unix_now_secs() -> u64 { .unwrap_or(0) } -fn load_cached_probe_output(probe_dir: &std::path::Path, login: &str) -> Option { - let raw = std::fs::read_to_string(probe_dir.join(CLAUDE_PROBE_CACHE_FILE)).ok()?; - let cache: ClaudeProbeCache = serde_json::from_str(&raw).ok()?; - let age = unix_now_secs().saturating_sub(cache.captured_at_unix); - if auto_precision::probe_screen_is_superseded(cache.captured_at_unix) - || login.is_empty() - || cache.login != login - || age > CLAUDE_PROBE_CACHE_TTL.as_secs() - || cache.output.trim().is_empty() - { - return None; - } - tracing::debug!(age_secs = age, "Reusing recent Claude CLI probe output"); - Some(cache.output) -} - -fn store_cached_probe_output(probe_dir: &std::path::Path, login: &str, output: &str) { - let cache = ClaudeProbeCache { - captured_at_unix: unix_now_secs(), - login: login.to_string(), - output: output.to_string(), - }; - // Atomic, because other processes read the cache without the probe lock. - let stored = serde_json::to_vec(&cache) - .map_err(anyhow::Error::from) - .and_then(|json| { - crate::atomic_file::write_atomic(&probe_dir.join(CLAUDE_PROBE_CACHE_FILE), &json) - }); - if let Err(err) = stored { - tracing::debug!(error = %err, "failed to persist Claude probe cache"); - } -} - -/// Identifies the Claude login a probe runs under without reading any -/// credential: the location, size and modification time of Claude Code's -/// `.credentials.json`, which every login, token refresh and account switch -/// rewrites. Without a credentials file a probe screen is never shared. -fn claude_login_fingerprint() -> Option { - let credentials = accounts::config_dir().ok()?.join(".credentials.json"); - login_fingerprint_at(&credentials) -} - -fn login_fingerprint_at(credentials: &std::path::Path) -> Option { - let metadata = std::fs::metadata(credentials).ok()?; - let modified = metadata - .modified() - .ok()? - .duration_since(std::time::UNIX_EPOCH) - .ok()?; - let identity = format!( - "{}|{}|{}", - credentials.display(), - metadata.len(), - modified.as_nanos() - ); - Some(crate::core::sha256_hex(identity.as_bytes())) -} - -/// The probe screen as it may appear in a log: secrets and email addresses -/// (the screen can show the signed-in account) are masked. -fn redacted_probe_screen(visible: &str) -> String { - let redacted = crate::core::SecretRedactor::redact(visible); - crate::core::PersonalInfoRedactor::redact_emails_in_text(Some(&redacted), true) - .unwrap_or(redacted) -} - -/// Only a parseable usage screen is worth sharing; errors are retried live. -fn claude_cli_output_is_shareable(output: &str) -> bool { - claude_cli_error_from_output(output).is_none() - && ClaudeProvider::new().parse_cli_output(output).is_ok() -} - -/// Cross-process guard around the Claude PTY probe. Claude Code refuses to -/// start a session whose `--session-id` is already running ("Session ID … -/// is already in use"), so two codexbar processes (for example the -/// `serve` daemon and a one-off `usage` call) must not probe concurrently. -struct ClaudeProbeLock(std::fs::File); - -impl ClaudeProbeLock { - /// Wait for the probe lock. `Ok(None)` means locking is unsupported here - /// and the probe runs unlocked. A probe still running elsewhere after the - /// wait is an error: probing alongside it would reuse its session id. - fn acquire(probe_dir: &std::path::Path) -> Result, ProviderError> { - Self::acquire_within(probe_dir, CLAUDE_PROBE_LOCK_WAIT) - } - - fn acquire_within( - probe_dir: &std::path::Path, - wait: Duration, - ) -> Result, ProviderError> { - let path = probe_dir.join(CLAUDE_PROBE_LOCK_FILE); - let file = match std::fs::OpenOptions::new() - .create(true) - .read(true) - .write(true) - .truncate(false) - .open(&path) - { - Ok(file) => file, - Err(err) => { - tracing::debug!(error = %err, "Claude probe lock file unavailable; continuing unlocked"); - return Ok(None); - } - }; - let deadline = Instant::now() + wait; - loop { - match file.try_lock() { - Ok(()) => return Ok(Some(Self(file))), - Err(std::fs::TryLockError::WouldBlock) => {} - Err(std::fs::TryLockError::Error(err)) => { - tracing::debug!(error = %err, "Claude probe lock unavailable; continuing unlocked"); - return Ok(None); - } - } - if Instant::now() >= deadline { - return Err(ProviderError::Other( - "Timed out waiting for another CodexBar process to finish its Claude CLI \ - usage probe." - .to_string(), - )); - } - std::thread::sleep(Duration::from_millis(200)); - } - } -} - -impl Drop for ClaudeProbeLock { - fn drop(&mut self) { - // Best-effort unlock; closing the handle releases it anyway. - let _unlocked = self.0.unlock(); - } -} - -fn claude_usage_probe_dir() -> Result { - let base = dirs::data_local_dir() - .or_else(dirs::home_dir) - .ok_or_else(|| { - ProviderError::Other("Could not resolve a local data directory".to_string()) - })?; - let dir = base.join("CodexBar").join("claude-usage-probe"); - std::fs::create_dir_all(&dir).map_err(|e| { - ProviderError::Other(format!( - "Failed to prepare Claude CLI probe directory: {}", - e - )) - })?; - Ok(dir) -} - -/// Persist and reuse one probe session id so repeated `/usage` PTY launches do -/// not register a fresh empty Claude account session each refresh (upstream #2263). -fn load_or_create_probe_session_id(probe_dir: &std::path::Path) -> String { - let path = probe_dir.join(CLAUDE_PROBE_SESSION_ID_FILE); - if let Ok(raw) = std::fs::read_to_string(&path) { - let trimmed = raw.trim(); - if uuid::Uuid::parse_str(trimmed).is_ok() { - return trimmed.to_ascii_lowercase(); - } - } - let id = uuid::Uuid::new_v4().to_string().to_ascii_lowercase(); - if let Err(err) = std::fs::write(&path, &id) { - tracing::debug!(error = %err, "failed to persist Claude probe session id"); - } - id -} - -/// Claude treats `--session-id` as create-only when a local transcript JSONL -/// already exists for that id. Clear probe-dir jsonl leftovers before reuse. -fn cleanup_probe_session_jsonl(probe_dir: &std::path::Path) { - let Ok(entries) = std::fs::read_dir(probe_dir) else { - return; - }; - for entry in entries.flatten() { - let path = entry.path(); - if path.extension().and_then(|e| e.to_str()) == Some("jsonl") { - // Best-effort cleanup: a locked or missing probe file just stays. - let _removed = std::fs::remove_file(&path); - } - } - cleanup_probe_transcript(probe_dir); -} - -/// Claude stores the transcript for a working directory under -/// `/projects//.jsonl`, where the -/// config dir is `CLAUDE_CONFIG_DIR` or `~/.claude`. Remove the probe session -/// transcripts there, otherwise the fixed `--session-id` fails with "already -/// in use" on the next run. -fn cleanup_probe_transcript(probe_dir: &std::path::Path) { - let Ok(config_dir) = accounts::config_dir() else { - return; - }; - cleanup_probe_transcripts_in(&config_dir.join("projects"), probe_dir); -} - -fn cleanup_probe_transcripts_in(projects_root: &std::path::Path, probe_dir: &std::path::Path) { - let project_dir = projects_root.join(claude_project_dir_name(probe_dir)); - let Ok(entries) = std::fs::read_dir(&project_dir) else { - return; - }; - for entry in entries.flatten() { - let path = entry.path(); - let is_file = entry.file_type().is_ok_and(|kind| kind.is_file()); - if is_file && path.extension().and_then(|e| e.to_str()) == Some("jsonl") { - // Best-effort cleanup: a locked transcript just stays. - let _removed = std::fs::remove_file(&path); - } - } - // Succeeds only when nothing else is left in the probe's project dir. - let _removed = std::fs::remove_dir(&project_dir); -} - -/// Longest project directory name Claude Code writes before it truncates the -/// name and appends a hash of the full path. -const CLAUDE_PROJECT_DIR_NAME_MAX: usize = 200; - -/// Claude Code's project directory name for a working directory: every UTF-16 -/// code unit that is not an ASCII letter or digit becomes `-` -/// (`C:\Users\x` -> `C--Users-x`), and long names are cut to 200 characters -/// plus `-`. (Claude Code also NFC-normalizes the -/// path first; Windows paths are normally NFC already.) -fn claude_project_dir_name(dir: &std::path::Path) -> String { - let path = dir.to_string_lossy(); - let sanitized: String = path - .encode_utf16() - .map(|unit| match u8::try_from(unit) { - Ok(byte) if byte.is_ascii_alphanumeric() => char::from(byte), - _ => '-', - }) - .collect(); - if sanitized.len() <= CLAUDE_PROJECT_DIR_NAME_MAX { - return sanitized; - } - format!( - "{}-{}", - &sanitized[..CLAUDE_PROJECT_DIR_NAME_MAX], - javascript_hash_base36(&path) - ) -} - -/// `Math.abs(hash).toString(36)` of the JavaScript string hash -/// `hash = (hash << 5) - hash + charCode`, kept in 32 bits. -fn javascript_hash_base36(text: &str) -> String { - let hash = text.encode_utf16().fold(0i32, |hash, unit| { - hash.wrapping_mul(31).wrapping_add(i32::from(unit)) - }); - let mut magnitude = i64::from(hash).unsigned_abs(); - let mut digits = Vec::new(); - loop { - digits.push(char::from_digit((magnitude % 36) as u32, 36).unwrap_or('0')); - magnitude /= 36; - if magnitude == 0 { - break; - } - } - digits.iter().rev().collect() -} - -/// Arguments shared by every Claude CLI `/usage` probe. -/// -/// The remote-control startup hook can otherwise change the interactive -/// session before the usage command is collected. Keep this override in one -/// helper so future CLI probe paths cannot silently omit it. -fn claude_usage_settings_args() -> [String; 2] { - [ - "--settings".to_string(), - // Issue #778 (Claude Code 2.1.293): `tui: default` stops the "Try the new - // fullscreen renderer?" offer. Older builds ignore unknown settings keys. - r#"{"remoteControlAtStartup":false,"tui":"default"}"#.to_string(), - ] -} - -fn claude_probe_launch_args(session_id: &str) -> Vec { - let mut args = vec![ - "--setting-sources".to_string(), - "user".to_string(), - "--allowed-tools".to_string(), - String::new(), - ]; - args.extend(claude_usage_settings_args()); - args.extend(["--session-id".to_string(), session_id.to_string()]); - args -} - -struct ClaudePtyProbeOptions { - script: &'static str, - timeout_secs: f64, - idle_timeout_secs: Option, - initial_delay_secs: f64, - script_char_delay_secs: f64, - script_line_delay_secs: f64, - screen_responder: Option, - /// Re-type the script at these offsets while no done marker is visible. - script_retry_delays_secs: &'static [f64], - script_done_substrings: &'static [&'static str], - script_echo_substrings: &'static [&'static str], - /// Idle window after the done marker appeared (trailing output only). - idle_timeout_after_done_secs: Option, - /// Reuse a recent usage screen another process stored, and share this - /// one (the `/usage` probe only, never the trust preflight). - share_output: bool, -} - -/// Offsets (seconds after launch) at which `/usage` is re-sent when Claude's -/// input widget was not ready for the first attempt. Claude Code needs roughly -/// 1-5 s to mount its prompt on Windows, and keystrokes before that are lost. -const CLAUDE_USAGE_RETRY_DELAYS_SECS: &[f64] = &[6.0, 9.5, 14.0]; -/// Output markers that prove `/usage` opened (limits view or activity stats). -const CLAUDE_USAGE_DONE_MARKERS: &[&str] = &[ - "current session", - "current week", - "total duration", - "favorite model:", - "total tokens:", -]; -/// The typed command as Claude echoes it into its prompt line. While this is -/// visible the first attempt is still being processed, so do not type again. -const CLAUDE_USAGE_ECHO_MARKERS: &[&str] = &["❯ /usage", "> /usage", "/usage show session cost"]; - -async fn run_claude_usage_pty_probe( - claude_path: std::path::PathBuf, - working_directory: std::path::PathBuf, -) -> Result { - run_claude_pty_probe( - claude_path, - working_directory, - ClaudePtyProbeOptions { - script: "/usage", - timeout_secs: 24.0, - idle_timeout_secs: Some(6.0), - initial_delay_secs: 3.0, - script_char_delay_secs: 0.04, - script_line_delay_secs: 0.0, - screen_responder: None, - script_retry_delays_secs: CLAUDE_USAGE_RETRY_DELAYS_SECS, - script_done_substrings: CLAUDE_USAGE_DONE_MARKERS, - script_echo_substrings: CLAUDE_USAGE_ECHO_MARKERS, - idle_timeout_after_done_secs: Some(1.5), - share_output: true, - }, - ) - .await -} - -async fn run_claude_trust_preflight( - claude_path: std::path::PathBuf, - working_directory: std::path::PathBuf, -) -> Result { - run_claude_pty_probe( - claude_path, - working_directory, - ClaudePtyProbeOptions { - script: "", - timeout_secs: 15.0, - idle_timeout_secs: Some(4.0), - initial_delay_secs: 0.6, - script_char_delay_secs: 0.0, - script_line_delay_secs: 0.0, - screen_responder: Some(trust_dialog::TRUST_RESPONDER), - script_retry_delays_secs: &[], - script_done_substrings: &[], - script_echo_substrings: &[], - idle_timeout_after_done_secs: None, - share_output: false, - }, - ) - .await -} - -fn resolve_claude_cli_path() -> Result { - locate_claude_binary().ok_or_else(|| { - ProviderError::NotInstalled( - "Claude CLI not found. Install from https://docs.claude.ai/claude-code".to_string(), - ) - }) -} - -async fn fetch_claude_cli_usage_text( - claude_path: std::path::PathBuf, -) -> Result { - let probe_dir = claude_usage_probe_dir()?; - if let Some(login) = claude_login_fingerprint() - && let Some(cached) = load_cached_probe_output(&probe_dir, &login) - { - return Ok(cached); - } - let combined = run_claude_usage_pty_probe(claude_path.clone(), probe_dir.clone()).await?; - - rerun_claude_usage_after_trust_prompt(claude_path, probe_dir, combined).await -} - -async fn rerun_claude_usage_after_trust_prompt( - claude_path: std::path::PathBuf, - probe_dir: std::path::PathBuf, - combined: String, -) -> Result { - if !is_workspace_trust_prompt(&cli_screen::render(&combined, true).to_lowercase()) { - return Ok(combined); - } - - run_claude_trust_preflight(claude_path.clone(), probe_dir.clone()).await?; - run_claude_usage_pty_probe(claude_path, probe_dir).await -} - -fn claude_cli_error_from_output(output: &str) -> Option { - let lowered = output.to_lowercase(); - claude_cli_auth_error(&lowered).or_else(|| claude_cli_environment_error(&lowered)) -} - -fn claude_cli_auth_error(lowered: &str) -> Option { - if claude_output_requires_login(lowered) { - return Some(ProviderError::AuthRequired); - } - if lowered.contains("token expired") || lowered.contains("token_expired") { - return Some(ProviderError::OAuthExpired( - "Token expired. Run `claude login` to refresh.".to_string(), - )); - } - if lowered.contains("authentication_error") { - return Some(ProviderError::OAuth( - "Authentication error. Run `claude login`.".to_string(), - )); - } - - None -} - -fn claude_output_requires_login(lowered: &str) -> bool { - lowered.contains("not logged in") || lowered.contains("login required") -} - -fn claude_cli_environment_error(lowered: &str) -> Option { - if lowered.contains("requires git-bash") { - return Some(ProviderError::Other( - "Claude CLI requires Git Bash on Windows. Install Git for Windows or set \ - CLAUDE_CODE_GIT_BASH_PATH to your bash.exe path." - .to_string(), - )); - } - if lowered.contains("running scripts is disabled") { - return Some(ProviderError::Other( - "Claude CLI could not start because PowerShell script execution is disabled. \ - Use claude.cmd or adjust the execution policy." - .to_string(), - )); - } - if lowered.contains("cannot run a document in the middle of a pipeline") { - return Some(ProviderError::Other( - "Claude CLI resolved to a Unix shell script on Windows. Reinstall Claude Code or \ - ensure claude.cmd is first on PATH." - .to_string(), - )); - } - - None -} - -/// Environment overrides for passive Claude CLI PTY probes. -fn claude_passive_probe_env( - mut base: std::collections::HashMap, -) -> std::collections::HashMap { - // Passive status/usage probes must not mutate or update the user's Claude CLI installation. - base.insert("NO_COLOR".to_string(), "1".to_string()); - base.insert("DISABLE_AUTOUPDATER".to_string(), "1".to_string()); - // Issue #778 (Claude Code 2.1.293): stops the "Claude in Chrome extension - // detected" offer, which would swallow `/usage`. Claude Code reads this - // variable; older builds ignore it. Probe launch only; config is untouched. - base.insert("CLAUDE_CODE_ENABLE_CFC".to_string(), "0".to_string()); - base -} - -async fn run_claude_pty_probe( - claude_path: std::path::PathBuf, - working_directory: std::path::PathBuf, - probe: ClaudePtyProbeOptions, -) -> Result { - tokio::task::spawn_blocking(move || { - // Keep ownership in the worker: cancelling the async refresh does not - // stop spawn_blocking or its CLI process from rotating credentials. - let _account_operation = accounts::CREDENTIAL_OPERATION.blocking_lock(); - let login: Option<&dyn Fn() -> Option> = - probe.share_output.then_some(&claude_login_fingerprint); - run_locked_probe(&working_directory, login, || { - cleanup_probe_session_jsonl(&working_directory); - let session_id = load_or_create_probe_session_id(&working_directory); - let env = claude_passive_probe_env(TtyCommandRunner::enriched_environment()); - - let mut options = TtyCommandOptions::new() - .with_timeout(probe.timeout_secs) - .with_initial_delay(probe.initial_delay_secs) - .with_script_char_delay(probe.script_char_delay_secs) - .with_script_line_delay(probe.script_line_delay_secs) - .with_working_directory(working_directory.clone()) - .with_extra_args(claude_probe_launch_args(&session_id)); - if let Some(idle) = probe.idle_timeout_secs { - options = options.with_idle_timeout(idle); - } - if let Some(idle) = probe.idle_timeout_after_done_secs { - options = options.with_idle_timeout_after_done(idle); - } - if let Some(responder) = probe.screen_responder { - options = options.with_screen_responder(responder); - } - if !probe.script_retry_delays_secs.is_empty() { - options = options.with_script_retries( - probe.script_retry_delays_secs.to_vec(), - probe - .script_done_substrings - .iter() - .map(|marker| (*marker).to_string()) - .collect(), - probe - .script_echo_substrings - .iter() - .map(|marker| (*marker).to_string()) - .collect(), - ); - } - options.env = env.into(); - - TtyCommandRunner::new() - .run(&claude_path.to_string_lossy(), probe.script, options) - .map(|result| result.text) - .map_err(claude_tty_error) - }) - }) - .await - .map_err(|e| ProviderError::Other(format!("Claude CLI probe failed: {}", e)))? -} - -/// Run one probe under the cross-process probe lock. `login` is set when the -/// screen may be shared and identifies the Claude login it belongs to: a -/// fresh screen another process stored for that login while this one waited -/// is reused, and a parseable screen is stored for the others unless the -/// login changed while the probe ran. -fn run_locked_probe( - probe_dir: &std::path::Path, - login: Option<&dyn Fn() -> Option>, - probe: impl FnOnce() -> Result, -) -> Result { - let _probe_lock = ClaudeProbeLock::acquire(probe_dir)?; - let before = login.and_then(|login| login()); - if let Some(before) = &before - && let Some(cached) = load_cached_probe_output(probe_dir, before) - { - return Ok(cached); - } - let output = probe()?; - if let (Some(before), Some(login)) = (&before, login) - && login().as_ref() == Some(before) - && claude_cli_output_is_shareable(&output) - { - store_cached_probe_output(probe_dir, before, &output); - } - Ok(output) -} - -fn claude_tty_error(error: crate::cli::tty_runner::TtyCommandError) -> ProviderError { - match error { - crate::cli::tty_runner::TtyCommandError::TimedOut => ProviderError::Timeout, - other => ProviderError::Other(format!("Claude CLI failed: {}", other)), - } -} - fn last_good_failure_policy_for_error(error: &str) -> LastGoodFailurePolicy { let lower = error.to_ascii_lowercase(); if lower.contains("credentials not found") @@ -888,8 +311,12 @@ impl ClaudeProvider { ) -> Result { let mut failures = Vec::new(); - if let Some(result) = self.try_auto_admin_api(ctx, &mut failures).await? { - return Ok(result); + if self.admin_fetcher.has_credentials(ctx) { + tracing::debug!("Attempting Admin API fetch for Claude"); + let admin = self.admin_fetcher.fetch(ctx).await; + if let Some(result) = record_auto_source(&mut failures, "Admin API", admin)? { + return Ok(result); + } } if let Some(result) = @@ -944,17 +371,6 @@ impl ClaudeProvider { Err(claude_auto_fetch_error(failures)) } - async fn try_auto_admin_api( - &self, - ctx: &FetchContext, - failures: &mut Vec<(&'static str, ProviderError)>, - ) -> Result, ProviderError> { - if !self.admin_fetcher.has_credentials(ctx) { - return Ok(None); - } - record_auto_source(failures, "Admin API", self.fetch_via_admin_api(ctx).await) - } - async fn fetch_via_oauth( &self, ctx: &FetchContext, @@ -970,14 +386,6 @@ impl ClaudeProvider { self.oauth_fetcher.fetch().await } - async fn fetch_via_admin_api( - &self, - ctx: &FetchContext, - ) -> Result { - tracing::debug!("Attempting Admin API fetch for Claude"); - self.admin_fetcher.fetch(ctx).await - } - async fn fetch_via_web( &self, ctx: &FetchContext, @@ -1049,26 +457,15 @@ impl ClaudeProvider { )); } - // Parse session percent: "X% used" or "X% left" - let mut session_percent: Option = None; - let mut weekly_percent: Option = None; - - // Look for "Current session" section - if let Some(session_pct) = extract_percent_near_label(&clean, "current session") { - session_percent = Some(session_pct); - } - - // Look for "Current week" section - if let Some(weekly_pct) = extract_percent_near_label(&clean, "current week (all models)") - .or_else(|| extract_percent_near_label(&clean, "current week")) - { - weekly_percent = Some(weekly_pct); - } + let mut session_percent = extract_percent_near_label(&clean, "current session"); + let mut weekly_percent = WEEKLY_LABELS + .iter() + .find_map(|label| extract_percent_near_label(&clean, label)); // Fallback: collect all percentages in order. Activity stats carry // their own percentages, which must never be read as plan limits. if session_percent.is_none() && !activity_stats { - let all_percents = extract_all_percents(&clean); + let all_percents: Vec = percent_matches(&clean).collect(); if !all_percents.is_empty() { session_percent = Some(all_percents[0]); } @@ -1092,8 +489,9 @@ impl ClaudeProvider { // Extract reset times let session_reset = extract_reset_description(&clean, "current session"); - let weekly_reset = extract_reset_description(&clean, "current week (all models)") - .or_else(|| extract_reset_description(&clean, "current week")); + let weekly_reset = WEEKLY_LABELS + .iter() + .find_map(|label| extract_reset_description(&clean, label)); let short_form_reset = if is_exhausted_short_form(&clean_lower) { extract_inline_reset_description(&clean) } else { @@ -1234,1358 +632,5 @@ fn should_fallback_from_claude_cli_error(error: &ProviderError) -> bool { } } -/// Locate the Claude CLI for shell integrations that need to reopen a session. -pub fn locate_claude_binary() -> Option { - if let Some(path) = std::env::var_os("CLAUDE_BINARY") - .map(std::path::PathBuf::from) - .filter(|path| path.is_file()) - { - return Some(path); - } - - #[cfg(windows)] - { - let candidates = [ - // Direct install - dirs::data_local_dir().map(|p| p.join("Programs").join("claude").join("claude.exe")), - // npm global (AppData\Roaming\npm) - dirs::data_local_dir().map(|p| p.join("npm").join("claude.cmd")), - dirs::home_dir().map(|h| { - h.join("AppData") - .join("Roaming") - .join("npm") - .join("claude.cmd") - }), - // npm global alternate (~\.npm-global) - dirs::home_dir().map(|h| h.join(".npm-global").join("claude.cmd")), - // Volta managed - dirs::data_local_dir().map(|p| { - p.join("Volta") - .join("tools") - .join("image") - .join("packages") - .join("@anthropic-ai") - .join("claude-code") - .join("bin") - .join("claude.cmd") - }), - // fnm managed (via shim) - dirs::data_local_dir().map(|p| p.join("fnm_multishells").join("claude.cmd")), - // PATH lookup - find_windows_claude_in_path(), - ]; - - candidates.into_iter().flatten().find(|p| p.exists()) - } - - #[cfg(not(windows))] - { - which::which("claude").ok() - } -} - -#[cfg(windows)] -fn find_windows_claude_in_path() -> Option { - const CREATE_NO_WINDOW: u32 = 0x08000000; - - let mut command = StdCommand::new("where"); - command - .arg("claude") - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .creation_flags(CREATE_NO_WINDOW); - let output = command.output().ok()?; - - if !output.status.success() { - return None; - } - - let mut matches: Vec<_> = String::from_utf8_lossy(&output.stdout) - .lines() - .map(str::trim) - .filter(|line| !line.is_empty()) - .map(std::path::PathBuf::from) - .collect(); - - matches.sort_by_key(|path| { - match path - .extension() - .and_then(|ext| ext.to_str()) - .map(|ext| ext.to_ascii_lowercase()) - .as_deref() - { - Some("cmd") => 0, - Some("bat") => 1, - Some("exe") => 2, - _ => 3, - } - }); - - matches.into_iter().find(|path| path.exists()) -} - -/// Detect the version of the claude CLI -fn detect_claude_version() -> Option { - let claude_path = locate_claude_binary()?; - - #[cfg(windows)] - const CREATE_NO_WINDOW: u32 = 0x08000000; - - let mut cmd = std::process::Command::new(claude_path); - cmd.args(["--version"]); - #[cfg(windows)] - cmd.creation_flags(CREATE_NO_WINDOW); - - let output = cmd.output().ok()?; - - if output.status.success() { - let version_str = String::from_utf8_lossy(&output.stdout); - super::extract_semver(&version_str) - } else { - None - } -} - -/// Strip ANSI escape codes from text -fn strip_ansi(text: &str) -> String { - let mut result = String::with_capacity(text.len()); - let mut chars = text.chars().peekable(); - - while let Some(c) = chars.next() { - if c == '\x1B' { - // Skip CSI sequences: ESC[...letter - if chars.peek() == Some(&'[') { - chars.next(); - let mut final_char = None; - while let Some(&next) = chars.peek() { - chars.next(); - if next.is_ascii_alphabetic() { - final_char = Some(next); - break; - } - } - if final_char == Some('C') { - result.push(' '); - } - // Skip OSC sequences: ESC]...BEL - } else if chars.peek() == Some(&']') { - for next in chars.by_ref() { - if next == '\x07' || next == '\\' { - break; - } - } - } - } else { - result.push(c); - } - } - - result -} - -fn is_non_interactive_slash_command_response(text: &str) -> bool { - let mentions_usage_and_exit = text.contains("/usage") && text.contains("/exit"); - let says_entered_commands = - text.contains("i see you've entered") || text.contains("you've entered two slash commands"); - let says_no_slash_command = text.contains("available custom slash commands") - && text.contains("don't see these commands"); - let says_usage_is_cli_only = text - .contains("token usage and statistics are typically displayed by the cli interface") - || text.contains("i don't have direct access to those metrics"); - - mentions_usage_and_exit - && (says_entered_commands || says_no_slash_command || says_usage_is_cli_only) -} - -fn is_workspace_trust_prompt(text: &str) -> bool { - text.contains("quick safety check") - && text.contains("trust this folder") - && text.contains("yes, i trust this folder") -} - -/// Current `/usage` panels show a local session summary above the plan limits, -/// so the presence of a limit section outranks the activity-stats markers. -fn has_plan_limit_section(text: &str) -> bool { - text.contains("current session") || text.contains("current week") -} - -fn is_cli_activity_stats_response(text: &str) -> bool { - let has_activity_overview = text.contains("favorite model:") || text.contains("total tokens:"); - let has_session_cost_summary = - text.contains("total duration") && text.contains("usage:") && text.contains("cache read"); - - has_activity_overview || has_session_cost_summary -} - -/// Extract percentage near a label (e.g., "Current session") -/// Returns the percentage as "used" (not remaining) -fn extract_percent_near_label(text: &str, label: &str) -> Option { - let label_normalized = normalized_for_label_search(label); - let lines: Vec<&str> = text.lines().collect(); - - // Find the line containing the label - for (idx, line) in lines.iter().enumerate() { - if normalized_for_label_search(line).contains(&label_normalized) { - // Look in the next few lines for a percentage - for (offset, next_line) in lines.iter().skip(idx).take(12).enumerate() { - if offset > 0 && starts_next_usage_section(next_line, &label_normalized) { - break; - } - if let Some(pct) = parse_percent_line(next_line) { - return Some(pct); - } - } - } - } - - None -} - -/// Extract all percentages from text in order -fn extract_all_percents(text: &str) -> Vec { - let re = match Regex::new( - r"(\d{1,3}(?:\.\d+)?)\s*%\s*(used|spent|consumed|left|remaining|available)", - ) { - Ok(r) => r, - Err(_) => return vec![], - }; - - let mut results = Vec::new(); - let lower = text.to_lowercase(); - - for caps in re.captures_iter(&lower) { - if let (Some(val_match), Some(kind_match)) = (caps.get(1), caps.get(2)) - && let Ok(val) = val_match.as_str().parse::() - { - let kind = kind_match.as_str(); - let used = if matches!(kind, "left" | "remaining" | "available") { - (100.0 - val).max(0.0) - } else { - val.min(100.0) - }; - results.push(used); - } - } - - results -} - -fn is_exhausted_short_form(clean_lower: &str) -> bool { - clean_lower.contains("out of extra usage") || clean_lower.contains("hit your limit") -} - -/// Extract email address from text -fn extract_email(text: &str) -> Option { - // Try explicit patterns first - let patterns = [ - r"Account:\s*([^\s@]+@[^\s@]+\.[^\s]+)", - r"Email:\s*([^\s@]+@[^\s@]+\.[^\s]+)", - r"([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,})", - ]; - - for pattern in patterns { - if let Ok(re) = Regex::new(pattern) - && let Some(caps) = re.captures(text) - && let Some(m) = caps.get(1) - { - return Some(m.as_str().trim().to_string()); - } - } - - None -} - -/// Extract login method / plan name from text -fn extract_login_method(text: &str) -> Option { - // Look for explicit "Login method:" line - if let Ok(re) = Regex::new(r"(?i)login\s+method:\s*(.+)") - && let Some(caps) = re.captures(text) - && let Some(m) = caps.get(1) - { - let method = m.as_str().trim(); - if !method.is_empty() { - return Some(clean_plan_name(method)); - } - } - - // Look for "Claude " patterns - if let Ok(re) = Regex::new(r"(?i)(claude\s+(?:max|pro|ultra|team|free)[a-z0-9\s._-]*)") - && let Some(caps) = re.captures(text) - && let Some(m) = caps.get(1) - { - let plan = m.as_str().trim(); - if !plan.to_lowercase().contains("code") { - return Some(clean_plan_name(plan)); - } - } - - None -} - -/// Extract reset description near a label -fn extract_reset_description(text: &str, label: &str) -> Option { - let label_normalized = normalized_for_label_search(label); - let lines: Vec<&str> = text.lines().collect(); - - for (idx, line) in lines.iter().enumerate() { - if normalized_for_label_search(line).contains(&label_normalized) { - // Look in the next few lines for "Resets" - for (offset, next_line) in lines.iter().skip(idx).take(14).enumerate() { - if offset > 0 && starts_next_usage_section(next_line, &label_normalized) { - break; - } - let lower = next_line.to_lowercase(); - if lower.contains("resets") { - // Extract the reset info - if let Some(pos) = lower.find("resets") { - let reset_part = &next_line[pos..]; - return Some(reset_part.trim().to_string()); - } - } - } - } - } - - None -} - -/// Extract a "resets ..." suffix from a short single-line status. -fn extract_inline_reset_description(text: &str) -> Option { - let lower = text.to_lowercase(); - let pos = lower.find("resets")?; - Some(text[pos..].trim().to_string()) -} - -/// Clean up a plan name by removing ANSI codes and extra whitespace -fn clean_plan_name(text: &str) -> String { - let cleaned = strip_ansi(text); - // Remove bracketed codes like [22m - let re = Regex::new(r"\[\d+m").unwrap_or_else(|_| Regex::new(".^").unwrap()); - let result = re.replace_all(&cleaned, ""); - result.trim().to_string() -} - #[cfg(test)] -mod tests { - use chrono::{DateTime, Utc}; - use std::collections::HashMap; - - use super::*; - - const LOGIN_A: &str = "login-a"; - - fn login_a() -> Option { - Some(LOGIN_A.to_string()) - } - - #[test] - fn probe_cache_roundtrip_and_expiry() { - let dir = tempfile::tempdir().unwrap(); - assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); - store_cached_probe_output(dir.path(), LOGIN_A, "Current session 12% used"); - assert_eq!( - load_cached_probe_output(dir.path(), LOGIN_A).as_deref(), - Some("Current session 12% used") - ); - let stale = ClaudeProbeCache { - captured_at_unix: unix_now_secs() - CLAUDE_PROBE_CACHE_TTL.as_secs() - 5, - login: LOGIN_A.to_string(), - output: "Current session 12% used".to_string(), - }; - std::fs::write( - dir.path().join(CLAUDE_PROBE_CACHE_FILE), - serde_json::to_string(&stale).unwrap(), - ) - .unwrap(); - assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); - } - - #[test] - fn probe_cache_is_never_shared_with_another_login() { - let dir = tempfile::tempdir().unwrap(); - store_cached_probe_output(dir.path(), LOGIN_A, "Current session 12% used"); - assert!(load_cached_probe_output(dir.path(), "login-b").is_none()); - - // Written before screens were scoped to a login. - let unscoped = format!( - r#"{{"captured_at_unix":{},"output":"Current session 12% used"}}"#, - unix_now_secs() - ); - std::fs::write(dir.path().join(CLAUDE_PROBE_CACHE_FILE), unscoped).unwrap(); - assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); - assert!(load_cached_probe_output(dir.path(), "").is_none()); - } - - #[test] - fn login_fingerprint_follows_credential_rewrites_without_reading_them() { - let dir = tempfile::tempdir().unwrap(); - let credentials = dir.path().join(".credentials.json"); - assert_eq!(login_fingerprint_at(&credentials), None); - - std::fs::write(&credentials, "{}").unwrap(); - let first = login_fingerprint_at(&credentials).expect("fingerprint"); - assert_eq!(login_fingerprint_at(&credentials).as_ref(), Some(&first)); - assert!(!first.contains(".credentials"), "only a digest is stored"); - - std::fs::write(&credentials, r#"{"another":"login"}"#).unwrap(); - assert_ne!(login_fingerprint_at(&credentials), Some(first)); - } - - const SHAREABLE_USAGE_SCREEN: &str = "Current session\n\ - ████████▌ 17% used\n\ - Resets 12pm (America/Bogota)\n"; - - #[test] - fn locked_probe_reuses_a_screen_stored_while_it_waited() { - let dir = tempfile::tempdir().unwrap(); - store_cached_probe_output(dir.path(), LOGIN_A, SHAREABLE_USAGE_SCREEN); - - let output = run_locked_probe(dir.path(), Some(&login_a), || { - panic!("a fresh shared screen must not launch another probe") - }) - .unwrap(); - assert_eq!(output, SHAREABLE_USAGE_SCREEN); - } - - #[test] - fn locked_probe_shares_only_parseable_usage_screens() { - let dir = tempfile::tempdir().unwrap(); - let output = run_locked_probe(dir.path(), Some(&login_a), || { - Ok("Not logged in".to_string()) - }); - assert_eq!(output.unwrap(), "Not logged in"); - assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); - - let output = run_locked_probe(dir.path(), Some(&login_a), || { - Ok(SHAREABLE_USAGE_SCREEN.into()) - }); - assert_eq!(output.unwrap(), SHAREABLE_USAGE_SCREEN); - assert_eq!( - load_cached_probe_output(dir.path(), LOGIN_A).as_deref(), - Some(SHAREABLE_USAGE_SCREEN) - ); - assert!( - !dir.path() - .read_dir() - .unwrap() - .flatten() - .any(|entry| entry.file_name().to_string_lossy().contains(".tmp-")), - "the atomic write left no staging file behind" - ); - } - - #[test] - fn locked_probe_keeps_a_screen_private_when_the_login_changed_meanwhile() { - let dir = tempfile::tempdir().unwrap(); - let calls = std::cell::Cell::new(0); - let switching_login = || { - calls.set(calls.get() + 1); - Some(format!("login-{}", calls.get())) - }; - let output = run_locked_probe(dir.path(), Some(&switching_login), || { - Ok(SHAREABLE_USAGE_SCREEN.into()) - }); - assert_eq!(output.unwrap(), SHAREABLE_USAGE_SCREEN); - assert_eq!(calls.get(), 2, "the login is read before and after"); - assert!(load_cached_probe_output(dir.path(), "login-1").is_none()); - assert!(load_cached_probe_output(dir.path(), "login-2").is_none()); - - let no_login = || None; - run_locked_probe(dir.path(), Some(&no_login), || { - Ok(SHAREABLE_USAGE_SCREEN.into()) - }) - .unwrap(); - assert!(!dir.path().join(CLAUDE_PROBE_CACHE_FILE).exists()); - } - - #[test] - fn unshared_probe_neither_reuses_nor_stores_screens() { - let dir = tempfile::tempdir().unwrap(); - store_cached_probe_output(dir.path(), LOGIN_A, SHAREABLE_USAGE_SCREEN); - let output = run_locked_probe(dir.path(), None, || Ok("trust preflight".into())); - assert_eq!(output.unwrap(), "trust preflight"); - - let other = tempfile::tempdir().unwrap(); - run_locked_probe(other.path(), None, || Ok(SHAREABLE_USAGE_SCREEN.into())).unwrap(); - assert!(!other.path().join(CLAUDE_PROBE_CACHE_FILE).exists()); - } - - #[test] - fn logged_probe_screen_masks_account_email_and_secrets() { - let screen = "Login: someone@example.com (Claude Max)\n\ - access_token=abcdef0123456789 sk-ant-abcdefgh12345678\n\ - Current session 12% used"; - let logged = redacted_probe_screen(screen); - assert!(!logged.contains("someone@example.com"), "{logged}"); - assert!(!logged.contains("abcdef0123456789"), "{logged}"); - assert!(!logged.contains("sk-ant-abcdefgh12345678"), "{logged}"); - assert!(logged.contains("Current session 12% used")); - } - - #[test] - fn probe_lock_wait_expiry_fails_instead_of_probing_alongside() { - let dir = tempfile::tempdir().unwrap(); - let held = ClaudeProbeLock::acquire_within(dir.path(), Duration::ZERO) - .expect("first lock") - .expect("file locking is supported"); - - let error = match ClaudeProbeLock::acquire_within(dir.path(), Duration::from_millis(300)) { - Ok(lock) => panic!("second lock acquired while held: {}", lock.is_some()), - Err(error) => error, - }; - assert!(error.to_string().contains("Timed out waiting"), "{error}"); - assert_eq!( - last_good_failure_policy_for_error(&error.to_string()), - LastGoodFailurePolicy::Preserve - ); - - drop(held); - assert!( - ClaudeProbeLock::acquire_within(dir.path(), Duration::ZERO) - .expect("lock after release") - .is_some() - ); - } - - #[test] - fn passive_probe_env_disables_autoupdater_and_color() { - let env = claude_passive_probe_env(HashMap::new()); - assert_eq!( - env.get("DISABLE_AUTOUPDATER").map(String::as_str), - Some("1") - ); - assert_eq!(env.get("NO_COLOR").map(String::as_str), Some("1")); - } - - #[test] - fn probe_avoids_chrome_and_fullscreen_startup_dialogs() { - let env = claude_passive_probe_env(HashMap::new()); - assert_eq!( - env.get("CLAUDE_CODE_ENABLE_CFC").map(String::as_str), - Some("0") - ); - let settings: serde_json::Value = - serde_json::from_str(&claude_usage_settings_args()[1]).unwrap(); - assert_eq!(settings["tui"], "default"); - assert_eq!(settings["remoteControlAtStartup"], false); - } - - #[test] - fn probe_session_id_is_reused_from_probe_directory() { - let dir = tempfile::tempdir().unwrap(); - let first = load_or_create_probe_session_id(dir.path()); - let second = load_or_create_probe_session_id(dir.path()); - assert_eq!(first, second); - assert!(uuid::Uuid::parse_str(&first).is_ok()); - let args = claude_probe_launch_args(&first); - // Positional structure only: the settings pair is pinned once by - // `claude_usage_settings_args` being the sole composer. - assert_eq!( - args[..4], - ["--setting-sources", "user", "--allowed-tools", ""] - ); - assert_eq!(args[4], claude_usage_settings_args()[0]); - assert_eq!(args[5], claude_usage_settings_args()[1]); - assert_eq!(args[6], "--session-id"); - assert_eq!(args[7], first); - } - - #[test] - fn usage_probe_settings_disable_remote_control_startup() { - assert_eq!( - claude_usage_settings_args(), - [ - "--settings".to_string(), - r#"{"remoteControlAtStartup":false,"tui":"default"}"#.to_string(), - ] - ); - } - - #[test] - fn probe_session_jsonl_cleanup_removes_transcript_files() { - let dir = tempfile::tempdir().unwrap(); - let jsonl = dir.path().join("session.jsonl"); - std::fs::write(&jsonl, "{}").unwrap(); - std::fs::write(dir.path().join("keep.txt"), "x").unwrap(); - cleanup_probe_session_jsonl(dir.path()); - assert!(!jsonl.exists()); - assert!(dir.path().join("keep.txt").exists()); - } - - #[test] - fn probe_project_dir_name_matches_claude_code() { - use std::path::Path; - assert_eq!( - claude_project_dir_name(Path::new( - r"C:\Users\user\AppData\Local\CodexBar\claude-usage-probe" - )), - "C--Users-user-AppData-Local-CodexBar-claude-usage-probe" - ); - assert_eq!( - claude_project_dir_name(Path::new("/Users/me/Library/Application Support/x")), - "-Users-me-Library-Application-Support-x" - ); - // One dash per UTF-16 code unit, so two for a character outside the BMP. - assert_eq!( - claude_project_dir_name(Path::new("C:\\Users\\J\u{f6}rg\u{1F600}\\probe")), - "C--Users-J-rg---probe" - ); - // Reference values from Claude Code's JavaScript implementation. - let long = format!( - r"C:\Users\user\AppData\Local\{}claude-usage-probe", - r"deep\".repeat(40) - ); - assert_eq!( - claude_project_dir_name(Path::new(&long)), - format!( - "C--Users-user-AppData-Local-{}de-ttzy4x", - "deep-".repeat(34) - ) - ); - assert_eq!(javascript_hash_base36("hello"), "1n1e4y"); - assert_eq!(javascript_hash_base36(""), "0"); - } - - #[test] - fn probe_transcript_cleanup_stays_inside_the_probe_project() { - use std::path::Path; - let projects = tempfile::tempdir().unwrap(); - let other = projects.path().join("C--work-repo"); - std::fs::create_dir_all(&other).unwrap(); - std::fs::write(other.join("session.jsonl"), "{}").unwrap(); - - let busy_probe = Path::new(r"C:\Users\user\AppData\Local\CodexBar\busy-probe"); - let busy = projects.path().join(claude_project_dir_name(busy_probe)); - std::fs::create_dir_all(busy.join("folder.jsonl")).unwrap(); - std::fs::write(busy.join("session.jsonl"), "{}").unwrap(); - std::fs::write(busy.join("notes.txt"), "x").unwrap(); - cleanup_probe_transcripts_in(projects.path(), busy_probe); - assert!(!busy.join("session.jsonl").exists()); - assert!(busy.join("notes.txt").exists()); - assert!(busy.join("folder.jsonl").is_dir(), "only files are removed"); - - let probe = Path::new(r"C:\Users\user\AppData\Local\CodexBar\claude-usage-probe"); - let project = projects.path().join(claude_project_dir_name(probe)); - std::fs::create_dir_all(&project).unwrap(); - std::fs::write(project.join("a.jsonl"), "{}").unwrap(); - std::fs::write(project.join("b.jsonl"), "{}").unwrap(); - cleanup_probe_transcripts_in(projects.path(), probe); - assert!(!project.exists(), "an emptied probe project dir is removed"); - - assert!(other.join("session.jsonl").exists(), "other projects stay"); - } - - #[test] - fn parses_current_cli_usage_screen() { - let provider = ClaudeProvider::new(); - let output = r#" -Status Config Usage - - Current session - ██████████████████████████████████████████████████ 100% used - Resets 12pm (America/Bogota) - - Current week (all models) - ████████████████████████▌ 49% used - Resets Apr 3, 2pm (America/Bogota) - - Extra usage - ██▍ 4% used - $3.31 / $70.00 spent · Resets Apr 1 (America/Bogota) -"#; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.source_label, "cli"); - assert_eq!(result.usage.primary.used_percent, 100.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Resets 12pm (America/Bogota)") - ); - - let weekly = result - .usage - .secondary - .expect("weekly usage should be present"); - assert_eq!(weekly.used_percent, 49.0); - assert_eq!( - weekly.reset_description.as_deref(), - Some("Resets Apr 3, 2pm (America/Bogota)") - ); - } - - #[test] - fn parses_exhausted_short_form_as_full_session_usage() { - let provider = ClaudeProvider::new(); - let output = "You're out of extra usage · resets 12pm (America/Bogota)"; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.primary.used_percent, 100.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("resets 12pm (America/Bogota)") - ); - } - - #[test] - fn parses_hit_limit_short_form_as_full_session_usage() { - let provider = ClaudeProvider::new(); - let output = "You've hit your limit \u{00b7} resets 3:20pm (Asia/Shanghai)"; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.primary.used_percent, 100.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("resets 3:20pm (Asia/Shanghai)") - ); - } - - #[test] - fn parses_remaining_available_and_decimal_percentages() { - let provider = ClaudeProvider::new(); - let output = r#" -Status Config Usage - - Current session - 12.5% remaining - Resets 8pm - - Current week (all models) - 4% available - Resets Apr 4, 2pm - - Current week (Sonnet only) - 1% consumed -"#; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.primary.used_percent, 87.5); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Resets 8pm") - ); - - let weekly = result - .usage - .secondary - .expect("weekly usage should be present"); - assert_eq!(weekly.used_percent, 96.0); - assert_eq!( - weekly.reset_description.as_deref(), - Some("Resets Apr 4, 2pm") - ); - - let sonnet = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == "claude-weekly-scoped-sonnet") - .expect("sonnet usage should be present"); - assert_eq!(sonnet.window.used_percent, 1.0); - } - - #[test] - fn parses_all_cli_model_scoped_weekly_limits() { - let provider = ClaudeProvider::new(); - let output = r#" -Current session -10% used -Resets 12pm (America/Bogota) - -Current week (all models) -20% used -Resets Apr 3, 2pm (America/Bogota) - -Current week (Sonnet only) -30% used -Resets Apr 4, 2pm (America/Bogota) - -Current week (Opus only) -40% used -Resets Apr 5, 2pm (America/Bogota) -"#; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.extra_rate_windows.len(), 2); - assert_eq!( - result.usage.extra_rate_windows[0].id, - "claude-weekly-scoped-sonnet" - ); - assert_eq!(result.usage.extra_rate_windows[0].title, "Sonnet only"); - assert_eq!(result.usage.extra_rate_windows[0].window.used_percent, 30.0); - assert_eq!( - result.usage.extra_rate_windows[1].id, - "claude-weekly-scoped-opus" - ); - assert!(result.usage.model_specific.is_none()); - } - - #[test] - fn scoped_weekly_parser_handles_non_ascii_labels_and_reset_prefixes() { - let now = "2026-04-02T18:00:00Z".parse::>().unwrap(); - let limits = extract_cli_scoped_weekly_limits( - "Current week (A€€)\n10% used\nİResets Apr 3 at 2pm (America/Bogota)", - now, - ); - - assert_eq!(limits.len(), 1); - assert_eq!(limits[0].title, "A€€"); - assert_eq!( - limits[0].window.resets_at, - Some("2026-04-03T19:00:00Z".parse().unwrap()) - ); - } - - #[test] - fn resolves_cli_reset_occurrences_in_the_reported_timezone() { - let now = "2026-04-02T18:00:00Z".parse::>().unwrap(); - - assert_eq!( - parse_claude_reset_date("Resets Apr 3, 2027, 2pm (America/Bogota)", now, None), - Some("2027-04-03T19:00:00Z".parse().unwrap()) - ); - assert_eq!( - parse_claude_reset_date("Resets Apr 3, 2pm (America/Bogota)", now, None), - Some("2026-04-03T19:00:00Z".parse().unwrap()) - ); - assert_eq!( - parse_claude_reset_date("Resets 12pm (America/Bogota)", now, None), - Some("2026-04-03T17:00:00Z".parse().unwrap()) - ); - assert_eq!( - parse_claude_reset_date("ResetsApr3at2pm(America/Bogota)", now, None), - Some("2026-04-03T19:00:00Z".parse().unwrap()) - ); - } - - #[test] - fn timezone_less_resets_use_the_supplied_system_zone() { - let now = "2026-03-07T18:00:00Z".parse::>().unwrap(); - - assert_eq!( - parse_claude_reset_date_in_system_zone( - "Resets Mar 8 at 3:30am", - now, - None, - "America/New_York".parse().unwrap(), - ), - Some("2026-03-08T07:30:00Z".parse().unwrap()) - ); - assert_eq!( - parse_claude_reset_date_in_system_zone( - "Resets Mar 8 at 3:30am (America/Los_Angeles)", - now, - None, - "America/New_York".parse().unwrap(), - ), - Some("2026-03-08T10:30:00Z".parse().unwrap()) - ); - } - - #[test] - fn parses_compact_usage_screen() { - let provider = ClaudeProvider::new(); - let output = r#" -Settings:StatusConfigUsage(tabtocycle) -Loadingusagedata... -Currentsession -6%used -Resets4:29am(Asia/Calcutta) -Currentweek(allmodels) -4%used -ResetsFeb12at1:29pm(Asia/Calcutta) -Currentweek(Sonnetonly) -1%used -ResetsFeb12at1:29pm(Asia/Calcutta) -"#; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.primary.used_percent, 6.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Resets4:29am(Asia/Calcutta)") - ); - assert_eq!( - result - .usage - .secondary - .expect("weekly usage should be present") - .used_percent, - 4.0 - ); - let sonnet = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == "claude-weekly-scoped-sonnet") - .expect("sonnet usage should be present"); - assert_eq!(result.usage.extra_rate_windows.len(), 1); - assert_eq!(sonnet.title, "Sonnet only"); - assert_eq!(sonnet.window.used_percent, 1.0); - } - - #[test] - fn does_not_promote_weekly_reset_to_session() { - let provider = ClaudeProvider::new(); - let output = r#" -Current session -17% used -Current week (all models) -4% used -Resets Dec 24 at 3:59pm (Europe/Paris) -"#; - - let result = provider.parse_cli_output(output).expect("should parse"); - - assert_eq!(result.usage.primary.used_percent, 17.0); - assert_eq!(result.usage.primary.reset_description, None); - assert_eq!( - result - .usage - .secondary - .expect("weekly usage should be present") - .reset_description - .as_deref(), - Some("Resets Dec 24 at 3:59pm (Europe/Paris)") - ); - } - - #[test] - fn rejects_cli_output_without_usage_markers() { - let provider = ClaudeProvider::new(); - let output = "Claude Code on Windows requires git-bash."; - - let err = provider - .parse_cli_output(output) - .expect_err("should reject non-usage output"); - - assert!(matches!(err, ProviderError::Parse(_))); - assert_eq!( - err.to_string(), - "Parse error: Claude CLI did not return usage data" - ); - } - - #[test] - fn cli_parse_usage_error_can_fallback_to_oauth() { - let err = ProviderError::Parse("Claude CLI did not return usage data".to_string()); - - assert!(should_fallback_from_claude_cli_error(&err)); - } - - #[test] - fn cli_auth_error_does_not_fallback_to_oauth() { - assert!(!should_fallback_from_claude_cli_error( - &ProviderError::AuthRequired - )); - } - - #[test] - fn auto_fetch_error_keeps_all_source_failures() { - let err = claude_auto_fetch_error(vec![ - ("OAuth", ProviderError::OAuth("token expired".to_string())), - ("Web", ProviderError::NoCookies), - ( - "CLI", - ProviderError::Parse("Empty output from Claude CLI".to_string()), - ), - ]); - - assert_eq!( - err.to_string(), - "Claude usage failed from all configured sources. OAuth: OAuth error: token expired; Web: No cookies available for web API; CLI: Parse error: Empty output from Claude CLI" - ); - } - - fn oauth_rate_limited() -> ProviderError { - ClaudeOAuthFetcher::rate_limited_error(Duration::from_secs(30)) - } - - #[test] - fn auto_fetch_error_asks_for_a_browser_sign_in_when_only_the_browser_can_help() { - // (CLI failure, retention policy of the plain summary) - let cases = [ - ( - ProviderError::Parse("Claude CLI did not return usage data".to_string()), - LastGoodFailurePolicy::Preserve, - ), - ( - ProviderError::Other("Claude CLI failed: exit status 1".to_string()), - LastGoodFailurePolicy::Replace, - ), - ]; - for (cli_failure, policy) in cases { - let err = claude_auto_fetch_error(vec![ - ("Web", ProviderError::NoCookies), - ("OAuth", oauth_rate_limited()), - ("CLI", cli_failure), - ]); - let ProviderError::BrowserSignInRequired { - message, - sign_in_url, - } = &err - else { - panic!("expected a browser sign-in signal, got {err:?}"); - }; - assert_eq!(sign_in_url, CLAUDE_BROWSER_SIGN_IN_URL); - assert_eq!(err.to_string(), *message); - assert!( - message.starts_with( - "Claude usage failed from all configured sources. Web: No cookies available for web API; OAuth: Transient OAuth error: Claude OAuth usage endpoint is rate limited." - ), - "{message}" - ); - assert!( - message - .ends_with("Sign in at https://claude.ai/login in your browser, then refresh."), - "{message}" - ); - // ClaudeProvider::error_state_kind defers to this for every - // variant except a missing CLI. - assert_eq!( - err.state_kind(), - crate::core::ProviderStateKind::NeedsAuthentication - ); - // The hint leaves the desktop retention policy unchanged. - let plain = message - .strip_suffix(browser_sign_in_hint().as_str()) - .map(str::trim_end) - .expect("hint is appended"); - assert_eq!(last_good_failure_policy_for_error(plain), policy); - assert_eq!(last_good_failure_policy_for_error(message), policy); - } - } - - #[test] - fn auto_fetch_error_keeps_other_failure_mixes_untyped() { - let cli_failure = - || ProviderError::Parse("Claude CLI did not return usage data".to_string()); - let mixes = [ - // A browser session was there; the Web source failed differently. - vec![ - ("Web", ProviderError::AuthRequired), - ("OAuth", oauth_rate_limited()), - ("CLI", cli_failure()), - ], - // Signed out of Claude Code, not rate limited. - vec![ - ("Web", ProviderError::NoCookies), - ( - "OAuth", - ProviderError::OAuth( - "Claude OAuth credentials not found. Run `claude` to authenticate." - .to_string(), - ), - ), - ("CLI", cli_failure()), - ], - // Another transient OAuth failure. - vec![ - ("Web", ProviderError::NoCookies), - ( - "OAuth", - ProviderError::OAuthTransient( - "Claude OAuth token expired and token refresh is cooling down after a failed attempt." - .to_string(), - ), - ), - ("CLI", cli_failure()), - ], - // The CLI was not tried. - vec![ - ("Web", ProviderError::NoCookies), - ("OAuth", oauth_rate_limited()), - ], - ]; - for failures in mixes { - let err = claude_auto_fetch_error(failures); - assert!(matches!(err, ProviderError::Other(_)), "{err:?}"); - assert!( - !err.to_string().contains(CLAUDE_BROWSER_SIGN_IN_URL), - "{err}" - ); - } - } - - #[test] - fn transient_transport_failure_stops_auto_fallback_and_preserves_last_good() { - let provider = ClaudeProvider::new(); - assert!(provider.retains_last_good_on_transport_failure()); - assert_eq!( - provider.last_good_failure_policy_for_error(&ProviderError::Timeout), - LastGoodFailurePolicy::Preserve - ); - - let mut failures = Vec::new(); - let result = record_auto_source(&mut failures, "Web", Err(ProviderError::Timeout)); - assert!(matches!(result, Err(ProviderError::Timeout))); - assert!(failures.is_empty()); - } - - #[test] - fn rejects_claude_2_1_non_interactive_slash_response() { - let provider = ClaudeProvider::new(); - let output = r#" -I see you've entered `/usage` and `/exit`. - -**Usage**: Token usage and statistics are typically displayed by the CLI interface itself. I don't have direct access to those metrics through my available tools. - -**Exit**: I'll end the session here. Goodbye! -"#; - - let err = provider - .parse_cli_output(output) - .expect_err("should reject non-interactive slash command response"); - - assert!(matches!(err, ProviderError::Other(_))); - assert_eq!( - err.to_string(), - "Claude CLI treated /usage as a normal prompt instead of opening the interactive usage screen. Use Auto, OAuth, or Web mode for Claude usage." - ); - } - - #[test] - fn rejects_legacy_non_interactive_slash_response() { - let provider = ClaudeProvider::new(); - let output = r#" -I see you've entered two slash commands: - -1. `/usage` - This appears to be a request to check usage information -2. `/exit` - This appears to be a request to exit - -However, looking at the available custom slash commands, I don't see these commands defined. -"#; - - let err = provider - .parse_cli_output(output) - .expect_err("should reject non-interactive slash command response"); - - assert!(matches!(err, ProviderError::Other(_))); - } - - #[test] - fn rejects_cli_activity_stats_without_plan_limits() { - let provider = ClaudeProvider::new(); - let output = r#" -❯ /usage - -Status Config Usage Stats - -Overview Models - -Favorite model: glm-4.6 Total tokens: 263.3k -Sessions: 6 Longest session: 18s -Active days: 2/10 Longest streak: 1 day -"#; - - let err = provider - .parse_cli_output(output) - .expect_err("should reject local activity stats"); - - assert!(matches!(err, ProviderError::Other(_))); - assert_eq!( - err.to_string(), - "Claude CLI /usage opened, but this Claude version returned local activity stats instead of plan limit percentages. Use Auto, OAuth, or Web mode for Claude limits." - ); - } - - #[test] - fn rejects_ansi_spaced_cli_activity_stats_without_plan_limits() { - let provider = ClaudeProvider::new(); - let output = "\x1b[2CTotal\x1b[1Ccost:\x1b[12C$0.0000\n\ - \x1b[2CTotal\x1b[1Cduration\x1b[1C(API):\x1b[2C0s\n\ - \x1b[2CUsage:\x1b[17C0\x1b[1Cinput,\x1b[1C0\x1b[1Coutput,\x1b[1C0\x1b[1Ccache\x1b[1Cread"; - - let err = provider - .parse_cli_output(output) - .expect_err("should reject ANSI-spaced local activity stats"); - - assert!(matches!(err, ProviderError::Other(_))); - } - - #[test] - fn accepts_plan_limits_followed_by_activity_stats() { - // Claude Code 2.1.27x on Windows prints the exit summary (cost, - // duration, cache tokens) after the /usage view when the probe ends. - let provider = ClaudeProvider::new(); - let output = r#" -❯ /usage - -Status Config Usage Stats - -Current session -███████░░░░░░░░░░░░░░░░░░░░░░ 19% used -Resets 3pm (Europe/Berlin) - -Current week (all models) -█████████░░░░░░░░░░░░░░░░░░░░ 31% used -Resets Sep 19, 4pm (Europe/Berlin) - -Total cost: $0.0000 -Total duration (API): 0s -Usage: 0 input, 0 output, 0 cache read -"#; - - let result = provider - .parse_cli_output(output) - .expect("plan limits should win over trailing activity stats"); - - assert_eq!(result.usage.primary.used_percent, 19.0); - assert_eq!( - result - .usage - .secondary - .as_ref() - .map(|window| window.used_percent), - Some(31.0) - ); - } - - // ── Upstream 0.50.1 #2516: revoked vs missing OAuth ──────────────────────── - - #[test] - fn oauth_revoked_error_is_detected() { - assert!(is_oauth_revoked_error(&ProviderError::OAuthRevoked( - "revoked".to_string() - ))); - assert!(!is_oauth_revoked_error(&ProviderError::OAuth( - "expired".to_string() - ))); - assert!(!is_oauth_revoked_error(&ProviderError::AuthRequired)); - } - - #[test] - fn rate_limited_and_revoked_oauth_reuse_the_cli_cache() { - let rate_limited = ProviderError::OAuthTransient( - "Claude OAuth usage endpoint is rate limited. Retrying in about 5m; credentials were preserved." - .to_string(), - ); - assert!(oauth::is_rate_limited_error(&rate_limited)); - assert!(oauth_failure_uses_cli_cache(&rate_limited)); - assert!(oauth_failure_uses_cli_cache(&ProviderError::OAuthRevoked( - "revoked".to_string() - ))); - // Other transient failures and plain expiry still probe the CLI. - assert!(!oauth_failure_uses_cli_cache( - &ProviderError::OAuthTransient("connection reset".to_string()) - )); - assert!(!oauth_failure_uses_cli_cache(&ProviderError::OAuth( - "expired".to_string() - ))); - assert!(!oauth_failure_uses_cli_cache(&ProviderError::AuthRequired)); - } - - #[test] - fn cli_result_cache_round_trips() { - let mut result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "cli"); - result.has_successful_claude_cli_quota = true; - cache_cli_result(result.clone()); - let cached = cached_cli_result().expect("cached result within TTL"); - assert!((cached.usage.primary.used_percent - 42.0).abs() < 0.01); - assert_eq!(cached.source_label, "cli"); - assert!(!cached.has_successful_claude_cli_quota); - - // A live non-CLI success clears the cache. Same test, because the - // global is shared and tests run in parallel without a lock. - clear_cli_result_cache(); - assert!(cached_cli_result().is_none()); - } - - #[test] - fn cli_quota_without_credential_identity_cannot_prove_account_action() { - let provider = ClaudeProvider::new(); - let result = provider - .parse_cli_output("Current session\n25% used\nCurrent week (all models)\n40% used") - .expect("CLI quota should parse"); - let result = mark_live_claude_cli_result(result); - - assert!(result.usage.account_email.is_none()); - assert!(result.has_successful_claude_cli_quota); - } - - #[test] - fn non_cli_fetch_result_does_not_prove_account_action() { - let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "oauth"); - - assert!(!result.has_successful_claude_cli_quota); - } - #[test] - fn cli_presence_maps_to_local_runtime_offline() { - assert_eq!( - ClaudeProvider::new().error_state_kind(&ProviderError::NotInstalled( - "Claude CLI not found. Install from https://docs.claude.ai/claude-code".to_string(), - )), - crate::core::ProviderStateKind::LocalRuntimeOffline - ); - // Other error kinds keep their default classification. - assert_eq!( - ClaudeProvider::new().error_state_kind(&ProviderError::AuthRequired), - crate::core::ProviderStateKind::NeedsAuthentication - ); - } - - #[test] - fn oauth_rate_limit_is_not_sign_in_required() { - let error = ProviderError::OAuthTransient( - "OAuth error: Claude OAuth usage endpoint is rate limited. Retrying in about 1s; credentials were preserved." - .to_string(), - ); - assert_eq!( - ClaudeProvider::new().error_state_kind(&error), - crate::core::ProviderStateKind::Unknown - ); - assert_eq!( - ClaudeProvider::new().last_good_failure_policy_for_error(&error), - LastGoodFailurePolicy::Preserve - ); - } - - #[test] - fn oauth_refresh_cooldown_is_not_sign_in_required() { - let error = ProviderError::OAuthTransient( - "Claude OAuth token expired and token refresh is cooling down after a failed attempt. Please retry shortly, or run `claude login`." - .to_string(), - ); - assert_eq!( - ClaudeProvider::new().error_state_kind(&error), - crate::core::ProviderStateKind::Unknown - ); - assert_eq!( - ClaudeProvider::new().last_good_failure_policy_for_error(&error), - LastGoodFailurePolicy::Preserve - ); - } - - #[test] - fn missing_oauth_credentials_still_require_sign_in() { - let error = ProviderError::OAuth( - "Claude OAuth credentials not found. Run `claude` to authenticate.".to_string(), - ); - assert_eq!( - ClaudeProvider::new().error_state_kind(&error), - crate::core::ProviderStateKind::NeedsAuthentication - ); - assert_eq!( - last_good_failure_policy_for_error(&error.to_string()), - LastGoodFailurePolicy::Replace - ); - } - - #[test] - fn untyped_oauth_rate_limit_text_is_not_transient() { - let error = ProviderError::OAuth("OAuth API returned rate limited".to_string()); - assert_eq!( - ClaudeProvider::new().error_state_kind(&error), - crate::core::ProviderStateKind::NeedsAuthentication - ); - assert_eq!( - ClaudeProvider::new().last_good_failure_policy_for_error(&error), - LastGoodFailurePolicy::Replace - ); - } -} +mod tests; diff --git a/rust/src/providers/claude/oauth/credentials_store.rs b/rust/src/providers/claude/oauth/credentials_store.rs index 19056585fa..b9f6770735 100644 --- a/rust/src/providers/claude/oauth/credentials_store.rs +++ b/rust/src/providers/claude/oauth/credentials_store.rs @@ -460,6 +460,27 @@ fn apply_refresh_to_credentials_json( #[cfg(test)] mod tests { + use super::{ + CredentialSource, apply_refresh_to_credentials_json, cached_refreshed_if_fresher, + parse_credentials_json, replacement_from_changed_fresh_keyring, store_refreshed, + }; + use crate::providers::claude::oauth::ClaudeOAuthCredentials; + use chrono::{DateTime, Duration, Utc}; + + fn creds( + access_token: &str, + refresh_token: Option<&str>, + expires_at: Option>, + ) -> ClaudeOAuthCredentials { + ClaudeOAuthCredentials { + access_token: access_token.to_string(), + refresh_token: refresh_token.map(str::to_string), + expires_at, + scopes: vec!["user:profile".to_string()], + rate_limit_tier: None, + } + } + #[test] fn a_new_saved_login_cannot_inherit_an_older_longer_lived_rotation() { let path = std::path::PathBuf::from(format!("saved-cache-{}", uuid::Uuid::new_v4())); @@ -482,11 +503,6 @@ mod tests { assert!(super::cached_refreshed_if_fresher(&old_source, &shorter_login).is_some()); assert!(super::cached_refreshed_if_fresher(&new_source, &shorter_login).is_none()); } - use super::{ - CredentialSource, apply_refresh_to_credentials_json, cached_refreshed_if_fresher, - parse_credentials_json, replacement_from_changed_fresh_keyring, store_refreshed, - }; - use crate::providers::claude::oauth::ClaudeOAuthCredentials; #[test] fn parses_claude_code_credentials_payload() { @@ -548,21 +564,17 @@ mod tests { #[test] fn changed_fresh_keyring_replaces_expired_file_credentials() { - let expired_file = ClaudeOAuthCredentials { - access_token: "expired-file-token".to_string(), - refresh_token: Some("expired-file-refresh".to_string()), - expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)), - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let expired_file = creds( + "expired-file-token", + Some("expired-file-refresh"), + Some(Utc::now() - Duration::hours(1)), + ); let keyring_source = CredentialSource::Keyring("test-user".to_string()); - let fresh_keyring = ClaudeOAuthCredentials { - access_token: "fresh-keyring-token".to_string(), - refresh_token: Some("fresh-keyring-refresh".to_string()), - expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let fresh_keyring = creds( + "fresh-keyring-token", + Some("fresh-keyring-refresh"), + Some(Utc::now() + Duration::hours(1)), + ); let adopted = replacement_from_changed_fresh_keyring( &expired_file, @@ -572,13 +584,11 @@ mod tests { assert_eq!(adopted.0.access_token, "fresh-keyring-token"); assert_eq!(adopted.1, keyring_source); - let expired_keyring = ClaudeOAuthCredentials { - access_token: "another-keyring-token".to_string(), - refresh_token: None, - expires_at: Some(chrono::Utc::now() - chrono::Duration::minutes(1)), - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let expired_keyring = creds( + "another-keyring-token", + None, + Some(Utc::now() - Duration::minutes(1)), + ); assert!( replacement_from_changed_fresh_keyring( &expired_file, @@ -645,22 +655,14 @@ mod tests { "env_source_not_shadowed_by_file_cache-unique-marker.json", )); - let file_cached_creds = ClaudeOAuthCredentials { - access_token: "file-refreshed-token".to_string(), - refresh_token: Some("file-refresh".to_string()), - expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)), - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let file_cached_creds = creds( + "file-refreshed-token", + Some("file-refresh"), + Some(Utc::now() + Duration::hours(1)), + ); store_refreshed(&file_source, &file_cached_creds); - let env_creds = ClaudeOAuthCredentials { - access_token: "env-token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let env_creds = creds("env-token", None, None); // Looking up under the Environment source must not see the File // source's cached (and "fresher"-by-the-naive-rule) entry. @@ -672,13 +674,7 @@ mod tests { // Sanity check: the file source's own cache entry is still there and // still considered fresher than a file-read with no expiry. - let file_disk_creds = ClaudeOAuthCredentials { - access_token: "file-disk-token".to_string(), - refresh_token: Some("file-disk-refresh".to_string()), - expires_at: None, - scopes: vec!["user:profile".to_string()], - rate_limit_tier: None, - }; + let file_disk_creds = creds("file-disk-token", Some("file-disk-refresh"), None); let same_source_result = cached_refreshed_if_fresher(&file_source, &file_disk_creds); assert_eq!( same_source_result.map(|c| c.access_token), diff --git a/rust/src/providers/claude/oauth/refresh.rs b/rust/src/providers/claude/oauth/refresh.rs index 98735b4bc9..10cf4ab49e 100644 --- a/rust/src/providers/claude/oauth/refresh.rs +++ b/rust/src/providers/claude/oauth/refresh.rs @@ -185,56 +185,35 @@ mod tests { // and 5xx stay transient (a retry can still heal them); 400/401 *without* // invalid_grant is likewise transient. #[test] - fn invalid_grant_on_400_or_401_is_terminal() { - for status in [400, 401] { + fn refresh_failure_kind_by_status_and_body() { + use RefreshFailureKind::{Terminal, Transient}; + const GRANT: &str = r#"{"error":"invalid_grant"}"#; + let rows = [ + (400, GRANT, Terminal), + (401, GRANT, Terminal), + // Case-insensitive match. + (400, r#"{"error":"INVALID_GRANT"}"#, Terminal), + // 403 is never terminal, even with invalid_grant in the body. + (403, GRANT, Transient), + // 400/401 with a different OAuth error are transient. + (400, r#"{"error":"invalid_client"}"#, Transient), + (401, r#"{"error":"invalid_client"}"#, Transient), + // 400/401 with no parseable error field are transient. + (400, "busy", Transient), + (401, "busy", Transient), + // Server and rate-limit errors. + (408, "busy", Transient), + (429, "busy", Transient), + (500, "busy", Transient), + (502, "busy", Transient), + (503, "busy", Transient), + ]; + for (status, body, kind) in rows { let failure = RefreshFailure::from_http_status( reqwest::StatusCode::from_u16(status).unwrap(), - r#"{"error":"invalid_grant"}"#, + body, ); - assert_eq!(failure.kind, RefreshFailureKind::Terminal, "HTTP {status}"); - } - // Case-insensitive match. - let failure = RefreshFailure::from_http_status( - reqwest::StatusCode::from_u16(400).unwrap(), - r#"{"error":"INVALID_GRANT"}"#, - ); - assert_eq!(failure.kind, RefreshFailureKind::Terminal); - } - - #[test] - fn forbidden_and_non_grant_4xx_stay_transient() { - // 403 is never terminal, even with invalid_grant in the body. - let failure = RefreshFailure::from_http_status( - reqwest::StatusCode::from_u16(403).unwrap(), - r#"{"error":"invalid_grant"}"#, - ); - assert_eq!(failure.kind, RefreshFailureKind::Transient); - // 400/401 with a different OAuth error are transient. - for status in [400, 401] { - let failure = RefreshFailure::from_http_status( - reqwest::StatusCode::from_u16(status).unwrap(), - r#"{"error":"invalid_client"}"#, - ); - assert_eq!(failure.kind, RefreshFailureKind::Transient, "HTTP {status}"); - } - // 400/401 with no parseable error field are transient. - for status in [400, 401] { - let failure = RefreshFailure::from_http_status( - reqwest::StatusCode::from_u16(status).unwrap(), - "busy", - ); - assert_eq!(failure.kind, RefreshFailureKind::Transient, "HTTP {status}"); - } - } - - #[test] - fn refresh_server_and_rate_limit_errors_stay_transient() { - for status in [408, 429, 500, 502, 503] { - let failure = RefreshFailure::from_http_status( - reqwest::StatusCode::from_u16(status).unwrap(), - "busy", - ); - assert_eq!(failure.kind, RefreshFailureKind::Transient, "HTTP {status}"); + assert_eq!(failure.kind, kind, "HTTP {status} {body}"); } } } diff --git a/rust/src/providers/claude/oauth/tests.rs b/rust/src/providers/claude/oauth/tests.rs index 3e3fa9ce5c..01a8bd8983 100644 --- a/rust/src/providers/claude/oauth/tests.rs +++ b/rust/src/providers/claude/oauth/tests.rs @@ -2,7 +2,7 @@ use super::{ ClaudeOAuthCredentials, ClaudeOAuthFetcher, OAuthUsageResponse, UsageWindow, credential_identity, is_rate_limited_error, }; -use crate::core::ProviderError; +use crate::core::{ProviderError, UsageSnapshot}; use base64::Engine; use reqwest::header::HeaderValue; use std::time::Duration; @@ -29,6 +29,19 @@ fn test_credentials(access_token: &str) -> ClaudeOAuthCredentials { } } +fn creds(rate_limit_tier: Option<&str>) -> ClaudeOAuthCredentials { + ClaudeOAuthCredentials { + scopes: vec![], + rate_limit_tier: rate_limit_tier.map(str::to_string), + ..test_credentials("token") + } +} + +fn snapshot(json: &str, credentials: &ClaudeOAuthCredentials) -> UsageSnapshot { + let response: OAuthUsageResponse = serde_json::from_str(json).expect("OAuth usage body"); + ClaudeOAuthFetcher::new().build_usage_snapshot(&response, credentials) +} + #[test] fn credential_identity_uses_jwt_subject_when_available() { let payload = @@ -53,56 +66,32 @@ fn opaque_credential_identity_is_a_non_secret_fingerprint() { } #[test] -fn keeps_sub_one_utilization_in_percent_units() { - let window = UsageWindow { - utilization: Some(0.23), - resets_at: None, - }; - - let rate = ClaudeOAuthFetcher::to_rate_window(&window, Some(300)).expect("rate window"); - - assert!((rate.used_percent - 0.23).abs() < f64::EPSILON); -} - -#[test] -fn one_percent_session_is_not_reported_as_full_quota() { - let window = UsageWindow { - utilization: Some(1.0), - resets_at: None, - }; - - let rate = ClaudeOAuthFetcher::to_rate_window(&window, Some(300)).expect("rate window"); - - assert!( - (rate.used_percent - 1.0).abs() < f64::EPSILON, - "session was {}, expected 1% (not 100%)", - rate.used_percent - ); -} - -#[test] -fn preserves_existing_percentage_utilization() { - let window = UsageWindow { - utilization: Some(23.0), - resets_at: None, - }; - - let rate = ClaudeOAuthFetcher::to_rate_window(&window, Some(300)).expect("rate window"); - - assert!((rate.used_percent - 23.0).abs() < f64::EPSILON); +fn keeps_utilization_in_percent_units() { + // 1.0 is a 1% session, not a full quota. + for utilization in [0.23, 1.0, 23.0] { + let window = UsageWindow { + utilization: Some(utilization), + resets_at: None, + }; + + let rate = ClaudeOAuthFetcher::to_rate_window(&window, Some(300)).expect("rate window"); + + assert!( + (rate.used_percent - utilization).abs() < f64::EPSILON, + "session was {}, expected {utilization}% (not 100%)", + rate.used_percent + ); + } } #[test] fn missing_oauth_session_is_informational_and_keeps_weekly_lane() { - let response: OAuthUsageResponse = serde_json::from_str( + let usage = snapshot( r#"{ "seven_day": {"utilization": 51.0, "resets_at": "2026-08-20T12:00:00Z"} }"#, - ) - .expect("OAuth response without a session lane should parse"); - - let usage = - ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &test_credentials("token")); + &test_credentials("token"), + ); assert!(usage.primary.is_informational); assert_eq!(usage.primary.window_minutes, Some(300)); @@ -115,7 +104,11 @@ fn missing_oauth_session_is_informational_and_keeps_weekly_lane() { #[test] fn parses_current_snake_case_oauth_usage_response() { - let response: OAuthUsageResponse = serde_json::from_str( + let credentials = ClaudeOAuthCredentials { + rate_limit_tier: Some("default_claude_ai".to_string()), + ..test_credentials("token") + }; + let usage = snapshot( r#"{ "five_hour": {"utilization": 1.0, "resets_at": "2026-05-22T22:10:00Z"}, "seven_day": {"utilization": 0.14, "resets_at": "2026-05-29T10:00:00Z"}, @@ -130,17 +123,8 @@ fn parses_current_snake_case_oauth_usage_response() { }], "extra_usage": {"is_enabled": true, "used_credits": 0, "monthly_limit": 1000, "currency": "USD"} }"#, - ) - .expect("snake_case OAuth response should parse"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec!["user:profile".to_string()], - rate_limit_tier: Some("default_claude_ai".to_string()), - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + &credentials, + ); assert_eq!(usage.primary.used_percent, 1.0); assert!((usage.secondary.expect("weekly").used_percent - 0.14).abs() < 0.001); @@ -155,7 +139,8 @@ fn parses_current_snake_case_oauth_usage_response() { #[test] fn weekly_all_limit_wins_over_stale_seven_day_utilization() { - let response: OAuthUsageResponse = serde_json::from_str( + let credentials = creds(Some("default_claude_max_5x")); + let usage = snapshot( r#"{ "five_hour": {"utilization": 8.0, "resets_at": "2026-07-20T04:29:59Z"}, "seven_day": {"utilization": 1.0, "resets_at": "2026-07-26T22:59:59Z"}, @@ -175,17 +160,8 @@ fn weekly_all_limit_wins_over_stale_seven_day_utilization() { } ] }"#, - ) - .expect("oauth body with weekly_all"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: Some("default_claude_max_5x".to_string()), - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + &credentials, + ); assert!((usage.primary.used_percent - 8.0).abs() < f64::EPSILON); // seven_day.utilization 1.0 would normalize to 100%; weekly_all wins. @@ -204,7 +180,8 @@ fn weekly_all_limit_wins_over_stale_seven_day_utilization() { fn issue_210_reporter_shape_secondary_is_one_percent_not_one_hundred() { // Mirrors the reporter JSON: session 8%, fable 2%, all-models should be 1% // while seven_day.utilization is the stale 1.0 (would display as 100%). - let response: OAuthUsageResponse = serde_json::from_str( + let credentials = creds(Some("default_claude_max_5x")); + let usage = snapshot( r#"{ "five_hour": { "utilization": 8.0, @@ -235,17 +212,8 @@ fn issue_210_reporter_shape_secondary_is_one_percent_not_one_hundred() { } ] }"#, - ) - .expect("issue 210 body"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: Some("default_claude_max_5x".to_string()), - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + &credentials, + ); assert_eq!(usage.login_method.as_deref(), Some("Claude Max 5x")); assert!((usage.primary.used_percent - 8.0).abs() < f64::EPSILON); @@ -270,7 +238,8 @@ fn issue_279_session_limits_win_over_stale_five_hour_after_rollover() { // can transiently report 1.0 (normalizes to 100%) even though // claude.ai shows only 5% for the fresh window. The limits[] entry // (kind=="session") carries the true value and must win. - let response: OAuthUsageResponse = serde_json::from_str( + let credentials = creds(Some("default_claude_max_5x")); + let usage = snapshot( r#"{ "five_hour": {"utilization": 1.0, "resets_at": "2026-08-13T12:49:59.578826Z"}, "seven_day": {"utilization": 0.01, "resets_at": "2026-07-26T22:59:59Z"}, @@ -289,17 +258,8 @@ fn issue_279_session_limits_win_over_stale_five_hour_after_rollover() { } ] }"#, - ) - .expect("issue 279 body"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: Some("default_claude_max_5x".to_string()), - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + &credentials, + ); // Primary session must be 5%, not the stale 100%. assert!( @@ -320,40 +280,29 @@ fn issue_279_session_limits_win_over_stale_five_hour_after_rollover() { fn session_falls_back_to_legacy_five_hour_without_limits_entry() { // When no limits[] session entry exists, the legacy five_hour field // is still the source of truth (backwards compatible). - let response: OAuthUsageResponse = serde_json::from_str( + let credentials = creds(None); + let usage = snapshot( r#"{ "five_hour": {"utilization": 10.0, "resets_at": "2026-08-13T12:49:59Z"} }"#, - ) - .expect("legacy-only body"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: None, - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + &credentials, + ); assert!((usage.primary.used_percent - 10.0).abs() < f64::EPSILON); assert_eq!(usage.primary.window_minutes, Some(300)); } #[test] -fn parses_retry_after_seconds() { - let header = HeaderValue::from_static("17"); - let duration = ClaudeOAuthFetcher::retry_after_duration(Some(&header)); - - assert_eq!(duration, Duration::from_secs(17)); -} - -#[test] -fn invalid_retry_after_uses_default_backoff() { - let header = HeaderValue::from_static("not-a-date"); - let duration = ClaudeOAuthFetcher::retry_after_duration(Some(&header)); - - assert_eq!(duration, ClaudeOAuthFetcher::DEFAULT_RATE_LIMIT_BACKOFF); +fn retry_after_parses_seconds_and_falls_back_to_default_backoff() { + for (value, expected) in [ + ("17", Duration::from_secs(17)), + ("not-a-date", ClaudeOAuthFetcher::DEFAULT_RATE_LIMIT_BACKOFF), + ] { + let header = HeaderValue::from_static(value); + let duration = ClaudeOAuthFetcher::retry_after_duration(Some(&header)); + + assert_eq!(duration, expected, "{value}"); + } } #[test] @@ -384,66 +333,39 @@ fn only_the_rate_limit_refusal_counts_as_rate_limited() { #[test] fn oauth_extras_put_scoped_weekly_before_routines() { - let response: OAuthUsageResponse = serde_json::from_str( + // With and without a scoped reset time. + for json in [ r#"{ - "five_hour": {"utilization": 10.0}, - "seven_day_routines": {"utilization": 5.0}, - "limits": [{ - "kind": "weekly_scoped", - "group": "weekly", - "percent": 7, - "resets_at": "2026-05-29T10:00:00Z", - "scope": {"model": {"display_name": "Fable"}} - }] - }"#, - ) - .expect("oauth body"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: None, - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); - - let ids: Vec<&str> = usage - .extra_rate_windows - .iter() - .map(|w| w.id.as_str()) - .collect(); - assert_eq!(ids, vec!["claude-weekly-scoped-fable", "claude-routines"]); -} - -#[test] -fn oauth_extras_keep_routines_in_raw_snapshot() { - let response: OAuthUsageResponse = serde_json::from_str( + "five_hour": {"utilization": 10.0}, + "seven_day_routines": {"utilization": 5.0}, + "limits": [{ + "kind": "weekly_scoped", + "group": "weekly", + "percent": 7, + "resets_at": "2026-05-29T10:00:00Z", + "scope": {"model": {"display_name": "Fable"}} + }] + }"#, r#"{ - "five_hour": {"utilization": 10.0}, - "seven_day_routines": {"utilization": 5.0}, - "limits": [{ - "kind": "weekly_scoped", - "group": "weekly", - "percent": 7, - "scope": {"model": {"display_name": "Fable"}} - }] - }"#, - ) - .expect("oauth body"); - - let credentials = ClaudeOAuthCredentials { - access_token: "token".to_string(), - refresh_token: None, - expires_at: None, - scopes: vec![], - rate_limit_tier: None, - }; - let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + "five_hour": {"utilization": 10.0}, + "seven_day_routines": {"utilization": 5.0}, + "limits": [{ + "kind": "weekly_scoped", + "group": "weekly", + "percent": 7, + "scope": {"model": {"display_name": "Fable"}} + }] + }"#, + ] { + let usage = snapshot(json, &creds(None)); - assert_eq!(usage.extra_rate_windows.len(), 2); - assert_eq!(usage.extra_rate_windows[0].id, "claude-weekly-scoped-fable"); - assert_eq!(usage.extra_rate_windows[1].id, "claude-routines"); + let ids: Vec<&str> = usage + .extra_rate_windows + .iter() + .map(|w| w.id.as_str()) + .collect(); + assert_eq!(ids, vec!["claude-weekly-scoped-fable", "claude-routines"]); + } } // ── Refresh-token backoff (upstream 0.48.0 #2650 mapping) ─── diff --git a/rust/src/providers/claude/quota_history.rs b/rust/src/providers/claude/quota_history.rs index 08e54c348f..962d677c49 100644 --- a/rust/src/providers/claude/quota_history.rs +++ b/rust/src/providers/claude/quota_history.rs @@ -158,24 +158,21 @@ pub fn aggregate_claude_quota_windows( history_coverage_established, } = options; let account_scope = account_scope.into(); + let empty = |account_scope| ClaudeQuotaHistoryReport { + account_scope, + windows: Vec::new(), + history_coverage_established, + }; let count = max_windows.clamp(1, 8); let duration_minutes = normalized_window_minutes(window_minutes); let duration = Duration::minutes(duration_minutes); let Some(live_reset_at) = live_reset_at.filter(|reset| reset.timestamp_millis() > 0) else { - return ClaudeQuotaHistoryReport { - account_scope, - windows: Vec::new(), - history_coverage_established, - }; + return empty(account_scope); }; let evidence = ResetEvidence::new(&account_scope, observations, now); if evidence.is_cancelled(live_reset_at) { - return ClaudeQuotaHistoryReport { - account_scope, - windows: Vec::new(), - history_coverage_established, - }; + return empty(account_scope); } let current_end = current_window_end(live_reset_at, now, duration); let boundaries = quota_boundaries(current_end, duration, &evidence, count); @@ -188,11 +185,7 @@ pub fn aggregate_claude_quota_windows( // transcript rows. Keep the account-scoped surface absent until a source // supplies explicit attribution for at least one row. if attributed_records.is_empty() { - return ClaudeQuotaHistoryReport { - account_scope, - windows: Vec::new(), - history_coverage_established, - }; + return empty(account_scope); } let deduped = deduplicate_claude_records(attributed_records); @@ -435,25 +428,36 @@ mod tests { value.parse().expect("RFC3339 timestamp") } + /// A row whose tokens and cost are complete exactly when present. fn record( at: &str, key: Option, tokens: Option, cost: Option, - tokens_are_complete: bool, - cost_is_complete: bool, ) -> ClaudeQuotaHistoryRecord { ClaudeQuotaHistoryRecord { timestamp: ts(at), tokens, cost_usd: cost, - tokens_are_complete, - cost_is_complete, + tokens_are_complete: tokens.is_some(), + cost_is_complete: cost.is_some(), dedup_key: key, attribution: ClaudeHistoryAttribution::Account("account-a".into()), } } + /// One closed weekly window ending at the 2026-09-22 reset. + fn options() -> ClaudeQuotaHistoryOptions<'static> { + ClaudeQuotaHistoryOptions { + live_reset_at: Some(ts("2026-09-22T12:00:00Z")), + window_minutes: None, + observations: &[], + now: ts("2026-09-21T12:00:00Z"), + max_windows: 1, + history_coverage_established: true, + } + } + #[test] fn reset_observations_are_account_scoped_and_rollovers_are_exact() { let now = ts("2026-09-21T12:00:00Z"); @@ -477,21 +481,14 @@ mod tests { ]; let report = aggregate_claude_quota_windows( "account-a", - &[record( - "2026-09-20T10:00:00Z", - None, - Some(1), - Some(0.1), - true, - true, - )], + &[record("2026-09-20T10:00:00Z", None, Some(1), Some(0.1))], ClaudeQuotaHistoryOptions { live_reset_at: Some(live_reset), window_minutes: Some(10_080), observations: &observations, now, max_windows: 3, - history_coverage_established: true, + ..options() }, ); assert_eq!(report.windows.len(), 3); @@ -506,22 +503,8 @@ mod tests { request_id: "req-1".into(), }; let records = deduplicate_claude_records(vec![ - record( - "2026-09-20T10:00:00Z", - Some(key.clone()), - Some(10), - None, - true, - false, - ), - record( - "2026-09-20T10:00:00Z", - Some(key), - Some(20), - Some(0.4), - true, - true, - ), + record("2026-09-20T10:00:00Z", Some(key.clone()), Some(10), None), + record("2026-09-20T10:00:00Z", Some(key), Some(20), Some(0.4)), record( "2026-09-20T10:01:00Z", Some(ClaudeQuotaDedupKey::Session { @@ -530,8 +513,6 @@ mod tests { }), Some(3), Some(0.1), - true, - true, ), ]); assert_eq!(records.len(), 2); @@ -544,17 +525,10 @@ mod tests { let report = aggregate_claude_quota_windows( "account-a", &[ - record("2026-09-20T10:00:00Z", None, Some(10), None, true, false), - record("2026-09-20T11:00:00Z", None, None, Some(0.5), false, true), + record("2026-09-20T10:00:00Z", None, Some(10), None), + record("2026-09-20T11:00:00Z", None, None, Some(0.5)), ], - ClaudeQuotaHistoryOptions { - live_reset_at: Some(ts("2026-09-22T12:00:00Z")), - window_minutes: None, - observations: &[], - now: ts("2026-09-21T12:00:00Z"), - max_windows: 1, - history_coverage_established: true, - }, + options(), ); let window = &report.windows[0]; assert_eq!(window.total_tokens, Some(10)); @@ -568,24 +542,10 @@ mod tests { let report = aggregate_claude_quota_windows( "account-a", &[ - record( - "2026-09-20T10:00:00Z", - None, - Some(u64::MAX), - None, - true, - false, - ), - record("2026-09-20T11:00:00Z", None, Some(1), None, true, false), + record("2026-09-20T10:00:00Z", None, Some(u64::MAX), None), + record("2026-09-20T11:00:00Z", None, Some(1), None), ], - ClaudeQuotaHistoryOptions { - live_reset_at: Some(ts("2026-09-22T12:00:00Z")), - window_minutes: None, - observations: &[], - now: ts("2026-09-21T12:00:00Z"), - max_windows: 1, - history_coverage_established: true, - }, + options(), ); let window = &report.windows[0]; @@ -602,10 +562,7 @@ mod tests { ClaudeQuotaHistoryOptions { live_reset_at: Some(before), window_minutes: Some(10_080), - observations: &[], - now: ts("2026-09-21T12:00:00Z"), - max_windows: 1, - history_coverage_established: true, + ..options() }, ); assert_eq!(before, ts("2026-09-22T12:00:00Z")); @@ -614,14 +571,7 @@ mod tests { #[test] fn account_scoped_history_excludes_unavailable_and_mismatched_rows() { - let mut unavailable = record( - "2026-09-20T10:00:00Z", - None, - Some(100), - Some(1.0), - true, - true, - ); + let mut unavailable = record("2026-09-20T10:00:00Z", None, Some(100), Some(1.0)); unavailable.attribution = ClaudeHistoryAttribution::Unavailable; let mut other_account = unavailable.clone(); other_account.attribution = ClaudeHistoryAttribution::Account("account-b".into()); @@ -630,14 +580,7 @@ mod tests { let report = aggregate_claude_quota_windows( "account-a", &[unavailable, other_account, matching], - ClaudeQuotaHistoryOptions { - live_reset_at: Some(ts("2026-09-22T12:00:00Z")), - window_minutes: None, - observations: &[], - now: ts("2026-09-21T12:00:00Z"), - max_windows: 1, - history_coverage_established: true, - }, + options(), ); let window = &report.windows[0]; assert_eq!(window.entry_count, 1); @@ -647,27 +590,9 @@ mod tests { #[test] fn account_scoped_history_is_absent_without_explicit_attribution() { - let mut record = record( - "2026-09-20T10:00:00Z", - None, - Some(100), - Some(1.0), - true, - true, - ); + let mut record = record("2026-09-20T10:00:00Z", None, Some(100), Some(1.0)); record.attribution = ClaudeHistoryAttribution::Unavailable; - let report = aggregate_claude_quota_windows( - "account-a", - &[record], - ClaudeQuotaHistoryOptions { - live_reset_at: Some(ts("2026-09-22T12:00:00Z")), - window_minutes: None, - observations: &[], - now: ts("2026-09-21T12:00:00Z"), - max_windows: 1, - history_coverage_established: true, - }, - ); + let report = aggregate_claude_quota_windows("account-a", &[record], options()); assert!(report.windows.is_empty()); } diff --git a/rust/src/providers/claude/reset_observations.rs b/rust/src/providers/claude/reset_observations.rs index 6ac5fd007d..762ec069da 100644 --- a/rust/src/providers/claude/reset_observations.rs +++ b/rust/src/providers/claude/reset_observations.rs @@ -51,16 +51,6 @@ pub struct ClaudeResetObservationMergeResult { pub changed: bool, } -pub fn default_store_path() -> Result { - let root = dirs::config_dir().ok_or_else(|| { - ClaudeResetObservationError::Read(std::io::Error::new( - std::io::ErrorKind::NotFound, - "configuration directory not found", - )) - })?; - Ok(root.join("CodexBar").join(STORE_RELATIVE_PATH)) -} - pub fn store_path(config_root: &Path) -> PathBuf { config_root.join(STORE_RELATIVE_PATH) } @@ -70,20 +60,8 @@ pub fn load_reset_observations( account_scope: &str, ) -> Result, ClaudeResetObservationError> { validate_scope(account_scope)?; - let path = store_path(config_root); - let raw = match secure_file::read_string(&path) { - Ok(raw) => raw, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), - Err(error) => return Err(ClaudeResetObservationError::Read(error)), - }; - let store: ClaudeResetObservationStore = - serde_json::from_str(&raw).map_err(ClaudeResetObservationError::Deserialize)?; - validate_store(&store)?; - Ok(store - .accounts - .get(account_scope) - .cloned() - .unwrap_or_default()) + let mut store = load_store(&store_path(config_root))?; + Ok(store.accounts.remove(account_scope).unwrap_or_default()) } pub fn merge_reset_observations( @@ -230,6 +208,31 @@ mod tests { ); } + #[test] + fn missing_store_loads_empty_and_invalid_scopes_fail() { + let root = tempdir().unwrap(); + assert_eq!(load_reset_observations(root.path(), "a").unwrap(), vec![]); + assert!(matches!( + load_reset_observations(root.path(), " "), + Err(ClaudeResetObservationError::EmptyAccountScope) + )); + let path = store_path(root.path()); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, br#"{"version":1}"#).unwrap(); + assert_eq!(load_reset_observations(root.path(), "a").unwrap(), vec![]); + let stray = observation("b", "2026-09-20T10:00:00Z", "2026-09-21T10:00:00Z"); + let store = serde_json::json!({"version": 1, "accounts": {"a": [stray]}}); + std::fs::write(&path, store.to_string()).unwrap(); + assert!(matches!( + load_reset_observations(root.path(), "z"), + Err(ClaudeResetObservationError::AccountScopeMismatch) + )); + assert!(matches!( + merge_and_persist_reset_observations(root.path(), "z", &[]), + Err(ClaudeResetObservationError::AccountScopeMismatch) + )); + } + #[test] fn malformed_and_future_stores_fail_closed() { let root = tempdir().unwrap(); diff --git a/rust/src/providers/claude/reset_opt_in_tests.rs b/rust/src/providers/claude/reset_opt_in_tests.rs index de0434b824..0b29f654bf 100644 --- a/rust/src/providers/claude/reset_opt_in_tests.rs +++ b/rust/src/providers/claude/reset_opt_in_tests.rs @@ -1,7 +1,9 @@ -//! Wire tests for the `cedar_ember` opt-in on the Claude Web usage request. +//! Wire tests for the Claude Web requests: the `cedar_ember` usage opt-in +//! and the status and parse mapping of each GET. use super::ClaudeWebApiFetcher; use crate::core::ProviderError; +use crate::providers::test_support::{mock_response, mock_response_expect}; use mockito::{Matcher, Mock, Server, ServerGuard}; const USAGE_PATH: &str = "/organizations/org-123/usage"; @@ -141,3 +143,184 @@ async fn ordinary_forbidden_retries_and_a_second_forbidden_is_an_auth_failure() first.assert_async().await; retry.assert_async().await; } + +fn describe(result: Result) -> String { + match result { + Ok(value) => format!("ok {value}"), + Err(ProviderError::AuthRequired) => "auth".to_string(), + Err(ProviderError::Parse(message)) => format!("parse {message}"), + Err(ProviderError::Other(message)) => format!("other {message}"), + Err(other) => format!("unexpected {other:?}"), + } +} + +#[tokio::test] +async fn organization_lookup_maps_statuses_after_the_account_fallback() { + let cloudflare = format!("other {CLOUDFLARE_MESSAGE}"); + let rows: [(usize, &str, &str); 6] = [ + (200, r#"[{"uuid": "org-9", "name": "Team"}]"#, "ok org-9"), + (200, "[]", "parse No organizations found"), + (401, "{}", "auth"), + (403, "permission denied", "auth"), + (403, "Just a moment...", &cloudflare), + ( + 500, + "{}", + "other Failed to get organizations: 500 Internal Server Error", + ), + ]; + for (status, body, expected) in rows { + let mut server = Server::new_async().await; + let account = mock_response(&mut server, "GET", "/account", 500, "{}").await; + let orgs = mock_response(&mut server, "GET", "/organizations", status, body).await; + let fetcher = ClaudeWebApiFetcher::new().with_base_url(server.url()); + let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-fixture-token"); + + let got = describe( + fetcher + .get_organization_id("sessionKey=sk-ant-fixture-token", &headers) + .await, + ); + + assert_eq!(got, expected, "status {status}"); + account.assert_async().await; + orgs.assert_async().await; + } +} + +#[tokio::test] +async fn organization_lookup_prefers_the_cookie_then_the_account() { + let mut server = Server::new_async().await; + let account = mock_response_expect( + &mut server, + "GET", + "/account", + 200, + r#"{"memberships": [{"organization": {"uuid": " org-acct "}}]}"#, + 1, + ) + .await; + let orgs = mock_response_expect(&mut server, "GET", "/organizations", 200, "[]", 0).await; + let fetcher = ClaudeWebApiFetcher::new().with_base_url(server.url()); + let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-fixture-token"); + + let from_cookie = fetcher + .get_organization_id( + "sessionKey=sk-ant-fixture-token; lastActiveOrg=org-cookie", + &headers, + ) + .await; + let from_account = fetcher + .get_organization_id("sessionKey=sk-ant-fixture-token", &headers) + .await; + + assert_eq!(describe(from_cookie), "ok org-cookie"); + assert_eq!(describe(from_account), "ok org-acct"); + account.assert_async().await; + orgs.assert_async().await; +} + +#[tokio::test] +async fn extra_usage_and_account_failures_keep_their_plain_status_text() { + let rows: [(usize, &str, &str, &str); 4] = [ + ( + 401, + "{}", + "other Failed to get extra usage: 401 Unauthorized", + "other Failed to get account: 401 Unauthorized", + ), + ( + 403, + "Just a moment...", + "other Failed to get extra usage: 403 Forbidden", + "other Failed to get account: 403 Forbidden", + ), + ( + 500, + "{}", + "other Failed to get extra usage: 500 Internal Server Error", + "other Failed to get account: 500 Internal Server Error", + ), + ( + 200, + "not json", + "parse Failed to parse extra usage: ", + "parse Failed to parse account: ", + ), + ]; + for (status, body, extra_expected, account_expected) in rows { + let mut server = Server::new_async().await; + let extra = mock_response( + &mut server, + "GET", + "/organizations/org-123/overage_spend_limit", + status, + body, + ) + .await; + let account = mock_response(&mut server, "GET", "/account", status, body).await; + let fetcher = ClaudeWebApiFetcher::new().with_base_url(server.url()); + let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-fixture-token"); + + let extra_got = describe( + fetcher + .get_extra_usage("org-123", &headers) + .await + .map(|usage| format!("{:?}", usage.monthly_credit_limit)), + ); + let account_got = describe( + fetcher + .get_account_info(&headers) + .await + .map(|info| format!("{:?}", info.email_address)), + ); + + if status == 200 { + assert!(extra_got.starts_with(extra_expected), "{extra_got}"); + assert!(account_got.starts_with(account_expected), "{account_got}"); + } else { + assert_eq!(extra_got, extra_expected, "status {status}"); + assert_eq!(account_got, account_expected, "status {status}"); + } + extra.assert_async().await; + account.assert_async().await; + } +} + +#[tokio::test] +async fn extra_usage_and_account_parse_successful_bodies() { + let mut server = Server::new_async().await; + let extra = mock_response( + &mut server, + "GET", + "/organizations/org-123/overage_spend_limit", + 200, + r#"{"monthly_credit_limit": 5000, "used_credits": 1200, "currency": "USD", "is_enabled": true}"#, + ) + .await; + let account = mock_response( + &mut server, + "GET", + "/account", + 200, + r#"{"email_address": "a@example.com", "rate_limit_tier": "default_claude_max_5x"}"#, + ) + .await; + let fetcher = ClaudeWebApiFetcher::new().with_base_url(server.url()); + let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-fixture-token"); + + let usage = fetcher.get_extra_usage("org-123", &headers).await.unwrap(); + let info = fetcher.get_account_info(&headers).await.unwrap(); + + assert_eq!(usage.monthly_credit_limit, Some(5000.0)); + assert_eq!(usage.used_credits, Some(1200.0)); + assert_eq!(usage.currency.as_deref(), Some("USD")); + assert_eq!(usage.is_enabled, Some(true)); + assert_eq!(info.email_address.as_deref(), Some("a@example.com")); + assert_eq!( + info.rate_limit_tier.as_deref(), + Some("default_claude_max_5x") + ); + extra.assert_async().await; + account.assert_async().await; +} diff --git a/rust/src/providers/claude/scoped_weekly.rs b/rust/src/providers/claude/scoped_weekly.rs index 35da92c032..601093ff2e 100644 --- a/rust/src/providers/claude/scoped_weekly.rs +++ b/rust/src/providers/claude/scoped_weekly.rs @@ -77,14 +77,7 @@ pub(super) fn weekly_all_window(limits: &[ScopedWeeklyLimit]) -> Option Option if limit.group.as_deref().is_some_and(|g| g != "session") { return None; } - let percent = limit.percent.filter(|value| value.is_finite())?; - let resets_at = limit_resets_at(limit); - Some(RateWindow::with_details( - percent.clamp(0.0, 100.0), - Some(5 * 60), - resets_at, - None, - )) + limit_window(limit, 5 * 60) }) } +/// A finite `percent`, clamped, as a window of `minutes`. +fn limit_window(limit: &ScopedWeeklyLimit, minutes: u32) -> Option { + let percent = limit.percent.filter(|value| value.is_finite())?; + Some(RateWindow::with_details( + percent.clamp(0.0, 100.0), + Some(minutes), + limit_resets_at(limit), + None, + )) +} + fn limit_resets_at(limit: &ScopedWeeklyLimit) -> Option> { limit .resets_at diff --git a/rust/src/providers/claude/tests.rs b/rust/src/providers/claude/tests.rs new file mode 100644 index 0000000000..0deb0f407e --- /dev/null +++ b/rust/src/providers/claude/tests.rs @@ -0,0 +1,1164 @@ +use chrono::{DateTime, Utc}; +use std::collections::HashMap; + +use super::cli_probe::{ + CLAUDE_PROBE_CACHE_FILE, CLAUDE_PROBE_CACHE_TTL, ClaudeProbeCache, ClaudeProbeLock, + claude_passive_probe_env, claude_probe_launch_args, claude_project_dir_name, + claude_usage_settings_args, cleanup_probe_session_jsonl, cleanup_probe_transcripts_in, + javascript_hash_base36, load_cached_probe_output, load_or_create_probe_session_id, + login_fingerprint_at, run_locked_probe, store_cached_probe_output, +}; +use super::*; + +const LOGIN_A: &str = "login-a"; + +fn login_a() -> Option { + Some(LOGIN_A.to_string()) +} + +#[test] +fn probe_cache_roundtrip_and_expiry() { + let dir = tempfile::tempdir().unwrap(); + assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); + store_cached_probe_output(dir.path(), LOGIN_A, "Current session 12% used"); + assert_eq!( + load_cached_probe_output(dir.path(), LOGIN_A).as_deref(), + Some("Current session 12% used") + ); + let stale = ClaudeProbeCache { + captured_at_unix: unix_now_secs() - CLAUDE_PROBE_CACHE_TTL.as_secs() - 5, + login: LOGIN_A.to_string(), + output: "Current session 12% used".to_string(), + }; + std::fs::write( + dir.path().join(CLAUDE_PROBE_CACHE_FILE), + serde_json::to_string(&stale).unwrap(), + ) + .unwrap(); + assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); +} + +#[test] +fn probe_cache_is_never_shared_with_another_login() { + let dir = tempfile::tempdir().unwrap(); + store_cached_probe_output(dir.path(), LOGIN_A, "Current session 12% used"); + assert!(load_cached_probe_output(dir.path(), "login-b").is_none()); + + // Written before screens were scoped to a login. + let unscoped = format!( + r#"{{"captured_at_unix":{},"output":"Current session 12% used"}}"#, + unix_now_secs() + ); + std::fs::write(dir.path().join(CLAUDE_PROBE_CACHE_FILE), unscoped).unwrap(); + assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); + assert!(load_cached_probe_output(dir.path(), "").is_none()); +} + +#[test] +fn login_fingerprint_follows_credential_rewrites_without_reading_them() { + let dir = tempfile::tempdir().unwrap(); + let credentials = dir.path().join(".credentials.json"); + assert_eq!(login_fingerprint_at(&credentials), None); + + std::fs::write(&credentials, "{}").unwrap(); + let first = login_fingerprint_at(&credentials).expect("fingerprint"); + assert_eq!(login_fingerprint_at(&credentials).as_ref(), Some(&first)); + assert!(!first.contains(".credentials"), "only a digest is stored"); + + std::fs::write(&credentials, r#"{"another":"login"}"#).unwrap(); + assert_ne!(login_fingerprint_at(&credentials), Some(first)); +} + +const SHAREABLE_USAGE_SCREEN: &str = "Current session\n\ + ████████▌ 17% used\n\ + Resets 12pm (America/Bogota)\n"; + +#[test] +fn locked_probe_reuses_a_screen_stored_while_it_waited() { + let dir = tempfile::tempdir().unwrap(); + store_cached_probe_output(dir.path(), LOGIN_A, SHAREABLE_USAGE_SCREEN); + + let output = run_locked_probe(dir.path(), Some(&login_a), || { + panic!("a fresh shared screen must not launch another probe") + }) + .unwrap(); + assert_eq!(output, SHAREABLE_USAGE_SCREEN); +} + +#[test] +fn locked_probe_shares_only_parseable_usage_screens() { + let dir = tempfile::tempdir().unwrap(); + let output = run_locked_probe(dir.path(), Some(&login_a), || { + Ok("Not logged in".to_string()) + }); + assert_eq!(output.unwrap(), "Not logged in"); + assert!(load_cached_probe_output(dir.path(), LOGIN_A).is_none()); + + let output = run_locked_probe(dir.path(), Some(&login_a), || { + Ok(SHAREABLE_USAGE_SCREEN.into()) + }); + assert_eq!(output.unwrap(), SHAREABLE_USAGE_SCREEN); + assert_eq!( + load_cached_probe_output(dir.path(), LOGIN_A).as_deref(), + Some(SHAREABLE_USAGE_SCREEN) + ); + assert!( + !dir.path() + .read_dir() + .unwrap() + .flatten() + .any(|entry| entry.file_name().to_string_lossy().contains(".tmp-")), + "the atomic write left no staging file behind" + ); +} + +#[test] +fn locked_probe_keeps_a_screen_private_when_the_login_changed_meanwhile() { + let dir = tempfile::tempdir().unwrap(); + let calls = std::cell::Cell::new(0); + let switching_login = || { + calls.set(calls.get() + 1); + Some(format!("login-{}", calls.get())) + }; + let output = run_locked_probe(dir.path(), Some(&switching_login), || { + Ok(SHAREABLE_USAGE_SCREEN.into()) + }); + assert_eq!(output.unwrap(), SHAREABLE_USAGE_SCREEN); + assert_eq!(calls.get(), 2, "the login is read before and after"); + assert!(load_cached_probe_output(dir.path(), "login-1").is_none()); + assert!(load_cached_probe_output(dir.path(), "login-2").is_none()); + + let no_login = || None; + run_locked_probe(dir.path(), Some(&no_login), || { + Ok(SHAREABLE_USAGE_SCREEN.into()) + }) + .unwrap(); + assert!(!dir.path().join(CLAUDE_PROBE_CACHE_FILE).exists()); +} + +#[test] +fn unshared_probe_neither_reuses_nor_stores_screens() { + let dir = tempfile::tempdir().unwrap(); + store_cached_probe_output(dir.path(), LOGIN_A, SHAREABLE_USAGE_SCREEN); + let output = run_locked_probe(dir.path(), None, || Ok("trust preflight".into())); + assert_eq!(output.unwrap(), "trust preflight"); + + let other = tempfile::tempdir().unwrap(); + run_locked_probe(other.path(), None, || Ok(SHAREABLE_USAGE_SCREEN.into())).unwrap(); + assert!(!other.path().join(CLAUDE_PROBE_CACHE_FILE).exists()); +} + +#[test] +fn logged_probe_screen_masks_account_email_and_secrets() { + let screen = "Login: someone@example.com (Claude Max)\n\ + access_token=abcdef0123456789 sk-ant-abcdefgh12345678\n\ + Current session 12% used"; + let logged = redacted_probe_screen(screen); + assert!(!logged.contains("someone@example.com"), "{logged}"); + assert!(!logged.contains("abcdef0123456789"), "{logged}"); + assert!(!logged.contains("sk-ant-abcdefgh12345678"), "{logged}"); + assert!(logged.contains("Current session 12% used")); +} + +#[test] +fn probe_lock_wait_expiry_fails_instead_of_probing_alongside() { + let dir = tempfile::tempdir().unwrap(); + let held = ClaudeProbeLock::acquire_within(dir.path(), Duration::ZERO) + .expect("first lock") + .expect("file locking is supported"); + + let error = match ClaudeProbeLock::acquire_within(dir.path(), Duration::from_millis(300)) { + Ok(lock) => panic!("second lock acquired while held: {}", lock.is_some()), + Err(error) => error, + }; + assert!(error.to_string().contains("Timed out waiting"), "{error}"); + assert_eq!( + last_good_failure_policy_for_error(&error.to_string()), + LastGoodFailurePolicy::Preserve + ); + + drop(held); + assert!( + ClaudeProbeLock::acquire_within(dir.path(), Duration::ZERO) + .expect("lock after release") + .is_some() + ); +} + +#[test] +fn passive_probe_env_disables_autoupdater_color_and_chrome() { + let env = claude_passive_probe_env(HashMap::new()); + for (key, value) in [ + ("DISABLE_AUTOUPDATER", "1"), + ("NO_COLOR", "1"), + ("CLAUDE_CODE_ENABLE_CFC", "0"), + ] { + assert_eq!(env.get(key).map(String::as_str), Some(value), "{key}"); + } +} + +#[test] +fn probe_session_id_is_reused_from_probe_directory() { + let dir = tempfile::tempdir().unwrap(); + let first = load_or_create_probe_session_id(dir.path()); + let second = load_or_create_probe_session_id(dir.path()); + assert_eq!(first, second); + assert!(uuid::Uuid::parse_str(&first).is_ok()); + let args = claude_probe_launch_args(&first); + // Positional structure only: the settings pair is pinned once by + // `claude_usage_settings_args` being the sole composer. + assert_eq!( + args[..4], + ["--setting-sources", "user", "--allowed-tools", ""] + ); + assert_eq!(args[4], claude_usage_settings_args()[0]); + assert_eq!(args[5], claude_usage_settings_args()[1]); + assert_eq!(args[6], "--session-id"); + assert_eq!(args[7], first); +} + +#[test] +fn usage_probe_settings_disable_remote_control_startup() { + assert_eq!( + claude_usage_settings_args(), + [ + "--settings".to_string(), + r#"{"remoteControlAtStartup":false,"tui":"default"}"#.to_string(), + ] + ); + let settings: serde_json::Value = + serde_json::from_str(&claude_usage_settings_args()[1]).unwrap(); + assert_eq!(settings["tui"], "default"); + assert_eq!(settings["remoteControlAtStartup"], false); +} + +#[test] +fn probe_session_jsonl_cleanup_removes_transcript_files() { + let dir = tempfile::tempdir().unwrap(); + let jsonl = dir.path().join("session.jsonl"); + std::fs::write(&jsonl, "{}").unwrap(); + std::fs::write(dir.path().join("keep.txt"), "x").unwrap(); + cleanup_probe_session_jsonl(dir.path()); + assert!(!jsonl.exists()); + assert!(dir.path().join("keep.txt").exists()); +} + +#[test] +fn probe_project_dir_name_matches_claude_code() { + use std::path::Path; + assert_eq!( + claude_project_dir_name(Path::new( + r"C:\Users\user\AppData\Local\CodexBar\claude-usage-probe" + )), + "C--Users-user-AppData-Local-CodexBar-claude-usage-probe" + ); + assert_eq!( + claude_project_dir_name(Path::new("/Users/me/Library/Application Support/x")), + "-Users-me-Library-Application-Support-x" + ); + // One dash per UTF-16 code unit, so two for a character outside the BMP. + assert_eq!( + claude_project_dir_name(Path::new("C:\\Users\\J\u{f6}rg\u{1F600}\\probe")), + "C--Users-J-rg---probe" + ); + // Reference values from Claude Code's JavaScript implementation. + let long = format!( + r"C:\Users\user\AppData\Local\{}claude-usage-probe", + r"deep\".repeat(40) + ); + assert_eq!( + claude_project_dir_name(Path::new(&long)), + format!( + "C--Users-user-AppData-Local-{}de-ttzy4x", + "deep-".repeat(34) + ) + ); + assert_eq!(javascript_hash_base36("hello"), "1n1e4y"); + assert_eq!(javascript_hash_base36(""), "0"); +} + +#[test] +fn probe_transcript_cleanup_stays_inside_the_probe_project() { + use std::path::Path; + let projects = tempfile::tempdir().unwrap(); + let other = projects.path().join("C--work-repo"); + std::fs::create_dir_all(&other).unwrap(); + std::fs::write(other.join("session.jsonl"), "{}").unwrap(); + + let busy_probe = Path::new(r"C:\Users\user\AppData\Local\CodexBar\busy-probe"); + let busy = projects.path().join(claude_project_dir_name(busy_probe)); + std::fs::create_dir_all(busy.join("folder.jsonl")).unwrap(); + std::fs::write(busy.join("session.jsonl"), "{}").unwrap(); + std::fs::write(busy.join("notes.txt"), "x").unwrap(); + cleanup_probe_transcripts_in(projects.path(), busy_probe); + assert!(!busy.join("session.jsonl").exists()); + assert!(busy.join("notes.txt").exists()); + assert!(busy.join("folder.jsonl").is_dir(), "only files are removed"); + + let probe = Path::new(r"C:\Users\user\AppData\Local\CodexBar\claude-usage-probe"); + let project = projects.path().join(claude_project_dir_name(probe)); + std::fs::create_dir_all(&project).unwrap(); + std::fs::write(project.join("a.jsonl"), "{}").unwrap(); + std::fs::write(project.join("b.jsonl"), "{}").unwrap(); + cleanup_probe_transcripts_in(projects.path(), probe); + assert!(!project.exists(), "an emptied probe project dir is removed"); + + assert!(other.join("session.jsonl").exists(), "other projects stay"); +} + +fn parse_ok(output: &str) -> ProviderFetchResult { + ClaudeProvider::new() + .parse_cli_output(output) + .expect("should parse") +} + +fn parse_err(output: &str) -> ProviderError { + ClaudeProvider::new() + .parse_cli_output(output) + .expect_err("should reject") +} + +#[test] +fn parses_current_cli_usage_screen() { + let output = r#" +Status Config Usage + + Current session + ██████████████████████████████████████████████████ 100% used + Resets 12pm (America/Bogota) + + Current week (all models) + ████████████████████████▌ 49% used + Resets Apr 3, 2pm (America/Bogota) + + Extra usage + ██▍ 4% used + $3.31 / $70.00 spent · Resets Apr 1 (America/Bogota) +"#; + + let result = parse_ok(output); + + assert_eq!(result.source_label, "cli"); + assert_eq!(result.usage.primary.used_percent, 100.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Resets 12pm (America/Bogota)") + ); + + let weekly = result + .usage + .secondary + .expect("weekly usage should be present"); + assert_eq!(weekly.used_percent, 49.0); + assert_eq!( + weekly.reset_description.as_deref(), + Some("Resets Apr 3, 2pm (America/Bogota)") + ); +} + +#[test] +fn parses_short_forms_as_full_session_usage() { + let rows = [ + ( + "You're out of extra usage · resets 12pm (America/Bogota)", + "resets 12pm (America/Bogota)", + ), + ( + "You've hit your limit \u{00b7} resets 3:20pm (Asia/Shanghai)", + "resets 3:20pm (Asia/Shanghai)", + ), + ]; + for (output, reset) in rows { + let result = parse_ok(output); + assert_eq!(result.usage.primary.used_percent, 100.0, "{output}"); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some(reset) + ); + } +} + +#[test] +fn parses_remaining_available_and_decimal_percentages() { + let output = r#" +Status Config Usage + + Current session + 12.5% remaining + Resets 8pm + + Current week (all models) + 4% available + Resets Apr 4, 2pm + + Current week (Sonnet only) + 1% consumed +"#; + + let result = parse_ok(output); + + assert_eq!(result.usage.primary.used_percent, 87.5); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Resets 8pm") + ); + + let weekly = result + .usage + .secondary + .expect("weekly usage should be present"); + assert_eq!(weekly.used_percent, 96.0); + assert_eq!( + weekly.reset_description.as_deref(), + Some("Resets Apr 4, 2pm") + ); + + let sonnet = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == "claude-weekly-scoped-sonnet") + .expect("sonnet usage should be present"); + assert_eq!(sonnet.window.used_percent, 1.0); +} + +#[test] +fn parses_all_cli_model_scoped_weekly_limits() { + let output = r#" +Current session +10% used +Resets 12pm (America/Bogota) + +Current week (all models) +20% used +Resets Apr 3, 2pm (America/Bogota) + +Current week (Sonnet only) +30% used +Resets Apr 4, 2pm (America/Bogota) + +Current week (Opus only) +40% used +Resets Apr 5, 2pm (America/Bogota) +"#; + + let result = parse_ok(output); + + assert_eq!(result.usage.extra_rate_windows.len(), 2); + assert_eq!( + result.usage.extra_rate_windows[0].id, + "claude-weekly-scoped-sonnet" + ); + assert_eq!(result.usage.extra_rate_windows[0].title, "Sonnet only"); + assert_eq!(result.usage.extra_rate_windows[0].window.used_percent, 30.0); + assert_eq!( + result.usage.extra_rate_windows[1].id, + "claude-weekly-scoped-opus" + ); + assert!(result.usage.model_specific.is_none()); +} + +#[test] +fn scoped_weekly_parser_handles_non_ascii_labels_and_reset_prefixes() { + let now = "2026-04-02T18:00:00Z".parse::>().unwrap(); + let limits = extract_cli_scoped_weekly_limits( + "Current week (A€€)\n10% used\nİResets Apr 3 at 2pm (America/Bogota)", + now, + ); + + assert_eq!(limits.len(), 1); + assert_eq!(limits[0].title, "A€€"); + assert_eq!( + limits[0].window.resets_at, + Some("2026-04-03T19:00:00Z".parse().unwrap()) + ); +} + +#[test] +fn resolves_cli_reset_occurrences_in_the_reported_timezone() { + let now = "2026-04-02T18:00:00Z".parse::>().unwrap(); + + assert_eq!( + parse_claude_reset_date("Resets Apr 3, 2027, 2pm (America/Bogota)", now, None), + Some("2027-04-03T19:00:00Z".parse().unwrap()) + ); + assert_eq!( + parse_claude_reset_date("Resets Apr 3, 2pm (America/Bogota)", now, None), + Some("2026-04-03T19:00:00Z".parse().unwrap()) + ); + assert_eq!( + parse_claude_reset_date("Resets 12pm (America/Bogota)", now, None), + Some("2026-04-03T17:00:00Z".parse().unwrap()) + ); + assert_eq!( + parse_claude_reset_date("ResetsApr3at2pm(America/Bogota)", now, None), + Some("2026-04-03T19:00:00Z".parse().unwrap()) + ); +} + +#[test] +fn timezone_less_resets_use_the_supplied_system_zone() { + let now = "2026-03-07T18:00:00Z".parse::>().unwrap(); + + assert_eq!( + parse_claude_reset_date_in_system_zone( + "Resets Mar 8 at 3:30am", + now, + None, + "America/New_York".parse().unwrap(), + ), + Some("2026-03-08T07:30:00Z".parse().unwrap()) + ); + assert_eq!( + parse_claude_reset_date_in_system_zone( + "Resets Mar 8 at 3:30am (America/Los_Angeles)", + now, + None, + "America/New_York".parse().unwrap(), + ), + Some("2026-03-08T10:30:00Z".parse().unwrap()) + ); +} + +#[test] +fn reset_dates_resolve_every_month_and_form() { + let now = "2026-09-24T12:00:00Z".parse::>().unwrap(); + let at = |text: &str, window: Option| { + parse_claude_reset_date(text, now, window).map(|date| date.to_rfc3339()) + }; + let months = [ + "Jan", "FEB", "mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "dEc", + ]; + for (index, month) in months.iter().enumerate() { + assert_eq!( + at(&format!("Resets {month} 5, 2027 at 3pm (UTC)"), None), + Some(format!("2027-{:02}-05T15:00:00+00:00", index + 1)), + "{month}" + ); + } + let rows = [ + ("Resets Foo 5, 2027 at 3pm (UTC)", None, None), + ("Resets Feb 30, 2027 at 3pm (UTC)", None, None), + ( + "Resets Sep 23 at 3pm (UTC)", + None, + Some("2027-09-23T15:00:00+00:00"), + ), + ( + "Resets Sep 23 at 3pm (UTC)", + Some(10_080), + Some("2026-09-23T15:00:00+00:00"), + ), + ( + "Resets Feb 29 at 3pm (UTC)", + None, + Some("2028-02-29T15:00:00+00:00"), + ), + ("Resets 3pm (UTC)", None, Some("2026-09-24T15:00:00+00:00")), + ("Resets 11am (UTC)", None, Some("2026-09-25T11:00:00+00:00")), + ( + "Resets 11am (UTC)", + Some(300), + Some("2026-09-24T11:00:00+00:00"), + ), + ( + "Resets Nov 1, 2026 at 1:30am (America/New_York)", + None, + Some("2026-11-01T05:30:00+00:00"), + ), + ( + "Resets Mar 8, 2026 at 2:30am (America/New_York)", + None, + None, + ), + ]; + for (text, window, expected) in rows { + assert_eq!(at(text, window).as_deref(), expected, "{text} {window:?}"); + } +} + +#[test] +fn parses_compact_usage_screen() { + let output = r#" +Settings:StatusConfigUsage(tabtocycle) +Loadingusagedata... +Currentsession +6%used +Resets4:29am(Asia/Calcutta) +Currentweek(allmodels) +4%used +ResetsFeb12at1:29pm(Asia/Calcutta) +Currentweek(Sonnetonly) +1%used +ResetsFeb12at1:29pm(Asia/Calcutta) +"#; + + let result = parse_ok(output); + + assert_eq!(result.usage.primary.used_percent, 6.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Resets4:29am(Asia/Calcutta)") + ); + assert_eq!( + result + .usage + .secondary + .expect("weekly usage should be present") + .used_percent, + 4.0 + ); + let sonnet = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == "claude-weekly-scoped-sonnet") + .expect("sonnet usage should be present"); + assert_eq!(result.usage.extra_rate_windows.len(), 1); + assert_eq!(sonnet.title, "Sonnet only"); + assert_eq!(sonnet.window.used_percent, 1.0); +} + +#[test] +fn does_not_promote_weekly_reset_to_session() { + let output = r#" +Current session +17% used +Current week (all models) +4% used +Resets Dec 24 at 3:59pm (Europe/Paris) +"#; + + let result = parse_ok(output); + + assert_eq!(result.usage.primary.used_percent, 17.0); + assert_eq!(result.usage.primary.reset_description, None); + assert_eq!( + result + .usage + .secondary + .expect("weekly usage should be present") + .reset_description + .as_deref(), + Some("Resets Dec 24 at 3:59pm (Europe/Paris)") + ); +} + +#[test] +fn cli_error_markers_map_to_fixed_errors() { + let git_bash = "Other(\"Claude CLI requires Git Bash on Windows. Install Git for Windows or set CLAUDE_CODE_GIT_BASH_PATH to your bash.exe path.\")"; + let cases = [ + ("Error: Not Logged In", "AuthRequired"), + ("login required to continue", "AuthRequired"), + ( + "TOKEN EXPIRED", + "OAuthExpired(\"Token expired. Run `claude login` to refresh.\")", + ), + ( + "{\"type\":\"token_expired\"}", + "OAuthExpired(\"Token expired. Run `claude login` to refresh.\")", + ), + ( + "authentication_error", + "OAuth(\"Authentication error. Run `claude login`.\")", + ), + ("Claude Code on Windows requires git-bash.", git_bash), + ( + "Running scripts is disabled on this system", + "Other(\"Claude CLI could not start because PowerShell script execution is disabled. Use claude.cmd or adjust the execution policy.\")", + ), + ( + "Cannot run a document in the middle of a pipeline", + "Other(\"Claude CLI resolved to a Unix shell script on Windows. Reinstall Claude Code or ensure claude.cmd is first on PATH.\")", + ), + // Auth markers win over environment markers. + ("requires git-bash; not logged in", "AuthRequired"), + ( + "requires git-bash; token expired", + "OAuthExpired(\"Token expired. Run `claude login` to refresh.\")", + ), + // Login wins over the other auth markers. + ("token expired; not logged in", "AuthRequired"), + ( + "authentication_error; token_expired", + "OAuthExpired(\"Token expired. Run `claude login` to refresh.\")", + ), + ("running scripts is disabled; requires git-bash", git_bash), + ]; + for (output, expected) in cases { + let error = claude_cli_error_from_output(output).expect(output); + assert_eq!(format!("{error:?}"), expected, "{output}"); + } + assert!(claude_cli_error_from_output("Current session 5% used").is_none()); +} + +#[test] +fn all_percents_fold_case_and_clamp() { + let text = "50% USED\n20 % Left\n101% used\n150% left\n5.5% remaining\n1000% used\n7%Spent 8% available"; + assert_eq!( + percent_matches(text).collect::>(), + vec![50.0, 80.0, 100.0, 0.0, 94.5, 0.0, 7.0, 92.0] + ); + assert!(percent_matches("no numbers here").next().is_none()); +} + +#[test] +fn label_sections_stop_at_their_window_and_the_next_section() { + let filler = |count: usize| vec!["filler"; count].join("\n"); + // Percent on the label line itself, and on the last line of the + // twelve-line window (label + 11). + assert_eq!( + extract_percent_near_label("Current session 30% used", "current session"), + Some(30.0) + ); + let at_last = format!("Current session\n{}\n40% used", filler(10)); + assert_eq!( + extract_percent_near_label(&at_last, "current session"), + Some(40.0) + ); + let past_window = format!("Current session\n{}\n40% used", filler(11)); + assert_eq!( + extract_percent_near_label(&past_window, "current session"), + None + ); + // The next "Current ..." heading ends the section, but the same + // heading does not. + let next_section = "Current session\nCurrent week\n40% used"; + assert_eq!( + extract_percent_near_label(next_section, "current session"), + None + ); + let same_label = "Current session\nCURRENT SESSION again\n40% used"; + assert_eq!( + extract_percent_near_label(same_label, "current session"), + Some(40.0) + ); + // A section without a value falls through to a later label line. + let later = "Current session\nCurrent week\n10% used\nCurrent session\n60% left"; + assert_eq!( + extract_percent_near_label(later, "current session"), + Some(40.0) + ); + assert_eq!( + extract_percent_near_label("Current week (all models)\n10% used", "current week"), + Some(10.0) + ); + + // Reset text uses a fourteen-line window (label + 13). + let reset_last = format!("Current week\n{}\nResets Mon 9am", filler(12)); + assert_eq!( + extract_reset_description(&reset_last, "current week").as_deref(), + Some("Resets Mon 9am") + ); + let reset_past = format!("Current week\n{}\nResets Mon 9am", filler(13)); + assert_eq!(extract_reset_description(&reset_past, "current week"), None); + assert_eq!( + extract_reset_description("Current week 5% used · resets Fri 1pm ", "current week") + .as_deref(), + Some("resets Fri 1pm") + ); + assert_eq!( + extract_reset_description( + "Current session\nCurrent week\nResets Mon", + "current session" + ), + None + ); + let later_reset = "Current session\nCurrent week\nCurrent session\nResets 5pm"; + assert_eq!( + extract_reset_description(later_reset, "current session").as_deref(), + Some("Resets 5pm") + ); +} + +#[test] +fn scoped_weekly_sections_use_a_fourteen_line_window() { + let now = Utc::now(); + let filler = |count: usize| vec!["filler"; count].join("\n"); + let inside = format!("Current week (Opus)\n{}\n25% used", filler(12)); + let limits = extract_cli_scoped_weekly_limits(&inside, now); + assert_eq!(limits.len(), 1); + assert_eq!(limits[0].window.used_percent, 25.0); + let outside = format!("Current week (Opus)\n{}\n25% used", filler(13)); + assert!(extract_cli_scoped_weekly_limits(&outside, now).is_empty()); + let next = "Current week (Opus)\nCurrent session\n25% used"; + assert!(extract_cli_scoped_weekly_limits(next, now).is_empty()); +} + +#[test] +fn cli_parse_usage_error_can_fallback_to_oauth() { + let err = ProviderError::Parse("Claude CLI did not return usage data".to_string()); + + assert!(should_fallback_from_claude_cli_error(&err)); +} + +#[test] +fn cli_auth_error_does_not_fallback_to_oauth() { + assert!(!should_fallback_from_claude_cli_error( + &ProviderError::AuthRequired + )); +} + +#[test] +fn auto_fetch_error_keeps_all_source_failures() { + let err = claude_auto_fetch_error(vec![ + ("OAuth", ProviderError::OAuth("token expired".to_string())), + ("Web", ProviderError::NoCookies), + ( + "CLI", + ProviderError::Parse("Empty output from Claude CLI".to_string()), + ), + ]); + + assert_eq!( + err.to_string(), + "Claude usage failed from all configured sources. OAuth: OAuth error: token expired; Web: No cookies available for web API; CLI: Parse error: Empty output from Claude CLI" + ); +} + +fn oauth_rate_limited() -> ProviderError { + ClaudeOAuthFetcher::rate_limited_error(Duration::from_secs(30)) +} + +#[test] +fn auto_fetch_error_asks_for_a_browser_sign_in_when_only_the_browser_can_help() { + // (CLI failure, retention policy of the plain summary) + let cases = [ + ( + ProviderError::Parse("Claude CLI did not return usage data".to_string()), + LastGoodFailurePolicy::Preserve, + ), + ( + ProviderError::Other("Claude CLI failed: exit status 1".to_string()), + LastGoodFailurePolicy::Replace, + ), + ]; + for (cli_failure, policy) in cases { + let err = claude_auto_fetch_error(vec![ + ("Web", ProviderError::NoCookies), + ("OAuth", oauth_rate_limited()), + ("CLI", cli_failure), + ]); + let ProviderError::BrowserSignInRequired { + message, + sign_in_url, + } = &err + else { + panic!("expected a browser sign-in signal, got {err:?}"); + }; + assert_eq!(sign_in_url, CLAUDE_BROWSER_SIGN_IN_URL); + assert_eq!(err.to_string(), *message); + assert!( + message.starts_with( + "Claude usage failed from all configured sources. Web: No cookies available for web API; OAuth: Transient OAuth error: Claude OAuth usage endpoint is rate limited." + ), + "{message}" + ); + assert!( + message.ends_with("Sign in at https://claude.ai/login in your browser, then refresh."), + "{message}" + ); + // ClaudeProvider::error_state_kind defers to this for every + // variant except a missing CLI. + assert_eq!( + err.state_kind(), + crate::core::ProviderStateKind::NeedsAuthentication + ); + // The hint leaves the desktop retention policy unchanged. + let plain = message + .strip_suffix(browser_sign_in_hint().as_str()) + .map(str::trim_end) + .expect("hint is appended"); + assert_eq!(last_good_failure_policy_for_error(plain), policy); + assert_eq!(last_good_failure_policy_for_error(message), policy); + } +} + +#[test] +fn auto_fetch_error_keeps_other_failure_mixes_untyped() { + let cli_failure = || ProviderError::Parse("Claude CLI did not return usage data".to_string()); + let mixes = [ + // A browser session was there; the Web source failed differently. + vec![ + ("Web", ProviderError::AuthRequired), + ("OAuth", oauth_rate_limited()), + ("CLI", cli_failure()), + ], + // Signed out of Claude Code, not rate limited. + vec![ + ("Web", ProviderError::NoCookies), + ( + "OAuth", + ProviderError::OAuth( + "Claude OAuth credentials not found. Run `claude` to authenticate." + .to_string(), + ), + ), + ("CLI", cli_failure()), + ], + // Another transient OAuth failure. + vec![ + ("Web", ProviderError::NoCookies), + ( + "OAuth", + ProviderError::OAuthTransient( + "Claude OAuth token expired and token refresh is cooling down after a failed attempt." + .to_string(), + ), + ), + ("CLI", cli_failure()), + ], + // The CLI was not tried. + vec![ + ("Web", ProviderError::NoCookies), + ("OAuth", oauth_rate_limited()), + ], + ]; + for failures in mixes { + let err = claude_auto_fetch_error(failures); + assert!(matches!(err, ProviderError::Other(_)), "{err:?}"); + assert!( + !err.to_string().contains(CLAUDE_BROWSER_SIGN_IN_URL), + "{err}" + ); + } +} + +#[test] +fn transient_transport_failure_stops_auto_fallback_and_preserves_last_good() { + let provider = ClaudeProvider::new(); + assert!(provider.retains_last_good_on_transport_failure()); + assert_eq!( + provider.last_good_failure_policy_for_error(&ProviderError::Timeout), + LastGoodFailurePolicy::Preserve + ); + + let mut failures = Vec::new(); + let result = record_auto_source(&mut failures, "Web", Err(ProviderError::Timeout)); + assert!(matches!(result, Err(ProviderError::Timeout))); + assert!(failures.is_empty()); +} + +#[test] +fn rejects_cli_output_that_is_not_a_usage_screen() { + let git_bash = "Claude Code on Windows requires git-bash."; + let claude_2_1 = r#" +I see you've entered `/usage` and `/exit`. + +**Usage**: Token usage and statistics are typically displayed by the CLI interface itself. I don't have direct access to those metrics through my available tools. + +**Exit**: I'll end the session here. Goodbye! +"#; + let legacy = r#" +I see you've entered two slash commands: + +1. `/usage` - This appears to be a request to check usage information +2. `/exit` - This appears to be a request to exit + +However, looking at the available custom slash commands, I don't see these commands defined. +"#; + let activity = r#" +❯ /usage + +Status Config Usage Stats + +Overview Models + +Favorite model: glm-4.6 Total tokens: 263.3k +Sessions: 6 Longest session: 18s +Active days: 2/10 Longest streak: 1 day +"#; + let ansi_activity = "\x1b[2CTotal\x1b[1Ccost:\x1b[12C$0.0000\n\ + \x1b[2CTotal\x1b[1Cduration\x1b[1C(API):\x1b[2C0s\n\ + \x1b[2CUsage:\x1b[17C0\x1b[1Cinput,\x1b[1C0\x1b[1Coutput,\x1b[1C0\x1b[1Ccache\x1b[1Cread"; + let rows: [(&str, &str, Option<&str>); 5] = [ + ( + git_bash, + "parse", + Some("Parse error: Claude CLI did not return usage data"), + ), + ( + claude_2_1, + "other", + Some( + "Claude CLI treated /usage as a normal prompt instead of opening the interactive usage screen. Use Auto, OAuth, or Web mode for Claude usage.", + ), + ), + (legacy, "other", None), + ( + activity, + "other", + Some( + "Claude CLI /usage opened, but this Claude version returned local activity stats instead of plan limit percentages. Use Auto, OAuth, or Web mode for Claude limits.", + ), + ), + (ansi_activity, "other", None), + ]; + for (output, kind, message) in rows { + let err = parse_err(output); + let got = match &err { + ProviderError::Parse(_) => "parse", + ProviderError::Other(_) => "other", + _ => "unexpected", + }; + assert_eq!(got, kind, "{output}"); + if let Some(message) = message { + assert_eq!(err.to_string(), message); + } + } +} + +#[test] +fn accepts_plan_limits_followed_by_activity_stats() { + // Claude Code 2.1.27x on Windows prints the exit summary (cost, + // duration, cache tokens) after the /usage view when the probe ends. + let output = r#" +❯ /usage + +Status Config Usage Stats + +Current session +███████░░░░░░░░░░░░░░░░░░░░░░ 19% used +Resets 3pm (Europe/Berlin) + +Current week (all models) +█████████░░░░░░░░░░░░░░░░░░░░ 31% used +Resets Sep 19, 4pm (Europe/Berlin) + +Total cost: $0.0000 +Total duration (API): 0s +Usage: 0 input, 0 output, 0 cache read +"#; + + let result = parse_ok(output); + + assert_eq!(result.usage.primary.used_percent, 19.0); + assert_eq!( + result + .usage + .secondary + .as_ref() + .map(|window| window.used_percent), + Some(31.0) + ); +} + +// ── Upstream 0.50.1 #2516: revoked vs missing OAuth ──────────────────────── + +#[test] +fn oauth_revoked_error_is_detected() { + assert!(is_oauth_revoked_error(&ProviderError::OAuthRevoked( + "revoked".to_string() + ))); + assert!(!is_oauth_revoked_error(&ProviderError::OAuth( + "expired".to_string() + ))); + assert!(!is_oauth_revoked_error(&ProviderError::AuthRequired)); +} + +#[test] +fn rate_limited_and_revoked_oauth_reuse_the_cli_cache() { + let rate_limited = ProviderError::OAuthTransient( + "Claude OAuth usage endpoint is rate limited. Retrying in about 5m; credentials were preserved." + .to_string(), + ); + assert!(oauth::is_rate_limited_error(&rate_limited)); + assert!(oauth_failure_uses_cli_cache(&rate_limited)); + assert!(oauth_failure_uses_cli_cache(&ProviderError::OAuthRevoked( + "revoked".to_string() + ))); + // Other transient failures and plain expiry still probe the CLI. + assert!(!oauth_failure_uses_cli_cache( + &ProviderError::OAuthTransient("connection reset".to_string()) + )); + assert!(!oauth_failure_uses_cli_cache(&ProviderError::OAuth( + "expired".to_string() + ))); + assert!(!oauth_failure_uses_cli_cache(&ProviderError::AuthRequired)); +} + +#[test] +fn cli_result_cache_round_trips() { + let mut result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "cli"); + result.has_successful_claude_cli_quota = true; + cache_cli_result(result.clone()); + let cached = cached_cli_result().expect("cached result within TTL"); + assert!((cached.usage.primary.used_percent - 42.0).abs() < 0.01); + assert_eq!(cached.source_label, "cli"); + assert!(!cached.has_successful_claude_cli_quota); + + // A live non-CLI success clears the cache. Same test, because the + // global is shared and tests run in parallel without a lock. + clear_cli_result_cache(); + assert!(cached_cli_result().is_none()); +} + +#[test] +fn cli_quota_without_credential_identity_cannot_prove_account_action() { + let result = mark_live_claude_cli_result(parse_ok( + "Current session\n25% used\nCurrent week (all models)\n40% used", + )); + + assert!(result.usage.account_email.is_none()); + assert!(result.has_successful_claude_cli_quota); +} + +#[test] +fn non_cli_fetch_result_does_not_prove_account_action() { + let result = ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(42.0)), "oauth"); + + assert!(!result.has_successful_claude_cli_quota); +} +#[test] +fn error_states_and_last_good_policies() { + use crate::core::ProviderStateKind; + let provider = ClaudeProvider::new(); + // `true` checks the free function on the rendered text; `false` the + // provider method. Each row keeps the path its original test used. + type Row = ( + ProviderError, + ProviderStateKind, + Option<(LastGoodFailurePolicy, bool)>, + ); + let rows: [Row; 6] = [ + ( + ProviderError::NotInstalled( + "Claude CLI not found. Install from https://docs.claude.ai/claude-code".to_string(), + ), + ProviderStateKind::LocalRuntimeOffline, + None, + ), + ( + ProviderError::AuthRequired, + ProviderStateKind::NeedsAuthentication, + None, + ), + (ProviderError::OAuthTransient( + "OAuth error: Claude OAuth usage endpoint is rate limited. Retrying in about 1s; credentials were preserved." + .to_string(), + ), ProviderStateKind::Unknown, Some((LastGoodFailurePolicy::Preserve, false))), + (ProviderError::OAuthTransient( + "Claude OAuth token expired and token refresh is cooling down after a failed attempt. Please retry shortly, or run `claude login`." + .to_string(), + ), ProviderStateKind::Unknown, Some((LastGoodFailurePolicy::Preserve, false))), + (ProviderError::OAuth( + "Claude OAuth credentials not found. Run `claude` to authenticate.".to_string(), + ), ProviderStateKind::NeedsAuthentication, Some((LastGoodFailurePolicy::Replace, true))), + (ProviderError::OAuth("OAuth API returned rate limited".to_string()), ProviderStateKind::NeedsAuthentication, Some((LastGoodFailurePolicy::Replace, false))), + ]; + for (error, state, policy) in rows { + assert_eq!(provider.error_state_kind(&error), state, "{error}"); + match policy { + Some((expected, true)) => assert_eq!( + last_good_failure_policy_for_error(&error.to_string()), + expected + ), + Some((expected, false)) => assert_eq!( + provider.last_good_failure_policy_for_error(&error), + expected, + "{error}" + ), + None => {} + } + } +} diff --git a/rust/src/providers/claude/web_api.rs b/rust/src/providers/claude/web_api.rs index dad46b0bf0..add9d5cfd5 100755 --- a/rust/src/providers/claude/web_api.rs +++ b/rust/src/providers/claude/web_api.rs @@ -58,6 +58,25 @@ fn classify_web_http_error( ProviderError::Other(format!("Failed to get {label}: {status}")) } +/// Pass a success through; otherwise read the body and classify the status. +async fn ensure_success( + response: reqwest::Response, + label: &str, +) -> Result { + let status = response.status(); + if status.is_success() { + return Ok(response); + } + let response_headers = response.headers().clone(); + let body = response.bytes().await?; + Err(classify_web_http_error( + label, + status, + &response_headers, + &body, + )) +} + /// Read the response body as text, then deserialize as JSON. On failure, include /// non-sensitive shape metadata so auth redirects, error envelopes, and schema /// changes are distinguishable without exposing account data in UI/log output. @@ -532,26 +551,7 @@ impl ClaudeWebApiFetcher { } let url = format!("{}/organizations", self.base_url); - - let response = self - .client - .get(&url) - .headers(headers.clone()) - .send() - .await?; - - let status = response.status(); - if !status.is_success() { - let response_headers = response.headers().clone(); - let body = response.bytes().await?; - return Err(classify_web_http_error( - "organizations", - status, - &response_headers, - &body, - )); - } - + let response = ensure_success(self.get(&url, headers).await?, "organizations").await?; let orgs: Vec = parse_json_with_body(response, "organizations").await?; orgs.into_iter() @@ -573,10 +573,7 @@ impl ClaudeWebApiFetcher { ) -> Result { let url = format!("{}/organizations/{}/usage", self.base_url, org_id); let opted_in = self - .client - .get(format!("{url}?{RESET_OPT_IN_QUERY}")) - .headers(headers.clone()) - .send() + .get(&format!("{url}?{RESET_OPT_IN_QUERY}"), headers) .await?; let response = match opted_in.status() { @@ -593,27 +590,15 @@ impl ClaudeWebApiFetcher { &body, )); } - self.get_plain_usage(&url, headers).await? + self.get(&url, headers).await? } - _ => self.get_plain_usage(&url, headers).await?, + _ => self.get(&url, headers).await?, }; - let status = response.status(); - if !status.is_success() { - let response_headers = response.headers().clone(); - let body = response.bytes().await?; - return Err(classify_web_http_error( - "usage", - status, - &response_headers, - &body, - )); - } - - parse_json_with_body(response, "usage").await + parse_json_with_body(ensure_success(response, "usage").await?, "usage").await } - async fn get_plain_usage( + async fn get( &self, url: &str, headers: &reqwest::header::HeaderMap, @@ -621,6 +606,24 @@ impl ClaudeWebApiFetcher { Ok(self.client.get(url).headers(headers.clone()).send().await?) } + /// GET and parse JSON. Unlike [`ensure_success`], a failure reports only + /// the status, without reading the body or mapping 401/403 to auth. + async fn get_json( + &self, + url: &str, + headers: &reqwest::header::HeaderMap, + label: &str, + ) -> Result { + let response = self.get(url, headers).await?; + if !response.status().is_success() { + return Err(ProviderError::Other(format!( + "Failed to get {label}: {}", + response.status() + ))); + } + parse_json_with_body(response, label).await + } + /// Get extra usage (credits) async fn get_extra_usage( &self, @@ -631,22 +634,7 @@ impl ClaudeWebApiFetcher { "{}/organizations/{}/overage_spend_limit", self.base_url, org_id ); - - let response = self - .client - .get(&url) - .headers(headers.clone()) - .send() - .await?; - - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Failed to get extra usage: {}", - response.status() - ))); - } - - parse_json_with_body(response, "extra usage").await + self.get_json(&url, headers, "extra usage").await } /// Best-effort prepaid Extra usage balance. Non-fatal on any failure. @@ -680,22 +668,7 @@ impl ClaudeWebApiFetcher { headers: &reqwest::header::HeaderMap, ) -> Result { let url = format!("{}/account", self.base_url); - - let response = self - .client - .get(&url) - .headers(headers.clone()) - .send() - .await?; - - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "Failed to get account: {}", - response.status() - ))); - } - - parse_json_with_body(response, "account").await + self.get_json(&url, headers, "account").await } /// Convert a usage window to a RateWindow @@ -772,11 +745,6 @@ impl ClaudeWebApiFetcher { fn format_reset_time(dt: DateTime) -> String { dt.format("%b %-d at %-I:%M%p").to_string() } - - /// Convert rate limit tier to plan name - fn tier_to_plan_name(tier: &str) -> String { - super::claude_plan_label(tier) - } } impl Default for ClaudeWebApiFetcher { @@ -873,513 +841,8 @@ fn append_web_extra_windows( } #[cfg(test)] -mod tests { - use super::{ - AccountResponse, ClaudeWebApiFetcher, UsageWindow, classify_web_http_error, cookie_value, - describe_json_body_shape, is_cookie_authentication_failure, - }; - use crate::core::ProviderError; - use reqwest::StatusCode; - use reqwest::header; - use std::sync::{Mutex, OnceLock}; - - fn env_lock() -> &'static Mutex<()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) - } - - #[test] - fn keeps_sub_one_utilization_in_percent_units() { - let window = UsageWindow { - utilization: Some(0.23), - resets_at: None, - }; - - let rate = ClaudeWebApiFetcher::new().to_rate_window(&window, Some(300)); - - assert!((rate.used_percent - 0.23).abs() < f64::EPSILON); - } - - #[test] - fn one_percent_session_is_not_reported_as_full_quota() { - let window = UsageWindow { - utilization: Some(1.0), - resets_at: None, - }; - - let rate = ClaudeWebApiFetcher::new().to_rate_window(&window, Some(300)); - - assert!( - (rate.used_percent - 1.0).abs() < f64::EPSILON, - "session was {}, expected 1% (not 100%)", - rate.used_percent - ); - } - - #[test] - fn null_five_hour_session_is_informational_placeholder() { - let placeholder = crate::core::RateWindow::no_active_session(); - assert!(placeholder.is_informational); - assert_eq!(placeholder.window_minutes, Some(300)); - assert!((placeholder.used_percent - 0.0).abs() < f64::EPSILON); - assert_eq!( - placeholder.reset_description.as_deref(), - Some("No active 5h session") - ); - - // Real idle session (object present at 0%) stays unflagged. - let idle = ClaudeWebApiFetcher::new().to_rate_window( - &UsageWindow { - utilization: Some(0.0), - resets_at: None, - }, - Some(300), - ); - assert!(!idle.is_informational); - } - - #[test] - fn preserves_existing_percentage_utilization() { - let window = UsageWindow { - utilization: Some(23.0), - resets_at: None, - }; - - let rate = ClaudeWebApiFetcher::new().to_rate_window(&window, Some(300)); - - assert!((rate.used_percent - 23.0).abs() < f64::EPSILON); - } - - #[test] - fn labels_max_5x_and_20x_plans() { - assert_eq!( - ClaudeWebApiFetcher::tier_to_plan_name("default_claude_max_5x"), - "Claude Max 5x" - ); - assert_eq!( - ClaudeWebApiFetcher::tier_to_plan_name("v2_default_claude_max_20x"), - "Claude Max 20x" - ); - } - - #[test] - fn resolves_raw_session_key_from_primary_env_var() { - let _guard = env_lock().lock().expect("env lock"); - // SAFETY: running under env_lock() so no other test thread touches the - // environment concurrently; single-threaded w.r.t. these keys. - unsafe { - std::env::remove_var("CLAUDE_AI_SESSION_KEY"); - std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); - std::env::set_var("CLAUDE_AI_SESSION_KEY", "sk-ant-primary"); - std::env::set_var("CLAUDE_WEB_SESSION_KEY", "sk-ant-secondary"); - } - - let session_key = ClaudeWebApiFetcher::resolve_session_key_from_env(); - - assert_eq!(session_key.as_deref(), Some("sk-ant-primary")); - - // SAFETY: same env_lock()-guarded mutation; restoring state after the - // assertions, before the lock is released. - unsafe { - std::env::remove_var("CLAUDE_AI_SESSION_KEY"); - std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); - } - } - - #[test] - fn resolves_session_key_assignment_from_env_var() { - let _guard = env_lock().lock().expect("env lock"); - // SAFETY: env_lock() held for this whole test, so set_var/remove_var - // cannot race another thread's environment access. - unsafe { - std::env::remove_var("CLAUDE_AI_SESSION_KEY"); - std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); - std::env::set_var("CLAUDE_WEB_SESSION_KEY", "sessionKey=sk-ant-cookie-format"); - } - - let session_key = ClaudeWebApiFetcher::resolve_session_key_from_env(); - - assert_eq!(session_key.as_deref(), Some("sk-ant-cookie-format")); - - // SAFETY: cleanup while still holding the env_lock() guard. - unsafe { - std::env::remove_var("CLAUDE_AI_SESSION_KEY"); - std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); - } - } - - #[test] - fn build_headers_include_required_browser_context() { - let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-cookie-format"); - - assert_eq!( - headers - .get(header::COOKIE) - .and_then(|value| value.to_str().ok()), - Some("sessionKey=sk-ant-cookie-format") - ); - assert_eq!( - headers - .get(header::ACCEPT) - .and_then(|value| value.to_str().ok()), - Some("application/json") - ); - assert_eq!( - headers - .get(header::ORIGIN) - .and_then(|value| value.to_str().ok()), - Some("https://claude.ai") - ); - assert_eq!( - headers - .get(header::REFERER) - .and_then(|value| value.to_str().ok()), - Some("https://claude.ai/settings/usage") - ); - assert_eq!( - headers - .get("anthropic-client-platform") - .and_then(|value| value.to_str().ok()), - Some("web_claude_ai") - ); - assert!(headers.contains_key(header::USER_AGENT)); - } - - #[test] - fn stale_cookie_recovery_retries_only_after_authentication_failure() { - assert!(is_cookie_authentication_failure( - &ProviderError::AuthRequired - )); - assert!(!is_cookie_authentication_failure(&ProviderError::Timeout)); - assert!(!is_cookie_authentication_failure(&ProviderError::Other( - "Failed to get organizations: 503 Service Unavailable".to_string(), - ))); - assert!(!is_cookie_authentication_failure(&classify_web_http_error( - "organizations", - StatusCode::FORBIDDEN, - &header::HeaderMap::new(), - b"Just a moment...", - ))); - } - - #[test] - fn malformed_response_shape_does_not_echo_body_contents() { - let shape = describe_json_body_shape( - "sessionKey=secret-session-token", - Some("text/html; charset=utf-8"), - ); - - assert_eq!( - shape, - "content_type=text/html; charset=utf-8, body_len=31, body_kind=non-json" - ); - assert!(!shape.contains("secret-session-token")); - - let object_shape = - describe_json_body_shape(r#"{"z":"secret-value","a":true}"#, Some("application/json")); - assert_eq!( - object_shape, - "content_type=application/json, body_len=29, json_keys=[a, z]" - ); - assert!(!object_shape.contains("secret-value")); - } - - #[test] - fn extracts_last_active_org_from_cookie_header() { - let org = cookie_value( - "foo=bar; sessionKey=sk-ant-session; lastActiveOrg=org-123; other=value", - "lastActiveOrg", - ); - - assert_eq!(org.as_deref(), Some("org-123")); - } - - #[test] - fn account_membership_prefers_nested_organization_uuid() { - let account: AccountResponse = serde_json::from_str( - r#"{ - "email_address": "user@example.com", - "memberships": [ - { - "uuid": "membership-id", - "organization": { "uuid": "org-id" } - } - ] - }"#, - ) - .unwrap(); - - assert_eq!(account.first_membership_org_id().as_deref(), Some("org-id")); - } - - #[test] - fn parses_extra_design_and_routines_aliases() { - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "five_hour": { "utilization": 0.1 }, - "seven_day_omelette": { "utilization": 26 }, - "seven_day_cowork": { "utilization": 11 } - }"#, - ) - .unwrap(); - - let fetcher = ClaudeWebApiFetcher::new(); - let design = usage - .seven_day_design - .as_ref() - .map(|w| fetcher.to_rate_window(w, Some(10080))) - .expect("design window"); - let routines = usage - .seven_day_routines - .as_ref() - .map(|w| fetcher.to_rate_window(w, Some(10080))) - .expect("routines window"); - - assert!((design.used_percent - 26.0).abs() < f64::EPSILON); - assert!((routines.used_percent - 11.0).abs() < f64::EPSILON); - } - - #[test] - fn maps_scoped_weekly_limits_even_when_inactive() { - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "limits": [{ - "kind": "weekly_scoped", - "group": "weekly", - "percent": 7, - "resets_at": "2026-07-16T10:00:00Z", - "scope": {"model": {"id": null, "display_name": "Fable"}}, - "is_active": false - }] - }"#, - ) - .unwrap(); - - let windows = super::super::scoped_weekly::scoped_weekly_windows(&usage.limits); - assert_eq!(windows.len(), 1); - assert_eq!(windows[0].id, "claude-weekly-scoped-fable"); - assert_eq!(windows[0].title, "Fable only"); - } - - #[test] - fn parses_duplicate_design_and_routines_aliases_with_preferred_key() { - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "seven_day_design": { "utilization": 31 }, - "seven_day_omelette": { "utilization": 26 }, - "seven_day_routines": { "utilization": 19 }, - "seven_day_cowork": { "utilization": 11 } - }"#, - ) - .unwrap(); - - let fetcher = ClaudeWebApiFetcher::new(); - let design = usage - .seven_day_design - .as_ref() - .map(|w| fetcher.to_rate_window(w, Some(10080))) - .expect("design window"); - let routines = usage - .seven_day_routines - .as_ref() - .map(|w| fetcher.to_rate_window(w, Some(10080))) - .expect("routines window"); - - assert!((design.used_percent - 31.0).abs() < f64::EPSILON); - assert!((routines.used_percent - 19.0).abs() < f64::EPSILON); - } - - #[test] - fn parses_oauth_apps_window_and_embedded_extra_usage() { - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "five_hour": { "utilization": 0.1 }, - "seven_day_oauth_apps": { "utilization": 42 }, - "extra_usage": { - "is_enabled": true, - "monthly_credit_limit": 2000, - "used_credits": 550, - "currency": "USD" - } - }"#, - ) - .unwrap(); - - let fetcher = ClaudeWebApiFetcher::new(); - let oauth_apps = usage - .seven_day_oauth_apps - .as_ref() - .map(|w| fetcher.to_rate_window(w, Some(10080))) - .expect("oauth apps window"); - let extra = usage.extra_usage.expect("extra usage"); - - assert!((oauth_apps.used_percent - 42.0).abs() < f64::EPSILON); - assert_eq!(extra.is_enabled, Some(true)); - assert_eq!(extra.monthly_credit_limit, Some(2000.0)); - assert_eq!(extra.used_credits, Some(550.0)); - } - - #[test] - fn issue_279_session_limits_win_over_stale_five_hour_after_rollover() { - // Right after a 5h window rollover the legacy five_hour.utilization - // can transiently report 1.0 (normalizes to 100%) even though - // claude.ai shows only 5% for the fresh window. The limits[] entry - // (kind=="session") carries the true value and must win. - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "five_hour": {"utilization": 1.0, "resets_at": "2026-08-13T12:49:59.578826Z"}, - "seven_day": {"utilization": 0.01, "resets_at": "2026-07-26T22:59:59Z"}, - "limits": [ - { - "kind": "session", - "group": "session", - "percent": 5, - "resets_at": "2026-08-13T12:49:59.578826Z" - }, - { - "kind": "weekly_all", - "group": "weekly", - "percent": 1, - "resets_at": "2026-07-26T22:59:59Z" - } - ] - }"#, - ) - .expect("issue 279 body"); - - let fetcher = ClaudeWebApiFetcher::new(); - let (primary, secondary, _) = fetcher.build_rate_windows(&usage); - - // Primary session must be 5%, not the stale 100%. - assert!( - (primary.used_percent - 5.0).abs() < f64::EPSILON, - "primary was {}, expected 5% (not 100%)", - primary.used_percent - ); - assert!((primary.used_percent - 100.0).abs() > 1.0); - assert_eq!(primary.window_minutes, Some(300)); - assert!(primary.resets_at.is_some()); - - // Weekly lane is unaffected (still prefers limits weekly_all). - let weekly = secondary.expect("weekly"); - assert!((weekly.used_percent - 1.0).abs() < f64::EPSILON); - } - - #[test] - fn session_falls_back_to_legacy_five_hour_without_limits_entry() { - // When no limits[] session entry exists, the legacy five_hour field - // is still the source of truth (backwards compatible). - let usage: super::UsageResponse = serde_json::from_str( - r#"{ - "five_hour": {"utilization": 10.0, "resets_at": "2026-08-13T12:49:59Z"} - }"#, - ) - .expect("legacy-only body"); - - let fetcher = ClaudeWebApiFetcher::new(); - let (primary, _, _) = fetcher.build_rate_windows(&usage); - - assert!((primary.used_percent - 10.0).abs() < f64::EPSILON); - assert_eq!(primary.window_minutes, Some(300)); - } - - #[test] - fn parse_prepaid_balance_converts_cents_to_dollars() { - let balance = super::parse_prepaid_balance(r#"{"amount": 2550, "currency": "usd"}"#) - .expect("prepaid balance"); - assert!((balance.amount_dollars - 25.5).abs() < f64::EPSILON); - assert_eq!(balance.currency_code, "USD"); - } - - #[test] - fn parse_prepaid_balance_rejects_negative_or_non_finite() { - assert!(super::parse_prepaid_balance(r#"{"amount": -1, "currency": "USD"}"#).is_none()); - assert!(super::parse_prepaid_balance(r#"{"amount": 10, "currency": " "}"#).is_none()); - } - - #[test] - fn apply_prepaid_balance_attaches_to_same_currency_cost() { - let existing = crate::core::CostSnapshot::new(1.0, "USD", "Monthly").with_limit(20.0); - let balance = super::PrepaidBalance { - amount_dollars: 12.34, - currency_code: "USD".into(), - }; - let cost = super::apply_prepaid_balance(balance, Some(existing)); - assert_eq!(cost.balance, Some(12.34)); - assert!((cost.used - 1.0).abs() < f64::EPSILON); - assert_eq!(cost.limit, Some(20.0)); - assert_eq!(cost.period, "Monthly"); - } - - #[test] - fn apply_prepaid_balance_creates_extra_usage_when_missing_or_mismatch() { - let balance = super::PrepaidBalance { - amount_dollars: 5.0, - currency_code: "USD".into(), - }; - let created = super::apply_prepaid_balance(balance.clone(), None); - assert_eq!(created.balance, Some(5.0)); - assert_eq!(created.period, "Extra usage"); - assert!((created.used - 0.0).abs() < f64::EPSILON); - - let eur = crate::core::CostSnapshot::new(2.0, "EUR", "Monthly"); - let replaced = super::apply_prepaid_balance(balance, Some(eur)); - assert_eq!(replaced.currency_code, "USD"); - assert_eq!(replaced.period, "Extra usage"); - assert_eq!(replaced.balance, Some(5.0)); - } - - #[test] - fn web_extras_order_oauth_scoped_then_routines() { - use crate::core::{NamedRateWindow, RateWindow, UsageSnapshot}; - - let mut snapshot = UsageSnapshot::new(RateWindow::new(10.0)); - super::append_web_extra_windows( - &mut snapshot, - Some(RateWindow::new(1.0)), - vec![NamedRateWindow::new( - "claude-weekly-scoped-fable", - "Fable only", - RateWindow::new(2.0), - )], - Some(RateWindow::new(3.0)), - ); - - let ids: Vec<&str> = snapshot - .extra_rate_windows - .iter() - .map(|w| w.id.as_str()) - .collect(); - assert_eq!( - ids, - vec![ - "claude-oauth-apps", - "claude-weekly-scoped-fable", - "claude-routines" - ] - ); - } - - #[test] - fn web_extras_keep_routines_in_raw_snapshot() { - use crate::core::{NamedRateWindow, RateWindow, UsageSnapshot}; - - let mut snapshot = UsageSnapshot::new(RateWindow::new(10.0)); - super::append_web_extra_windows( - &mut snapshot, - Some(RateWindow::new(1.0)), - vec![NamedRateWindow::new( - "claude-weekly-scoped-fable", - "Fable only", - RateWindow::new(2.0), - )], - Some(RateWindow::new(3.0)), - ); - - assert_eq!(snapshot.extra_rate_windows.len(), 3); - assert_eq!(snapshot.extra_rate_windows[2].id, "claude-routines"); - } -} +#[path = "web_api_tests.rs"] +mod tests; #[cfg(test)] #[path = "cloudflare_tests.rs"] diff --git a/rust/src/providers/claude/web_api_tests.rs b/rust/src/providers/claude/web_api_tests.rs new file mode 100644 index 0000000000..2b591aff3a --- /dev/null +++ b/rust/src/providers/claude/web_api_tests.rs @@ -0,0 +1,424 @@ +use super::{ + AccountResponse, ClaudeWebApiFetcher, UsageWindow, classify_web_http_error, cookie_value, + describe_json_body_shape, is_cookie_authentication_failure, +}; +use crate::core::ProviderError; +use reqwest::StatusCode; +use reqwest::header; +use std::sync::{Mutex, OnceLock}; + +fn env_lock() -> &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) +} + +#[test] +fn keeps_utilization_in_percent_units() { + // 1.0 is a 1% session, not a full quota. + for utilization in [0.23, 1.0, 23.0] { + let window = UsageWindow { + utilization: Some(utilization), + resets_at: None, + }; + + let rate = ClaudeWebApiFetcher::new().to_rate_window(&window, Some(300)); + + assert!( + (rate.used_percent - utilization).abs() < f64::EPSILON, + "session was {}, expected {utilization}% (not 100%)", + rate.used_percent + ); + } +} + +#[test] +fn null_five_hour_session_is_informational_placeholder() { + let placeholder = crate::core::RateWindow::no_active_session(); + assert!(placeholder.is_informational); + assert_eq!(placeholder.window_minutes, Some(300)); + assert!((placeholder.used_percent - 0.0).abs() < f64::EPSILON); + assert_eq!( + placeholder.reset_description.as_deref(), + Some("No active 5h session") + ); + + // Real idle session (object present at 0%) stays unflagged. + let idle = ClaudeWebApiFetcher::new().to_rate_window( + &UsageWindow { + utilization: Some(0.0), + resets_at: None, + }, + Some(300), + ); + assert!(!idle.is_informational); +} + +#[test] +fn labels_max_5x_and_20x_plans() { + assert_eq!( + crate::providers::claude::claude_plan_label("default_claude_max_5x"), + "Claude Max 5x" + ); + assert_eq!( + crate::providers::claude::claude_plan_label("v2_default_claude_max_20x"), + "Claude Max 20x" + ); +} + +#[test] +fn resolves_session_key_from_env_vars() { + let _guard = env_lock().lock().expect("env lock"); + // (CLAUDE_AI_SESSION_KEY, CLAUDE_WEB_SESSION_KEY, resolved key) + let rows = [ + (Some("sk-ant-primary"), "sk-ant-secondary", "sk-ant-primary"), + ( + None, + "sessionKey=sk-ant-cookie-format", + "sk-ant-cookie-format", + ), + ]; + for (primary, secondary, expected) in rows { + // SAFETY: env_lock() held for this whole test, so set_var/remove_var + // cannot race another thread's environment access. + unsafe { + std::env::remove_var("CLAUDE_AI_SESSION_KEY"); + std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); + if let Some(primary) = primary { + std::env::set_var("CLAUDE_AI_SESSION_KEY", primary); + } + std::env::set_var("CLAUDE_WEB_SESSION_KEY", secondary); + } + + let session_key = ClaudeWebApiFetcher::resolve_session_key_from_env(); + + assert_eq!(session_key.as_deref(), Some(expected)); + } + + // SAFETY: cleanup while still holding the env_lock() guard. + unsafe { + std::env::remove_var("CLAUDE_AI_SESSION_KEY"); + std::env::remove_var("CLAUDE_WEB_SESSION_KEY"); + } +} + +#[test] +fn build_headers_include_required_browser_context() { + let headers = ClaudeWebApiFetcher::build_headers("sessionKey=sk-ant-cookie-format"); + + for (name, value) in [ + (header::COOKIE.as_str(), "sessionKey=sk-ant-cookie-format"), + (header::ACCEPT.as_str(), "application/json"), + (header::ORIGIN.as_str(), "https://claude.ai"), + (header::REFERER.as_str(), "https://claude.ai/settings/usage"), + ("anthropic-client-platform", "web_claude_ai"), + ] { + assert_eq!( + headers.get(name).and_then(|value| value.to_str().ok()), + Some(value), + "{name}" + ); + } + assert!(headers.contains_key(header::USER_AGENT)); +} + +#[test] +fn stale_cookie_recovery_retries_only_after_authentication_failure() { + assert!(is_cookie_authentication_failure( + &ProviderError::AuthRequired + )); + assert!(!is_cookie_authentication_failure(&ProviderError::Timeout)); + assert!(!is_cookie_authentication_failure(&ProviderError::Other( + "Failed to get organizations: 503 Service Unavailable".to_string(), + ))); + assert!(!is_cookie_authentication_failure(&classify_web_http_error( + "organizations", + StatusCode::FORBIDDEN, + &header::HeaderMap::new(), + b"Just a moment...", + ))); +} + +#[test] +fn malformed_response_shape_does_not_echo_body_contents() { + let shape = describe_json_body_shape( + "sessionKey=secret-session-token", + Some("text/html; charset=utf-8"), + ); + + assert_eq!( + shape, + "content_type=text/html; charset=utf-8, body_len=31, body_kind=non-json" + ); + assert!(!shape.contains("secret-session-token")); + + let object_shape = + describe_json_body_shape(r#"{"z":"secret-value","a":true}"#, Some("application/json")); + assert_eq!( + object_shape, + "content_type=application/json, body_len=29, json_keys=[a, z]" + ); + assert!(!object_shape.contains("secret-value")); +} + +#[test] +fn extracts_last_active_org_from_cookie_header() { + let org = cookie_value( + "foo=bar; sessionKey=sk-ant-session; lastActiveOrg=org-123; other=value", + "lastActiveOrg", + ); + + assert_eq!(org.as_deref(), Some("org-123")); +} + +#[test] +fn account_membership_prefers_nested_organization_uuid() { + let account: AccountResponse = serde_json::from_str( + r#"{ + "email_address": "user@example.com", + "memberships": [ + { + "uuid": "membership-id", + "organization": { "uuid": "org-id" } + } + ] + }"#, + ) + .unwrap(); + + assert_eq!(account.first_membership_org_id().as_deref(), Some("org-id")); +} + +#[test] +fn parses_design_and_routines_aliases_preferring_the_named_key() { + let rows = [ + ( + r#"{ + "five_hour": { "utilization": 0.1 }, + "seven_day_omelette": { "utilization": 26 }, + "seven_day_cowork": { "utilization": 11 } + }"#, + 26.0, + 11.0, + ), + ( + r#"{ + "seven_day_design": { "utilization": 31 }, + "seven_day_omelette": { "utilization": 26 }, + "seven_day_routines": { "utilization": 19 }, + "seven_day_cowork": { "utilization": 11 } + }"#, + 31.0, + 19.0, + ), + ]; + let fetcher = ClaudeWebApiFetcher::new(); + for (json, design_percent, routines_percent) in rows { + let usage: super::UsageResponse = serde_json::from_str(json).unwrap(); + let design = usage + .seven_day_design + .as_ref() + .map(|w| fetcher.to_rate_window(w, Some(10080))) + .expect("design window"); + let routines = usage + .seven_day_routines + .as_ref() + .map(|w| fetcher.to_rate_window(w, Some(10080))) + .expect("routines window"); + + assert!((design.used_percent - design_percent).abs() < f64::EPSILON); + assert!((routines.used_percent - routines_percent).abs() < f64::EPSILON); + } +} + +#[test] +fn maps_scoped_weekly_limits_even_when_inactive() { + let usage: super::UsageResponse = serde_json::from_str( + r#"{ + "limits": [{ + "kind": "weekly_scoped", + "group": "weekly", + "percent": 7, + "resets_at": "2026-07-16T10:00:00Z", + "scope": {"model": {"id": null, "display_name": "Fable"}}, + "is_active": false + }] + }"#, + ) + .unwrap(); + + let windows = super::super::scoped_weekly::scoped_weekly_windows(&usage.limits); + assert_eq!(windows.len(), 1); + assert_eq!(windows[0].id, "claude-weekly-scoped-fable"); + assert_eq!(windows[0].title, "Fable only"); +} + +#[test] +fn parses_oauth_apps_window_and_embedded_extra_usage() { + let usage: super::UsageResponse = serde_json::from_str( + r#"{ + "five_hour": { "utilization": 0.1 }, + "seven_day_oauth_apps": { "utilization": 42 }, + "extra_usage": { + "is_enabled": true, + "monthly_credit_limit": 2000, + "used_credits": 550, + "currency": "USD" + } + }"#, + ) + .unwrap(); + + let fetcher = ClaudeWebApiFetcher::new(); + let oauth_apps = usage + .seven_day_oauth_apps + .as_ref() + .map(|w| fetcher.to_rate_window(w, Some(10080))) + .expect("oauth apps window"); + let extra = usage.extra_usage.expect("extra usage"); + + assert!((oauth_apps.used_percent - 42.0).abs() < f64::EPSILON); + assert_eq!(extra.is_enabled, Some(true)); + assert_eq!(extra.monthly_credit_limit, Some(2000.0)); + assert_eq!(extra.used_credits, Some(550.0)); +} + +#[test] +fn issue_279_session_limits_win_over_stale_five_hour_after_rollover() { + // Right after a 5h window rollover the legacy five_hour.utilization + // can transiently report 1.0 (normalizes to 100%) even though + // claude.ai shows only 5% for the fresh window. The limits[] entry + // (kind=="session") carries the true value and must win. + let usage: super::UsageResponse = serde_json::from_str( + r#"{ + "five_hour": {"utilization": 1.0, "resets_at": "2026-08-13T12:49:59.578826Z"}, + "seven_day": {"utilization": 0.01, "resets_at": "2026-07-26T22:59:59Z"}, + "limits": [ + { + "kind": "session", + "group": "session", + "percent": 5, + "resets_at": "2026-08-13T12:49:59.578826Z" + }, + { + "kind": "weekly_all", + "group": "weekly", + "percent": 1, + "resets_at": "2026-07-26T22:59:59Z" + } + ] + }"#, + ) + .expect("issue 279 body"); + + let fetcher = ClaudeWebApiFetcher::new(); + let (primary, secondary, _) = fetcher.build_rate_windows(&usage); + + // Primary session must be 5%, not the stale 100%. + assert!( + (primary.used_percent - 5.0).abs() < f64::EPSILON, + "primary was {}, expected 5% (not 100%)", + primary.used_percent + ); + assert!((primary.used_percent - 100.0).abs() > 1.0); + assert_eq!(primary.window_minutes, Some(300)); + assert!(primary.resets_at.is_some()); + + // Weekly lane is unaffected (still prefers limits weekly_all). + let weekly = secondary.expect("weekly"); + assert!((weekly.used_percent - 1.0).abs() < f64::EPSILON); +} + +#[test] +fn session_falls_back_to_legacy_five_hour_without_limits_entry() { + // When no limits[] session entry exists, the legacy five_hour field + // is still the source of truth (backwards compatible). + let usage: super::UsageResponse = serde_json::from_str( + r#"{ + "five_hour": {"utilization": 10.0, "resets_at": "2026-08-13T12:49:59Z"} + }"#, + ) + .expect("legacy-only body"); + + let fetcher = ClaudeWebApiFetcher::new(); + let (primary, _, _) = fetcher.build_rate_windows(&usage); + + assert!((primary.used_percent - 10.0).abs() < f64::EPSILON); + assert_eq!(primary.window_minutes, Some(300)); +} + +#[test] +fn parse_prepaid_balance_converts_cents_to_dollars() { + let balance = super::parse_prepaid_balance(r#"{"amount": 2550, "currency": "usd"}"#) + .expect("prepaid balance"); + assert!((balance.amount_dollars - 25.5).abs() < f64::EPSILON); + assert_eq!(balance.currency_code, "USD"); +} + +#[test] +fn parse_prepaid_balance_rejects_negative_or_non_finite() { + assert!(super::parse_prepaid_balance(r#"{"amount": -1, "currency": "USD"}"#).is_none()); + assert!(super::parse_prepaid_balance(r#"{"amount": 10, "currency": " "}"#).is_none()); +} + +#[test] +fn apply_prepaid_balance_attaches_to_same_currency_cost() { + let existing = crate::core::CostSnapshot::new(1.0, "USD", "Monthly").with_limit(20.0); + let balance = super::PrepaidBalance { + amount_dollars: 12.34, + currency_code: "USD".into(), + }; + let cost = super::apply_prepaid_balance(balance, Some(existing)); + assert_eq!(cost.balance, Some(12.34)); + assert!((cost.used - 1.0).abs() < f64::EPSILON); + assert_eq!(cost.limit, Some(20.0)); + assert_eq!(cost.period, "Monthly"); +} + +#[test] +fn apply_prepaid_balance_creates_extra_usage_when_missing_or_mismatch() { + let balance = super::PrepaidBalance { + amount_dollars: 5.0, + currency_code: "USD".into(), + }; + let created = super::apply_prepaid_balance(balance.clone(), None); + assert_eq!(created.balance, Some(5.0)); + assert_eq!(created.period, "Extra usage"); + assert!((created.used - 0.0).abs() < f64::EPSILON); + + let eur = crate::core::CostSnapshot::new(2.0, "EUR", "Monthly"); + let replaced = super::apply_prepaid_balance(balance, Some(eur)); + assert_eq!(replaced.currency_code, "USD"); + assert_eq!(replaced.period, "Extra usage"); + assert_eq!(replaced.balance, Some(5.0)); +} + +#[test] +fn web_extras_order_oauth_scoped_then_routines() { + use crate::core::{NamedRateWindow, RateWindow, UsageSnapshot}; + + let mut snapshot = UsageSnapshot::new(RateWindow::new(10.0)); + super::append_web_extra_windows( + &mut snapshot, + Some(RateWindow::new(1.0)), + vec![NamedRateWindow::new( + "claude-weekly-scoped-fable", + "Fable only", + RateWindow::new(2.0), + )], + Some(RateWindow::new(3.0)), + ); + + let ids: Vec<&str> = snapshot + .extra_rate_windows + .iter() + .map(|w| w.id.as_str()) + .collect(); + assert_eq!( + ids, + vec![ + "claude-oauth-apps", + "claude-weekly-scoped-fable", + "claude-routines" + ] + ); +} diff --git a/rust/src/providers/codex/api.rs b/rust/src/providers/codex/api.rs index 5baf943126..b376525659 100755 --- a/rust/src/providers/codex/api.rs +++ b/rust/src/providers/codex/api.rs @@ -3,44 +3,25 @@ //! Uses OAuth tokens stored by the Codex CLI in ~/.codex/auth.json use super::{pat, weekly_reset}; -use crate::core::{ - CostSnapshot, NamedRateWindow, ProviderError, RateWindow, RateWindowCadence, UsageSnapshot, -}; -use crate::providers::openai::OpenAISubscriptionFetchResult; -use base64::Engine; -use chrono::{DateTime, TimeZone, Utc}; -use serde::Deserialize; -use std::collections::HashMap; +use crate::core::{CostSnapshot, ProviderError, UsageSnapshot}; +use chrono::Utc; use std::path::PathBuf; -use std::sync::{Arc, Mutex, OnceLock}; -use std::time::{Duration, Instant, SystemTime}; -use tokio::sync::Mutex as AsyncMutex; +use std::time::Instant; + +mod credentials; +mod parse; +mod reset_credits; + +use credentials::CodexCredentials; +use parse::format_plan_type; +use reset_credits::apply_reset_credits_window; +pub(super) use reset_credits::{ResetCredit, ResetCredits, next_available_reset_credit_expiry}; #[path = "subscription.rs"] mod subscription; const DEFAULT_BASE_URL: &str = "https://chatgpt.com/backend-api"; const USAGE_PATH: &str = "/wham/usage"; -const RESET_CREDITS_PATH: &str = "/wham/rate-limit-reset-credits"; -const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(5); -/// Upstream 0.69.0 #4088: the Codex CLI owns `auth.json` and may be publishing a -/// replacement while we read it. A failed or stale read is repeated up to this many -/// times, `CREDENTIAL_READ_RETRY_DELAY` apart, before the error is reported. -const CREDENTIAL_READ_RETRIES: u32 = 2; -const CREDENTIAL_READ_RETRY_DELAY: Duration = Duration::from_millis(50); -const RESET_CREDITS_CACHE_TTL: Duration = Duration::from_secs(600); -const EXTERNAL_OAUTH_REFRESH_WINDOW: chrono::TimeDelta = chrono::Duration::minutes(5); - -static CREDENTIAL_CACHE: OnceLock>> = OnceLock::new(); -static RESET_CREDITS_CACHE: OnceLock>>>> = - OnceLock::new(); - -#[derive(Default)] -struct ResetCreditsCache { - loaded_at: Option, - value: Option, - confirmation_failure_at: Option, -} /// Codex API client pub struct CodexApi { @@ -110,14 +91,14 @@ impl CodexApi { { usage = usage.with_login_method(format_plan_type(&plan_type)); } - let usage = self - .enrich_subscription_metadata( - &self.resolve_base_url(), - &token, - account_id.as_deref(), - usage, - ) - .await; + let usage = subscription::enrich_subscription_metadata( + self, + &self.resolve_base_url(), + &token, + account_id.as_deref(), + usage, + ) + .await; Ok((usage, cost, account_identity)) } @@ -198,22 +179,22 @@ impl CodexApi { %error, "Codex weekly reset confirmation failed; preserving first successful usage" ); - let result = Self::preserve_after_confirmation_failure( - &state, - first_usage, - first_cost, - ); + let usage = weekly_reset::preserve_weekly(&state, first_usage); weekly_reset::save(&scope, &state); - let (usage, cost) = result; - let usage = self - .enrich_subscription_metadata( - &base_url, - &creds.access_token, - creds.account_id.as_deref(), - usage, - ) - .await; - return Ok((usage, cost, creds.account_id.clone(), displayed_credits)); + let usage = subscription::enrich_subscription_metadata( + self, + &base_url, + &creds.access_token, + creds.account_id.as_deref(), + usage, + ) + .await; + return Ok(( + usage, + first_cost, + creds.account_id.clone(), + displayed_credits, + )); } }; let confirmation_credits = if initial_credits.is_some() { @@ -256,51 +237,36 @@ impl CodexApi { } } }; - let usage = self - .enrich_subscription_metadata( - &base_url, - &creds.access_token, - creds.account_id.as_deref(), - usage, - ) - .await; + let usage = subscription::enrich_subscription_metadata( + self, + &base_url, + &creds.access_token, + creds.account_id.as_deref(), + usage, + ) + .await; let usage = apply_reset_credits_window(usage, displayed_credits.as_ref()); Ok((usage, cost, creds.account_id.clone(), displayed_credits)) } - /// Subscription metadata is optional enrichment. Usage remains usable when - /// the endpoint is unavailable, malformed, unauthorized, or points at a - /// custom backend. A successful empty cancellation response is the only - /// result allowed to clear dates on the fresh snapshot. - async fn enrich_subscription_metadata( + /// Bearer GET shared by the ChatGPT backend endpoints; an empty account id + /// sends no `ChatGPT-Account-Id` header. + pub(super) fn authed_get( &self, - base_url: &str, - access_token: &str, - account_id: Option<&str>, - usage: UsageSnapshot, - ) -> UsageSnapshot { - subscription::enrich_subscription_metadata(self, base_url, access_token, account_id, usage) - .await - } - - async fn fetch_subscription_metadata( - &self, - base_url: &str, + url: &str, access_token: &str, account_id: Option<&str>, - ) -> OpenAISubscriptionFetchResult { - subscription::fetch_subscription_metadata(self, base_url, access_token, account_id).await - } - - fn preserve_after_confirmation_failure( - state: &weekly_reset::AccountState, - first_usage: UsageSnapshot, - first_cost: Option, - ) -> (UsageSnapshot, Option) { - ( - weekly_reset::preserve_weekly(state, first_usage), - first_cost, - ) + ) -> reqwest::RequestBuilder { + let request = self + .client + .get(url) + .header("Authorization", format!("Bearer {access_token}")) + .header("User-Agent", "CodexBar") + .header("Accept", "application/json"); + match account_id.filter(|id| !id.is_empty()) { + Some(account_id) => request.header("ChatGPT-Account-Id", account_id), + None => request, + } } async fn fetch_usage_once( @@ -309,19 +275,11 @@ impl CodexApi { base_url: &str, ) -> Result<(UsageSnapshot, Option, Option), ProviderError> { let url = format!("{}{}", base_url, USAGE_PATH); - let mut request = self - .client - .get(&url) - .header("Authorization", format!("Bearer {}", creds.access_token)) - .header("User-Agent", "CodexBar") - .header("Accept", "application/json") - .timeout(std::time::Duration::from_secs(30)); - if let Some(account_id) = &creds.account_id - && !account_id.is_empty() - { - request = request.header("ChatGPT-Account-Id", account_id); - } - let response = request.send().await?; + let response = self + .authed_get(&url, &creds.access_token, creds.account_id.as_deref()) + .timeout(std::time::Duration::from_secs(30)) + .send() + .await?; if !response.status().is_success() { return Err(super::authenticated_http_error(response, "Codex API").await); } @@ -337,368 +295,6 @@ impl CodexApi { Ok((usage, cost, reset_credits)) } - fn reset_credits_cache_slot( - &self, - creds: &CodexCredentials, - base_url: &str, - ) -> Arc> { - // The Codex home is part of the scope: two homes never share an - // observation, even when they hold the same account and token. - let auth_path = self.get_auth_path(); - let home = weekly_reset::scope_key(None, &auth_path); - let account = weekly_reset::scope_key(creds.account_id.as_deref(), &auth_path); - let token = weekly_reset::scope_key(Some(&creds.access_token), &auth_path); - let key = format!( - "{}|{home}|{account}|{token}", - base_url.trim_end_matches('/') - ); - let cache = RESET_CREDITS_CACHE.get_or_init(|| Mutex::new(HashMap::new())); - let mut cache = cache - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - Arc::clone( - cache - .entry(key) - .or_insert_with(|| Arc::new(AsyncMutex::new(ResetCreditsCache::default()))), - ) - } - - async fn fetch_rate_limit_reset_credits_cached( - &self, - creds: &CodexCredentials, - base_url: &str, - ) -> Option { - let slot = self.reset_credits_cache_slot(creds, base_url); - let mut cache = slot.lock().await; - if cache - .loaded_at - .is_some_and(|loaded| loaded.elapsed() < RESET_CREDITS_CACHE_TTL) - { - return cache.value.clone(); - } - cache.value = self - .fetch_rate_limit_reset_credits(creds, base_url) - .await - .ok(); - cache.loaded_at = Some(Instant::now()); - cache.confirmation_failure_at = None; - cache.value.clone() - } - - /// Reset credits for the initial weekly-reset decision. A pending delayed - /// candidate is revalidated against the current inventory, which must be - /// observed after the candidate was stored: the ten-minute cache can still - /// hold the very observation that created it, so that case reads fresh. - async fn initial_reset_credits( - &self, - creds: &CodexCredentials, - base_url: &str, - started: Instant, - observed: Option, - candidate_pending: bool, - ) -> Option { - if !candidate_pending { - return observed; - } - self.fresh_reset_credits_for_confirmation(creds, base_url, started) - .await - } - - async fn fresh_reset_credits_for_confirmation( - &self, - creds: &CodexCredentials, - base_url: &str, - started: Instant, - ) -> Option { - let slot = self.reset_credits_cache_slot(creds, base_url); - let mut cache = slot.lock().await; - cache.value.as_ref()?; - if cache - .confirmation_failure_at - .is_some_and(|failed| failed.elapsed() < RESET_CREDITS_CACHE_TTL) - { - return None; - } - if cache.loaded_at.is_some_and(|loaded| loaded >= started) { - return cache.value.clone(); - } - let fresh = self - .fetch_rate_limit_reset_credits(creds, base_url) - .await - .ok(); - if let Some(value) = fresh.as_ref() { - cache.value = Some(value.clone()); - cache.loaded_at = Some(Instant::now()); - cache.confirmation_failure_at = None; - } else { - cache.confirmation_failure_at = Some(Instant::now()); - } - fresh - } - - async fn fetch_rate_limit_reset_credits_fresh( - &self, - creds: &CodexCredentials, - base_url: &str, - ) -> Option { - let slot = self.reset_credits_cache_slot(creds, base_url); - let mut cache = slot.lock().await; - if cache - .confirmation_failure_at - .is_some_and(|failed| failed.elapsed() < RESET_CREDITS_CACHE_TTL) - { - return None; - } - let fresh = self - .fetch_rate_limit_reset_credits(creds, base_url) - .await - .ok(); - if let Some(value) = fresh.as_ref() { - cache.value = Some(value.clone()); - cache.loaded_at = Some(Instant::now()); - cache.confirmation_failure_at = None; - } else { - cache.confirmation_failure_at = Some(Instant::now()); - } - fresh - } - - async fn fetch_rate_limit_reset_credits( - &self, - creds: &CodexCredentials, - base_url: &str, - ) -> Result { - let mut request = self - .client - .get(format!("{}{}", base_url, RESET_CREDITS_PATH)) - .header("Authorization", format!("Bearer {}", creds.access_token)) - .header("User-Agent", "CodexBar") - .header("Accept", "application/json"); - if let Some(account_id) = &creds.account_id - && !account_id.is_empty() - { - request = request.header("ChatGPT-Account-Id", account_id); - } - let response = request.send().await?; - if !response.status().is_success() { - return Err(super::authenticated_http_error(response, "Codex reset credits").await); - } - decode_reset_credits(&response.bytes().await?) - } - - /// Load credentials, tolerating a brief owner publication of `auth.json`. - /// - /// Upstream 0.69.0 #4088 (`CodexOAuthFetchStrategy.loadCredentials` on the - /// usage path, `retryStale: true`): every failed read is repeated. That covers - /// a missing (`NotInstalled`), unreadable (`Other`), malformed or incomplete - /// (`Parse`) file, and a credential the gate rejects as stale (`AuthRequired`, - /// such as a token inside its renewal window), because the CLI may be - /// publishing its renewal. This only rereads the file: no token is redeemed, - /// nothing is written, and the credential cache semantics are unchanged. After - /// the last read the error keeps its category, so unchanged stale credentials - /// still need their owner's renewal. - async fn load_credentials(&self) -> Result { - Self::reread_during_owner_publication(|| self.load_credentials_once()).await - } - - /// The bounded reread behind [`Self::load_credentials`]: `read` runs up to - /// `1 + CREDENTIAL_READ_RETRIES` times, `CREDENTIAL_READ_RETRY_DELAY` apart, - /// until it succeeds, and the last result is returned unchanged. Dropping the - /// returned future cancels the pending delay and any further read (upstream - /// checks task cancellation before each read). - async fn reread_during_owner_publication( - mut read: impl FnMut() -> Result, - ) -> Result { - let mut retries_remaining = CREDENTIAL_READ_RETRIES; - loop { - match read() { - Err(_) if retries_remaining > 0 => { - retries_remaining -= 1; - tokio::time::sleep(CREDENTIAL_READ_RETRY_DELAY).await; - } - result => return result, - } - } - } - - fn load_credentials_once(&self) -> Result { - let auth_path = self.get_auth_path(); - - let metadata = - std::fs::metadata(&auth_path).map_err(|error| self.credential_file_error(error))?; - let modified = metadata.modified().ok(); - if let Some(cached) = Self::cached_credentials(&auth_path, modified) { - Self::enforce_external_oauth_gate(&cached)?; - return Ok(cached); - } - - let content = std::fs::read_to_string(&auth_path) - .map_err(|error| self.credential_file_error(error))?; - - let credentials = Self::parse_credentials_json(&content)?; - Self::enforce_external_oauth_gate(&credentials)?; - Self::store_cached_credentials(auth_path, modified, credentials.clone()); - Ok(credentials) - } - - fn missing_credentials_error(&self) -> ProviderError { - // Upstream 0.50.0 #2679: when the CLI targets Amazon Bedrock or - // another custom backend without ChatGPT auth, sign-in guidance - // is wrong — rate limits simply are not available there. - if self.uses_custom_backend() { - return ProviderError::NotInstalled( - "Codex uses a custom backend (chatgpt_base_url / model_provider) without \ - ChatGPT auth. ChatGPT rate limits are unavailable for this setup." - .to_string(), - ); - } - - ProviderError::NotInstalled( - "Codex auth.json not found. Run `codex login` in a terminal to sign in.".to_string(), - ) - } - - fn credential_file_error(&self, error: std::io::Error) -> ProviderError { - if error.kind() == std::io::ErrorKind::NotFound { - return self.missing_credentials_error(); - } - - ProviderError::Other(format!("Failed to read Codex credentials: {error}")) - } - - fn parse_credentials_json(content: &str) -> Result { - let json: serde_json::Value = serde_json::from_str(content) - .map_err(|e| ProviderError::Parse(format!("Invalid Codex credentials JSON: {}", e)))?; - - // Check for OPENAI_API_KEY first - if let Some(api_key) = json.get("OPENAI_API_KEY").and_then(|v| v.as_str()) { - let trimmed = api_key.trim(); - if !trimmed.is_empty() { - return Ok(CodexCredentials { - access_token: trimmed.to_string(), - account_id: None, - is_external_oauth: false, - access_token_expires_at: None, - last_refresh: None, - }); - } - } - - // Otherwise, look for tokens object (external OAuth source) - let tokens = json.get("tokens").ok_or_else(|| { - ProviderError::Parse("Codex auth.json exists but contains no tokens.".to_string()) - })?; - - let access_token = tokens - .get("access_token") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .ok_or_else(|| { - ProviderError::Parse("Missing access_token in Codex credentials".to_string()) - })? - .to_string(); - - let account_id = tokens - .get("account_id") - .and_then(|v| v.as_str()) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()); - - // Upstream 0.50.1 #2944: an OAuth token set with a refresh_token is an - // external (CLI-owned) OAuth source. The `last_refresh` timestamp is - // retained only as provenance for the opt-in safety gate. - let has_refresh_token = tokens - .get("refresh_token") - .and_then(|v| v.as_str()) - .is_some_and(|s| !s.trim().is_empty()); - let last_refresh = json - .get("last_refresh") - .and_then(|v| v.as_str()) - .and_then(parse_timestamp); - - let access_token_expires_at = parse_access_token_expiry(&access_token); - - Ok(CodexCredentials { - access_token, - account_id, - is_external_oauth: has_refresh_token, - access_token_expires_at, - last_refresh, - }) - } - - /// Upstream 0.50.1 #2944: when `codex_external_oauth_sources_allowed` is - /// OFF (the default), external OAuth credential files without refresh - /// provenance fail closed instead of being used silently. An external - /// OAuth source is an auth.json `tokens` object with a `refresh_token` - /// (CLI-owned OAuth, not an API key). Win-CodexBar never refreshes or - /// writes this source: the gate only decides whether the read-only usage - /// request may use it. When the access token is a JWT, its native expiry - /// is the validity authority; opaque tokens are sent to the server. - fn enforce_external_oauth_gate(credentials: &CodexCredentials) -> Result<(), ProviderError> { - if !credentials.is_external_oauth { - return Ok(()); - } - // The opt-in only matters without refresh provenance. Skip the settings - // load otherwise: credential reads repeat while the owner publishes. - let external_sources_allowed = credentials.last_refresh.is_some() - || crate::settings::Settings::load().codex_external_oauth_sources_allowed; - Self::enforce_external_oauth_gate_at(credentials, external_sources_allowed, Utc::now()) - } - - fn enforce_external_oauth_gate_at( - credentials: &CodexCredentials, - external_sources_allowed: bool, - now: DateTime, - ) -> Result<(), ProviderError> { - if !credentials.is_external_oauth { - return Ok(()); - } - if !external_sources_allowed && credentials.last_refresh.is_none() { - return Err(ProviderError::AuthRequired); - } - if let Some(expires_at) = credentials.access_token_expires_at - && expires_at - now <= EXTERNAL_OAUTH_REFRESH_WINDOW - { - return Err(ProviderError::AuthRequired); - } - Ok(()) - } - - fn credential_cache() -> &'static Mutex> { - CREDENTIAL_CACHE.get_or_init(|| Mutex::new(None)) - } - - fn cached_credentials( - path: &std::path::Path, - modified: Option, - ) -> Option { - let guard = Self::credential_cache().lock().ok()?; - let cached = guard.as_ref()?; - if cached.path == path - && cached.modified == modified - && cached.loaded_at.elapsed() <= CREDENTIAL_CACHE_TTL - { - return Some(cached.credentials.clone()); - } - None - } - - fn store_cached_credentials( - path: PathBuf, - modified: Option, - credentials: CodexCredentials, - ) { - if let Ok(mut guard) = Self::credential_cache().lock() { - *guard = Some(CachedCodexCredentials { - path, - modified, - loaded_at: Instant::now(), - credentials, - }); - } - } - fn get_auth_path(&self) -> PathBuf { self.codex_dir().join("auth.json") } @@ -734,471 +330,6 @@ impl CodexApi { }; parse_chatgpt_base_url(&content).is_some() || config_uses_non_chatgpt_provider(&content) } - - fn build_result_from_json( - &self, - json: &serde_json::Value, - ) -> Result<(UsageSnapshot, Option), ProviderError> { - // Extract plan type - let plan_type = json - .get("plan_type") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - - // Extract rate limit info - handle multiple possible structures - let (primary, secondary, monthly, code_review, code_review_verified) = - self.extract_rate_limits(json); - - // Build login method string - let login_method = plan_type.as_deref().map(format_plan_type); - - let mut usage = UsageSnapshot::new(primary); - if let Some(sec) = secondary { - usage = usage.with_secondary(sec); - } - // F5 (upstream 0.48.0): monthly (30-day) windows go to tertiary so the - // bridge and frontend can show a monthly reset instead of swallowing it. - if let Some(mo) = monthly { - usage = usage.with_tertiary(mo); - } - if let Some(cr) = code_review { - usage = if code_review_verified { - usage.with_code_review(cr) - } else { - usage.with_model_specific(cr) - }; - } - for extra in self.extract_additional_rate_limits(json) { - usage.extra_rate_windows.push(extra); - } - if let Some(method) = login_method { - usage = usage.with_login_method(method); - } - - // Extract credits if present - let cost = self.extract_credits(json); - - Ok((usage, cost)) - } - - #[cfg(test)] - pub(crate) fn build_result_from_json_for_test( - &self, - json: &serde_json::Value, - ) -> Result<(UsageSnapshot, Option), ProviderError> { - self.build_result_from_json(json) - } - - fn extract_rate_limits( - &self, - json: &serde_json::Value, - ) -> ( - RateWindow, - Option, - Option, - Option, - bool, - ) { - // Try rate_limit object - if let Some(rate_limit) = json.get("rate_limit") { - let primary_opt = rate_limit - .get("primary_window") - .and_then(|w| self.parse_window_if_present(w)); - - let secondary_opt = rate_limit - .get("secondary_window") - .and_then(|w| self.parse_window_if_present(w)); - - let code_review = rate_limit - .get("code_review_window") - .and_then(|w| self.parse_window_if_present(w)); - - let (primary, secondary) = normalize_named_windows(primary_opt, secondary_opt); - - // F5 (upstream 0.48.0): named windows carry only session/weekly/code_review. - // Monthly is extracted separately (from array windows) — return None here. - let code_review_verified = code_review.is_some(); - return (primary, secondary, None, code_review, code_review_verified); - } - - // Try rate_limits array - if let Some(rate_limits) = json.get("rate_limits").and_then(|v| v.as_array()) { - let windows = rate_limits - .iter() - .filter_map(|window| self.parse_window_if_present(window)) - .collect::>(); - let (primary, secondary, monthly, code_review) = normalize_array_windows(windows); - // F5 (upstream 0.48.0): route monthly to its own tertiary lane. - let mut usage = UsageSnapshot::new(primary); - if let Some(sec) = secondary { - usage = usage.with_secondary(sec); - } - if let Some(mo) = monthly { - usage = usage.with_tertiary(mo); - } - if let Some(cr) = code_review { - usage = usage.with_model_specific(cr); - } - return ( - usage.primary, - usage.secondary, - usage.tertiary, - usage.model_specific, - false, - ); - } - - // Try direct fields - let used_percent = json - .get("used_percent") - .or_else(|| json.get("usage_percent")) - .and_then(json_f64); - let primary = RateWindow::new(used_percent.unwrap_or(0.0)) - .with_usage_known(valid_used_percent(used_percent)); - - (primary, None, None, None, false) - } - - fn parse_window(&self, window: &serde_json::Value) -> RateWindow { - let used_percent = window - .get("used_percent") - .or_else(|| window.get("usage_percent")) - .and_then(json_f64); - - let window_minutes = window - .get("limit_window_seconds") - .and_then(json_i64) - .and_then(|seconds| u32::try_from(seconds / 60).ok()); - - let reset_at = window - .get("reset_at") - .and_then(json_i64) - .and_then(|ts| Utc.timestamp_opt(ts, 0).single()); - - RateWindow::with_details( - used_percent.unwrap_or(0.0), - window_minutes, - reset_at, - format_reset_countdown(reset_at), - ) - .with_usage_known(valid_used_percent(used_percent)) - } - - fn parse_window_if_present(&self, window: &serde_json::Value) -> Option { - (!window.is_null() && !is_placeholder_window(window)).then(|| self.parse_window(window)) - } - - fn extract_additional_rate_limits(&self, json: &serde_json::Value) -> Vec { - json.get("additional_rate_limits") - .and_then(|v| v.as_array()) - .into_iter() - .flatten() - .filter_map(|entry| self.parse_additional_rate_limit(entry)) - .collect() - } - - fn parse_additional_rate_limit(&self, entry: &serde_json::Value) -> Option { - let metered_feature = entry - .get("metered_feature") - .and_then(|v| v.as_str()) - .map(str::trim) - .filter(|v| !v.is_empty()); - let limit_name = entry - .get("limit_name") - .and_then(|v| v.as_str()) - .map(str::trim) - .filter(|v| !v.is_empty()); - - let rate_limit = entry.get("rate_limit").unwrap_or(entry); - let primary = rate_limit.get("primary_window"); - let secondary = rate_limit.get("secondary_window"); - let window = primary.or(secondary)?; - if is_placeholder_window(window) { - return None; - } - - let parsed = self.parse_window(window); - let feature = metered_feature.unwrap_or_default(); - let limit = limit_name.unwrap_or_default(); - let is_spark = feature.eq_ignore_ascii_case("codex_spark") - || feature.eq_ignore_ascii_case("spark") - || limit.to_ascii_lowercase().contains("spark"); - - if is_spark { - let is_weekly = secondary.is_some() && primary.is_none() - || parsed - .window_minutes - .is_some_and(|mins| mins >= 7 * 24 * 60); - let (id, title) = if is_weekly { - ("codex-spark-weekly", "Codex Spark Weekly") - } else { - ("codex-spark", "Codex Spark 5-hour") - }; - return Some(NamedRateWindow::new(id, title, parsed)); - } - - let label = limit_name.or(metered_feature)?; - let slug = slugify(label); - if slug.is_empty() { - return None; - } - - Some(NamedRateWindow::new( - format!("codex-{slug}"), - titleize_limit_label(label), - parsed, - )) - } - - fn extract_credits(&self, json: &serde_json::Value) -> Option { - let credits = json.get("credits")?; - - let has_credits = credits - .get("has_credits") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - if !has_credits { - return None; - } - - let unlimited = credits - .get("unlimited") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - if unlimited { - return None; - } - - let balance = credits - .get("balance") - .and_then(|v| v.as_f64()) - .unwrap_or(0.0); - - Some(CostSnapshot::new(balance, "USD", "Credits")) - } - - fn build_result( - &self, - response: UsageResponse, - ) -> Result<(UsageSnapshot, Option), ProviderError> { - let (primary, secondary) = normalize_named_windows( - response - .rate_limit - .as_ref() - .and_then(|rate_limit| rate_limit.primary_window.as_ref()) - .map(rate_window_from_snapshot), - response - .rate_limit - .as_ref() - .and_then(|rate_limit| rate_limit.secondary_window.as_ref()) - .map(rate_window_from_snapshot), - ); - - // Extract code review rate window - let code_review = response - .rate_limit - .as_ref() - .and_then(|rate_limit| rate_limit.code_review_window.as_ref()) - .map(rate_window_from_snapshot); - - // Build usage snapshot - let login_method = response.plan_type.as_ref().map(|pt| match pt.as_str() { - "guest" => "Guest".to_string(), - "free" => "ChatGPT Free".to_string(), - "go" => "ChatGPT Go".to_string(), - "plus" => "ChatGPT Plus".to_string(), - "pro" => "ChatGPT Pro".to_string(), - "team" => "ChatGPT Team".to_string(), - "business" => "ChatGPT Business".to_string(), - "enterprise" => "ChatGPT Enterprise".to_string(), - "education" | "edu" => "ChatGPT Education".to_string(), - other => format!("ChatGPT {}", capitalize(other)), - }); - - let mut usage = UsageSnapshot::new(primary); - if let Some(sec) = secondary { - usage = usage.with_secondary(sec); - } - if let Some(cr) = code_review { - usage = usage.with_code_review(cr); - } - if let Some(method) = login_method { - usage = usage.with_login_method(method); - } - - // Build cost snapshot if credits are present - let credit_limit = response.individual_limit.as_ref().or_else(|| { - response - .rate_limit - .as_ref() - .and_then(|rate_limit| rate_limit.individual_limit.as_ref()) - }); - let cost = response.credits.as_ref().and_then(|credits| { - if credits.has_credits() { - let balance = credits.balance.unwrap_or(0.0); - if credits.unlimited() { - None // Unlimited credits, no need to show - } else if let Some(limit) = - credit_limit.and_then(|limit| limit.to_cost_snapshot(balance)) - { - Some(limit) - } else { - Some(CostSnapshot::new(balance, "USD", "Credits")) - } - } else { - None - } - }); - - Ok((usage, cost)) - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum CodexWindowRole { - Session, - Weekly, - Monthly, - Unknown, -} - -fn codex_window_role(window: &RateWindow) -> CodexWindowRole { - match window - .window_minutes - .map(RateWindowCadence::from_minutes) - .unwrap_or(RateWindowCadence::Unknown) - { - RateWindowCadence::Session => CodexWindowRole::Session, - RateWindowCadence::Monthly => CodexWindowRole::Monthly, - RateWindowCadence::Weekly => CodexWindowRole::Weekly, - RateWindowCadence::Unknown => CodexWindowRole::Unknown, - } -} - -/// Normalize the named `primary_window`/`secondary_window` fields by duration. -fn normalize_named_windows( - primary: Option, - secondary: Option, -) -> (RateWindow, Option) { - match (primary, secondary) { - (None, None) => (RateWindow::no_active_session(), None), - (Some(window), None) => { - if codex_window_role(&window) == CodexWindowRole::Weekly { - (RateWindow::no_active_session(), Some(window)) - } else { - (window, None) - } - } - (None, Some(window)) => { - if codex_window_role(&window) == CodexWindowRole::Weekly { - (RateWindow::no_active_session(), Some(window)) - } else { - (window, None) - } - } - (Some(primary), Some(secondary)) => { - match (codex_window_role(&primary), codex_window_role(&secondary)) { - (CodexWindowRole::Weekly, CodexWindowRole::Session) => (secondary, Some(primary)), - (CodexWindowRole::Weekly, CodexWindowRole::Unknown) => { - (RateWindow::no_active_session(), Some(primary)) - } - (CodexWindowRole::Unknown, CodexWindowRole::Session) => (secondary, Some(primary)), - (CodexWindowRole::Session, CodexWindowRole::Weekly) - | (CodexWindowRole::Unknown, CodexWindowRole::Weekly) => (primary, Some(secondary)), - _ => (primary, Some(secondary)), - } - } - } -} - -/// Normalize an array of Codex windows without relying on the API's ordering. -/// Normalize an array of Codex windows without relying on the API's ordering. -/// -/// Returns (session, weekly, monthly, code_review). F5 (upstream 0.48.0): -/// monthly (30-day) windows are routed to their own lane so surfaces can -/// display a monthly reset instead of swallowing it into the weekly label. -fn normalize_array_windows( - windows: Vec, -) -> ( - RateWindow, - Option, - Option, - Option, -) { - if windows.is_empty() { - return (RateWindow::no_active_session(), None, None, None); - } - - // Preserve the old positional fallback when the API provides no role - // metadata at all. There is no safe way to infer session vs weekly then. - if !windows - .iter() - .any(|window| codex_window_role(window) != CodexWindowRole::Unknown) - { - let mut windows = windows.into_iter(); - return ( - windows.next().unwrap_or_else(RateWindow::no_active_session), - windows.next(), - windows.next(), - windows.next(), - ); - } - - let mut session = None; - let mut weekly = None; - let mut monthly = None; - let mut remaining = Vec::new(); - - for window in windows { - match codex_window_role(&window) { - CodexWindowRole::Session if session.is_none() => session = Some(window), - CodexWindowRole::Weekly if weekly.is_none() => weekly = Some(window), - CodexWindowRole::Monthly if monthly.is_none() => monthly = Some(window), - _ => remaining.push(window), - } - } - - ( - session.unwrap_or_else(RateWindow::no_active_session), - weekly, - monthly, - remaining.into_iter().next(), - ) -} - -fn rate_window_from_snapshot(window: &WindowSnapshot) -> RateWindow { - let reset_at = timestamp_to_datetime(window.reset_at); - let used_percent = f64::from(window.used_percent); - RateWindow::with_details( - used_percent, - window - .limit_window_seconds - .and_then(|seconds| u32::try_from(seconds / 60).ok()), - reset_at, - format_reset_countdown(reset_at), - ) - .with_usage_known(valid_used_percent(Some(used_percent))) -} - -fn format_plan_type(plan_type: &str) -> String { - match plan_type { - "guest" => "Guest".to_string(), - "free" => "ChatGPT Free".to_string(), - "go" => "Codex Go".to_string(), - "plus" => "ChatGPT Plus".to_string(), - "pro" => "ChatGPT Pro".to_string(), - "pro_lite" | "prolite" | "pro-lite" => "Pro Lite".to_string(), - "team" => "ChatGPT Team".to_string(), - "business" => "ChatGPT Business".to_string(), - "enterprise" => "ChatGPT Enterprise".to_string(), - "education" | "edu" => "ChatGPT Education".to_string(), - "free_workspace" | "freeWorkspace" => "Free Workspace".to_string(), - "quorum" => "Codex Quorum".to_string(), - "k12" => "Codex K12".to_string(), - other => format!("ChatGPT {}", capitalize(other)), - } } impl Default for CodexApi { @@ -1207,316 +338,6 @@ impl Default for CodexApi { } } -// --- Data structures --- - -#[derive(Clone)] -struct CodexCredentials { - access_token: String, - account_id: Option, - /// True when the source is an external OAuth token set (has a - /// `refresh_token`), as opposed to an `OPENAI_API_KEY`. The Codex CLI owns - /// refresh and persistence for this source; this app only reads it. The - /// `codex_external_oauth_sources_allowed` setting gates that read - /// (upstream 0.50.1 #2944). - is_external_oauth: bool, - /// Native access-token JWT expiry. When available, this is authoritative - /// for validity; the Codex CLI still owns the refresh lifecycle. - access_token_expires_at: Option>, - /// `last_refresh` timestamp from auth.json, when present. Its presence - /// supplies provenance when the external-source opt-in setting is OFF; - /// its age is not an access-token expiry signal. - last_refresh: Option>, -} - -struct CachedCodexCredentials { - path: PathBuf, - modified: Option, - loaded_at: Instant, - credentials: CodexCredentials, -} - -#[derive(Debug, Deserialize)] -struct UsageResponse { - plan_type: Option, - rate_limit: Option, - credits: Option, - #[serde(default, alias = "individualLimit")] - individual_limit: Option, -} - -#[derive(Debug, Deserialize)] -struct RateLimitDetails { - primary_window: Option, - secondary_window: Option, - code_review_window: Option, - #[serde(default, alias = "individualLimit")] - individual_limit: Option, -} - -#[derive(Debug, Deserialize)] -struct WindowSnapshot { - used_percent: i32, - reset_at: Option, - limit_window_seconds: Option, -} - -#[derive(Debug, Deserialize)] -struct CreditDetails { - has_credits: Option, - unlimited: Option, - balance: Option, -} - -#[derive(Debug, Deserialize)] -struct SpendControlLimitSnapshot { - limit: Option, - used: Option, - #[serde(default, alias = "remainingPercent")] - remaining_percent: Option, - #[serde(default, alias = "resetsAt")] - resets_at: Option, -} - -#[derive(Debug, Clone, Deserialize)] -pub(super) struct ResetCredit { - #[serde(default)] - pub(super) id: Option, - #[serde(default, alias = "resetType")] - pub(super) reset_type: Option, - #[serde(default)] - pub(super) status: Option, - #[serde(default)] - pub(super) expires_at: Option, -} - -#[derive(Debug, Clone, Deserialize)] -pub(super) struct ResetCredits { - #[serde(default)] - pub(super) credits: Vec, - pub(super) available_count: u32, -} - -fn decode_reset_credits(data: &[u8]) -> Result { - serde_json::from_slice(data) - .map_err(|e| ProviderError::Parse(format!("Failed to parse Codex reset credits: {e}"))) -} - -fn parse_credit_expiry(raw: &str) -> Option> { - DateTime::parse_from_rfc3339(raw) - .ok() - .map(|dt| dt.with_timezone(&Utc)) -} - -fn is_available_credit(credit: &ResetCredit) -> bool { - match credit.status.as_deref() { - None | Some("") => true, - Some(status) => status.eq_ignore_ascii_case("available"), - } -} - -pub(super) fn next_available_reset_credit_expiry( - credits: &[ResetCredit], - now: DateTime, -) -> Option> { - credits - .iter() - .filter(|credit| is_available_credit(credit)) - .filter_map(|credit| credit.expires_at.as_deref().and_then(parse_credit_expiry)) - .filter(|expires_at| *expires_at > now) - .min() -} - -fn reset_credits_rate_window(reset: &ResetCredits, now: DateTime) -> RateWindow { - let description = format!( - "{} reset credit{} available", - reset.available_count, - if reset.available_count == 1 { "" } else { "s" } - ); - let mut window = RateWindow::informational(description); - window.resets_at = next_available_reset_credit_expiry(&reset.credits, now); - window -} - -fn apply_reset_credits_window( - mut usage: UsageSnapshot, - reset: Option<&ResetCredits>, -) -> UsageSnapshot { - usage - .extra_rate_windows - .retain(|window| window.id != "reset-credits"); - if let Some(reset) = reset.filter(|reset| reset.available_count > 0) { - let window = reset_credits_rate_window(reset, Utc::now()); - usage = usage.with_extra_rate_window("reset-credits", "Reset credits", window); - } - usage -} - -impl CreditDetails { - // Helper to safely check has_credits - fn has_credits(&self) -> bool { - self.has_credits.unwrap_or(false) - } - - fn unlimited(&self) -> bool { - self.unlimited.unwrap_or(false) - } -} - -impl SpendControlLimitSnapshot { - fn to_cost_snapshot(&self, balance: f64) -> Option { - let limit = self - .limit - .filter(|limit| limit.is_finite() && *limit >= 0.0)?; - let used = self - .used - .filter(|used| used.is_finite() && *used >= 0.0) - .or_else(|| { - self.remaining_percent - .filter(|pct| pct.is_finite() && *pct >= 0.0) - .map(|remaining| limit * (1.0 - (remaining / 100.0))) - }) - .unwrap_or_else(|| (limit - balance).max(0.0)); - let mut cost = - CostSnapshot::new(used.clamp(0.0, limit), "USD", "Monthly credits").with_limit(limit); - if let Some(resets_at) = timestamp_to_datetime(self.resets_at) { - cost = cost.with_resets_at(resets_at); - } - Some(cost) - } -} - -// --- Helper functions --- - -fn timestamp_to_datetime(timestamp: Option) -> Option> { - timestamp.and_then(|ts| Utc.timestamp_opt(ts, 0).single()) -} - -/// Parse the native `exp` claim from an access-token JWT. Opaque or malformed -/// tokens return `None` and are handled by the read-only usage request. -fn parse_access_token_expiry(token: &str) -> Option> { - let payload = token.split('.').nth(1)?; - let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(payload) - .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload)) - .ok()?; - let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; - let exp = json.get("exp")?.as_i64()?; - Utc.timestamp_opt(exp, 0).single() -} -fn parse_timestamp(raw: &str) -> Option> { - let trimmed = raw.trim(); - if trimmed.is_empty() { - return None; - } - DateTime::parse_from_rfc3339(trimmed) - .ok() - .map(|dt| dt.with_timezone(&Utc)) - .or_else(|| { - chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%dT%H:%M:%S%.f") - .ok() - .map(|naive| DateTime::::from_naive_utc_and_offset(naive, Utc)) - }) -} - -fn json_f64(value: &serde_json::Value) -> Option { - value - .as_f64() - .or_else(|| value.as_i64().map(|value| value as f64)) - .or_else(|| value.as_str()?.trim().parse::().ok()) -} - -fn json_i64(value: &serde_json::Value) -> Option { - value - .as_i64() - .or_else(|| value.as_str()?.trim().parse::().ok()) -} - -fn valid_used_percent(value: Option) -> bool { - value.is_some_and(|value| value.is_finite() && (0.0..=100.0).contains(&value)) -} - -fn is_placeholder_window(window: &serde_json::Value) -> bool { - let has_usage = window - .get("used_percent") - .or_else(|| window.get("usage_percent")) - .and_then(json_f64) - .is_some(); - let has_duration = window - .get("limit_window_seconds") - .and_then(json_i64) - .is_some(); - let has_reset = window.get("reset_at").and_then(json_i64).is_some(); - - !has_usage && !has_duration && !has_reset -} - -fn slugify(label: &str) -> String { - let mut slug = String::new(); - let mut previous_dash = false; - - for ch in label.chars() { - if ch.is_ascii_alphanumeric() { - slug.push(ch.to_ascii_lowercase()); - previous_dash = false; - } else if !previous_dash && !slug.is_empty() { - slug.push('-'); - previous_dash = true; - } - } - - while slug.ends_with('-') { - slug.pop(); - } - slug -} - -fn titleize_limit_label(label: &str) -> String { - label - .split(['_', '-', ' ']) - .filter(|part| !part.is_empty()) - .map(|part| { - let mut chars = part.chars(); - match chars.next() { - Some(first) => first - .to_uppercase() - .chain(chars.flat_map(char::to_lowercase)) - .collect(), - None => String::new(), - } - }) - .collect::>() - .join(" ") -} - -fn format_reset_countdown(reset_at: Option>) -> Option { - let dt = reset_at?; - let now = Utc::now(); - if dt <= now { - return Some("now".to_string()); - } - let diff = dt - now; - let total_mins = diff.num_minutes(); - let hours = diff.num_hours(); - let mins = total_mins % 60; - if hours >= 24 { - let days = hours / 24; - let rem_h = hours % 24; - if rem_h == 0 { - Some(format!("{}d", days)) - } else { - Some(format!("{}d {}h", days, rem_h)) - } - } else if hours > 0 { - if mins == 0 { - Some(format!("{}h", hours)) - } else { - Some(format!("{}h {}m", hours, mins)) - } - } else { - Some(format!("{}m", mins)) - } -} - /// Whether config.toml selects a non-ChatGPT model provider (e.g. Bedrock), /// meaning the CLI never authenticates against ChatGPT. fn config_uses_non_chatgpt_provider(config_content: &str) -> bool { @@ -1580,1044 +401,8 @@ fn normalize_base_url(url: &str) -> String { trimmed } -fn capitalize(s: &str) -> String { - let mut chars = s.chars(); - match chars.next() { - None => String::new(), - Some(first) => first.to_uppercase().chain(chars).collect(), - } -} - #[cfg(test)] mod credential_retry_tests; #[cfg(test)] -mod tests { - use super::*; - use serde_json::json; - - #[test] - fn non_chatgpt_model_provider_is_detected_for_guidance() { - // Upstream 0.50.0 #2679: Bedrock and other custom backends get - // rate-limit guidance instead of login instructions. - assert!(config_uses_non_chatgpt_provider( - "model_provider = \"bedrock\"\n" - )); - assert!(config_uses_non_chatgpt_provider( - "# relay\nmodel_provider = 'ollama'" - )); - assert!(!config_uses_non_chatgpt_provider( - "model_provider = \"openai\"" - )); - assert!(!config_uses_non_chatgpt_provider( - "model = \"gpt-5\"\napproval_policy = \"never\"" - )); - } - - #[test] - fn parses_codex_credentials_without_retaining_refresh_token() { - let credentials = CodexApi::parse_credentials_json( - r#"{ - "tokens": { - "access_token": "access", - "refresh_token": "refresh", - "account_id": "acct_123" - } - }"#, - ) - .expect("credentials"); - - assert_eq!(credentials.access_token, "access"); - assert_eq!(credentials.account_id.as_deref(), Some("acct_123")); - } - - #[test] - fn decodes_reset_credits() { - let credits = decode_reset_credits( - br#"{"available_count":2,"credits":[{"id":"a","status":"available","expires_at":"2026-08-01T12:00:00Z"}]}"#, - ) - .expect("reset credits"); - assert_eq!(credits.available_count, 2); - assert_eq!(credits.credits.len(), 1); - assert_eq!(credits.credits[0].status.as_deref(), Some("available")); - assert_eq!( - credits.credits[0].expires_at.as_deref(), - Some("2026-08-01T12:00:00Z") - ); - } - - #[test] - fn missing_reset_credit_count_is_unavailable_not_zero() { - assert!(decode_reset_credits(br#"{"credits":[]}"#).is_err()); - } - - #[test] - fn next_expiry_picks_soonest_available() { - let now = DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let credits = vec![ - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-07-10T00:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-07-05T00:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-07-20T00:00:00Z".into()), - }, - ]; - let expiry = next_available_reset_credit_expiry(&credits, now).expect("expiry"); - assert_eq!( - expiry, - DateTime::parse_from_rfc3339("2026-07-05T00:00:00Z") - .unwrap() - .with_timezone(&Utc) - ); - } - - #[test] - fn next_expiry_skips_past_and_non_available() { - let now = DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let credits = vec![ - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-06-01T00:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: Some("used".into()), - expires_at: Some("2026-07-03T00:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: Some("AVAILABLE".into()), - expires_at: Some("2026-07-08T00:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: None, - expires_at: Some("2026-07-09T00:00:00Z".into()), - }, - ]; - let expiry = next_available_reset_credit_expiry(&credits, now).expect("expiry"); - assert_eq!( - expiry, - DateTime::parse_from_rfc3339("2026-07-08T00:00:00Z") - .unwrap() - .with_timezone(&Utc) - ); - } - - #[test] - fn reset_credits_window_sets_informational_and_expiry() { - let now = DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let reset = ResetCredits { - available_count: 2, - credits: vec![ - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-07-15T12:00:00Z".into()), - }, - ResetCredit { - id: None, - reset_type: None, - status: Some("available".into()), - expires_at: Some("2026-07-10T12:00:00Z".into()), - }, - ], - }; - let window = reset_credits_rate_window(&reset, now); - assert!(window.is_informational); - assert_eq!( - window.reset_description.as_deref(), - Some("2 reset credits available") - ); - assert_eq!( - window.resets_at, - Some( - DateTime::parse_from_rfc3339("2026-07-10T12:00:00Z") - .unwrap() - .with_timezone(&Utc) - ) - ); - } - - #[test] - fn reset_credits_window_count_only_without_expiry() { - let now = DateTime::parse_from_rfc3339("2026-07-01T00:00:00Z") - .unwrap() - .with_timezone(&Utc); - let reset = ResetCredits { - available_count: 1, - credits: vec![], - }; - let window = reset_credits_rate_window(&reset, now); - assert!(window.is_informational); - assert_eq!( - window.reset_description.as_deref(), - Some("1 reset credit available") - ); - assert!(window.resets_at.is_none()); - } - - fn write_codex_home(base_url: &str) -> tempfile::TempDir { - let dir = tempfile::tempdir().expect("temp codex home"); - std::fs::write( - dir.path().join("auth.json"), - r#"{"tokens":{"access_token":"test-token","account_id":"acct_test"}}"#, - ) - .expect("auth.json"); - std::fs::write( - dir.path().join("config.toml"), - format!("chatgpt_base_url = \"{base_url}\""), - ) - .expect("config.toml"); - dir - } - - #[tokio::test] - async fn reset_credit_cache_single_flight_and_unknown_are_ten_minute_observations() { - let mut server = mockito::Server::new_async().await; - let request = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(503) - .create_async() - .await; - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let creds = api.load_credentials().await.unwrap(); - let base = server.url(); - let (first, second) = tokio::join!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base), - api.fetch_rate_limit_reset_credits_cached(&creds, &base), - ); - assert!(first.is_none() && second.is_none()); - assert!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .is_none() - ); - request.assert_async().await; - assert!(RESET_CREDITS_CACHE_TTL == Duration::from_secs(600)); - } - - #[tokio::test] - async fn reset_credit_cache_expires_and_token_rotation_uses_new_scope() { - let mut server = mockito::Server::new_async().await; - let first = server - .mock("GET", "/wham/rate-limit-reset-credits") - .match_header("authorization", "Bearer test-token") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":2,"credits":[]}"#) - .create_async() - .await; - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let mut creds = api.load_credentials().await.unwrap(); - let base = server.url(); - assert_eq!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .unwrap() - .available_count, - 2 - ); - assert_eq!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .unwrap() - .available_count, - 2 - ); - first.assert_async().await; - first.remove_async().await; - let second = server - .mock("GET", "/wham/rate-limit-reset-credits") - .match_header("authorization", "Bearer test-token") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":1,"credits":[]}"#) - .create_async() - .await; - api.reset_credits_cache_slot(&creds, &base) - .lock() - .await - .loaded_at = Some(Instant::now() - RESET_CREDITS_CACHE_TTL); - assert_eq!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .unwrap() - .available_count, - 1 - ); - second.assert_async().await; - second.remove_async().await; - creds.access_token = "rotated-token".into(); - let rotated = server - .mock("GET", "/wham/rate-limit-reset-credits") - .match_header("authorization", "Bearer rotated-token") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":3,"credits":[]}"#) - .create_async() - .await; - assert_eq!( - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .unwrap() - .available_count, - 3 - ); - rotated.assert_async().await; - } - - #[tokio::test] - async fn reset_credit_cache_is_scoped_to_the_codex_home() { - // Same base URL, account and token, different Codex homes: each home - // makes its own observation instead of reading the other's. - let mut server = mockito::Server::new_async().await; - let base = server.url(); - let first = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":2,"credits":[]}"#) - .create_async() - .await; - let first_home = write_codex_home(&base); - let first_api = CodexApi::new().with_codex_home(first_home.path()); - let first_creds = first_api.load_credentials().await.unwrap(); - let first_count = first_api - .fetch_rate_limit_reset_credits_cached(&first_creds, &base) - .await - .map(|credits| credits.available_count); - assert_eq!(first_count, Some(2)); - first.assert_async().await; - first.remove_async().await; - - let second = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(503) - .create_async() - .await; - let second_home = write_codex_home(&base); - let second_api = CodexApi::new().with_codex_home(second_home.path()); - let second_creds = second_api.load_credentials().await.unwrap(); - assert_eq!(first_creds.access_token, second_creds.access_token); - assert_eq!(first_creds.account_id, second_creds.account_id); - assert!( - second_api - .fetch_rate_limit_reset_credits_cached(&second_creds, &base) - .await - .is_none() - ); - second.assert_async().await; - } - - #[tokio::test] - async fn suspicious_weekly_reset_uses_independent_credit_observations() { - let mut server = mockito::Server::new_async().await; - let cached_response = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":2,"credits":[]}"#) - .create_async() - .await; - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let creds = api.load_credentials().await.unwrap(); - let base = server.url(); - api.fetch_rate_limit_reset_credits_cached(&creds, &base) - .await - .unwrap(); - cached_response.assert_async().await; - cached_response.remove_async().await; - - let started = Instant::now(); - let initial_response = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":1,"credits":[]}"#) - .create_async() - .await; - let initial = api - .fresh_reset_credits_for_confirmation(&creds, &base, started) - .await - .unwrap(); - assert_eq!(initial.available_count, 1); - initial_response.assert_async().await; - initial_response.remove_async().await; - - let confirmation_response = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":0,"credits":[]}"#) - .create_async() - .await; - let confirmation = api - .fetch_rate_limit_reset_credits_fresh(&creds, &base) - .await - .unwrap(); - assert_eq!(confirmation.available_count, 0); - confirmation_response.assert_async().await; - } - - #[tokio::test] - async fn pending_delayed_candidate_revalidates_with_a_fresh_credit_observation() { - let mut server = mockito::Server::new_async().await; - let candidate_observation = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":2,"credits":[]}"#) - .create_async() - .await; - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let creds = api.load_credentials().await.unwrap(); - let base = server.url(); - let cached = api - .fetch_rate_limit_reset_credits_cached(&creds, &base) - .await; - candidate_observation.assert_async().await; - candidate_observation.remove_async().await; - - // A later refresh: the ten-minute cache still holds the observation - // that created the candidate. - let started = Instant::now(); - let changed = server - .mock("GET", "/wham/rate-limit-reset-credits") - .expect(1) - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":1,"credits":[]}"#) - .create_async() - .await; - let without_candidate = api - .initial_reset_credits(&creds, &base, started, cached.clone(), false) - .await; - assert_eq!( - without_candidate.map(|credits| credits.available_count), - Some(2) - ); - let with_candidate = api - .initial_reset_credits(&creds, &base, started, cached, true) - .await; - assert_eq!( - with_candidate.map(|credits| credits.available_count), - Some(1) - ); - changed.assert_async().await; - } - - #[tokio::test] - async fn fetch_usage_attaches_reset_credits_from_http() { - let mut server = mockito::Server::new_async().await; - let soonest = (Utc::now() + chrono::Duration::days(5)).to_rfc3339(); - let later = (Utc::now() + chrono::Duration::days(12)).to_rfc3339(); - - let usage_mock = server - .mock("GET", "/wham/usage") - .match_header("authorization", "Bearer test-token") - .with_status(200) - .with_header("content-type", "application/json") - .with_body( - r#"{"plan_type":"plus","rate_limit":{"primary_window":{"used_percent":10,"limit_window_seconds":18000}}}"#, - ) - .create_async() - .await; - - let reset_body = format!( - r#"{{"available_count":2,"credits":[ - {{"status":"available","expires_at":"{later}"}}, - {{"status":"available","expires_at":"{soonest}"}} - ]}}"# - ); - let reset_mock = server - .mock("GET", "/wham/rate-limit-reset-credits") - .match_header("authorization", "Bearer test-token") - .with_status(200) - .with_header("content-type", "application/json") - .with_body(reset_body) - .create_async() - .await; - - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let (usage, _, _) = api.fetch_usage().await.expect("fetch_usage"); - - usage_mock.assert_async().await; - reset_mock.assert_async().await; - - let extra = usage - .extra_rate_windows - .iter() - .find(|w| w.id == "reset-credits") - .expect("reset-credits window attached"); - assert_eq!(extra.title, "Reset credits"); - assert!(extra.window.is_informational); - assert_eq!( - extra.window.reset_description.as_deref(), - Some("2 reset credits available") - ); - let expected = DateTime::parse_from_rfc3339(&soonest) - .unwrap() - .with_timezone(&Utc); - assert_eq!(extra.window.resets_at, Some(expected)); - } - - #[tokio::test] - async fn authenticated_codex_http_distinguishes_401_from_403() { - for (status, expects_authentication) in [(401, true), (403, false)] { - let mut server = mockito::Server::new_async().await; - let mock = server - .mock("GET", "/wham/usage") - .with_status(status) - .with_body("fixture refusal") - .create_async() - .await; - - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let error = match api.fetch_usage().await { - Ok(_) => panic!("expected HTTP {status} to fail"), - Err(error) => error, - }; - - if expects_authentication { - assert!(matches!(error, ProviderError::AuthRequired)); - } else { - let message = error.to_string(); - assert!(message.contains("403")); - assert!(message.contains("fixture refusal")); - assert!(!matches!(error, ProviderError::AuthRequired)); - } - mock.assert_async().await; - } - } - - #[tokio::test] - async fn old_opaque_external_oauth_reaches_usage_request() { - let mut server = mockito::Server::new_async().await; - let usage_mock = server - .mock("GET", "/wham/usage") - .match_header("authorization", "Bearer opaque-token") - .with_status(200) - .with_header("content-type", "application/json") - .with_body( - r#"{"plan_type":"plus","rate_limit":{"primary_window":{"used_percent":10,"limit_window_seconds":18000}}}"#, - ) - .create_async() - .await; - let reset_mock = server - .mock("GET", "/wham/rate-limit-reset-credits") - .match_header("authorization", "Bearer opaque-token") - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":0,"credits":[]}"#) - .create_async() - .await; - - let creds = CodexApi::parse_credentials_json( - r#"{ - "tokens": { - "access_token": "opaque-token", - "refresh_token": "refresh", - "account_id": "acct_test" - }, - "last_refresh": "2026-01-01T00:00:00Z" - }"#, - ) - .expect("credentials"); - assert!(CodexApi::enforce_external_oauth_gate_at(&creds, false, Utc::now()).is_ok()); - - let api = CodexApi::new(); - let (usage, _, _) = api - .fetch_usage_once(&creds, &server.url()) - .await - .expect("opaque OAuth usage request"); - assert_eq!(usage.primary.used_percent, 10.0); - usage_mock.assert_async().await; - reset_mock.assert_async().await; - } - - #[tokio::test] - async fn fetch_usage_skips_reset_credits_when_available_count_zero() { - let mut server = mockito::Server::new_async().await; - - let usage_mock = server - .mock("GET", "/wham/usage") - .with_status(200) - .with_header("content-type", "application/json") - .with_body( - r#"{"plan_type":"plus","rate_limit":{"primary_window":{"used_percent":10,"limit_window_seconds":18000}}}"#, - ) - .create_async() - .await; - - let reset_mock = server - .mock("GET", "/wham/rate-limit-reset-credits") - .with_status(200) - .with_header("content-type", "application/json") - .with_body(r#"{"available_count":0,"credits":[]}"#) - .create_async() - .await; - - let home = write_codex_home(&server.url()); - let api = CodexApi::new().with_codex_home(home.path()); - let (usage, _, _) = api.fetch_usage().await.expect("fetch_usage"); - - usage_mock.assert_async().await; - reset_mock.assert_async().await; - - assert!( - usage - .extra_rate_windows - .iter() - .all(|w| w.id != "reset-credits"), - "available_count=0 must not attach reset-credits" - ); - } - - #[test] - fn keeps_weekly_window_in_secondary_when_session_is_absent() { - let api = CodexApi::new(); - let (usage, _) = api - .build_result_from_json(&json!({ - "rate_limit": { - "secondary_window": { - "used_percent": 25, - "limit_window_seconds": 604800, - "reset_at": 1783036800 - } - } - })) - .expect("codex usage"); - - assert!(usage.primary.is_informational); - assert_eq!(usage.primary.window_minutes, Some(300)); - assert_eq!( - usage.primary.reset_description.as_deref(), - Some("No active 5h session") - ); - - let weekly = usage.secondary.expect("weekly window"); - assert!(!weekly.is_informational); - assert_eq!(weekly.used_percent, 25.0); - assert_eq!(weekly.window_minutes, Some(10080)); - } - - #[test] - fn identifies_rate_limit_array_windows_by_duration() { - let api = CodexApi::new(); - let (usage, _) = api - .build_result_from_json(&json!({ - "rate_limits": [ - { - "used_percent": 25, - "limit_window_seconds": 604800, - "reset_at": 1783036800 - }, - { - "used_percent": 10, - "limit_window_seconds": 18000, - "reset_at": 1783018800 - } - ] - })) - .expect("codex usage"); - - assert!(!usage.primary.is_informational); - assert_eq!(usage.primary.used_percent, 10.0); - assert_eq!(usage.primary.window_minutes, Some(300)); - - let weekly = usage.secondary.expect("weekly window"); - assert_eq!(weekly.used_percent, 25.0); - assert_eq!(weekly.window_minutes, Some(10080)); - } - - #[test] - fn identifies_weekly_only_rate_limit_array_without_a_session() { - let api = CodexApi::new(); - let (usage, _) = api - .build_result_from_json(&json!({ - "rate_limits": [{ - "used_percent": 25, - "limit_window_seconds": 604800, - "reset_at": 1783036800 - }] - })) - .expect("codex usage"); - - assert!(usage.primary.is_informational); - assert_eq!(usage.secondary.expect("weekly window").used_percent, 25.0); - } - - #[test] - fn maps_codex_spark_additional_rate_limits() { - let api = CodexApi::new(); - let (usage, _) = api - .build_result_from_json(&json!({ - "plan_type": "pro", - "rate_limit": { - "primary_window": { "used_percent": 20, "limit_window_seconds": 18000 }, - "secondary_window": { "used_percent": 40, "limit_window_seconds": 604800 } - }, - "additional_rate_limits": [ - { - "limit_name": "Codex Spark", - "metered_feature": "codex_spark", - "rate_limit": { - "primary_window": { "used_percent": "17", "limit_window_seconds": 18000 } - } - }, - { - "limit_name": "Codex Spark Weekly", - "metered_feature": "codex_spark", - "rate_limit": { - "secondary_window": { "used_percent": 62, "limit_window_seconds": 604800 } - } - } - ] - })) - .expect("codex usage"); - - assert_eq!(usage.extra_rate_windows.len(), 2); - assert_eq!(usage.extra_rate_windows[0].id, "codex-spark"); - assert_eq!(usage.extra_rate_windows[0].title, "Codex Spark 5-hour"); - assert_eq!(usage.extra_rate_windows[0].window.used_percent, 17.0); - assert_eq!(usage.extra_rate_windows[1].id, "codex-spark-weekly"); - assert_eq!(usage.extra_rate_windows[1].title, "Codex Spark Weekly"); - assert_eq!(usage.extra_rate_windows[1].window.used_percent, 62.0); - } - - #[test] - fn ignores_placeholder_additional_rate_limits() { - let api = CodexApi::new(); - let (usage, _) = api - .build_result_from_json(&json!({ - "rate_limit": { - "primary_window": { "used_percent": 0, "limit_window_seconds": 18000 } - }, - "additional_rate_limits": [ - { - "limit_name": "placeholder", - "metered_feature": "placeholder", - "rate_limit": { "primary_window": {} } - } - ] - })) - .expect("codex usage"); - - assert!(usage.extra_rate_windows.is_empty()); - } - - #[test] - fn maps_top_level_individual_credit_limit_to_cost_snapshot() { - let api = CodexApi::new(); - let (_, cost) = api - .build_result(UsageResponse { - plan_type: None, - rate_limit: None, - credits: Some(CreditDetails { - has_credits: Some(true), - unlimited: Some(false), - balance: Some(7.5), - }), - individual_limit: Some(SpendControlLimitSnapshot { - limit: Some(20.0), - used: Some(12.5), - remaining_percent: None, - resets_at: Some(1783036800), - }), - }) - .expect("codex result"); - let cost = cost.expect("cost"); - assert_eq!(cost.used, 12.5); - assert_eq!(cost.limit, Some(20.0)); - assert!(cost.resets_at.is_some()); - } - - #[test] - fn maps_nested_individual_credit_limit_to_cost_snapshot() { - let api = CodexApi::new(); - let (_, cost) = api - .build_result(UsageResponse { - plan_type: None, - rate_limit: Some(RateLimitDetails { - primary_window: None, - secondary_window: None, - code_review_window: None, - individual_limit: Some(SpendControlLimitSnapshot { - limit: Some(100.0), - used: None, - remaining_percent: Some(60.0), - resets_at: None, - }), - }), - credits: Some(CreditDetails { - has_credits: Some(true), - unlimited: Some(false), - balance: Some(60.0), - }), - individual_limit: None, - }) - .expect("codex result"); - let cost = cost.expect("cost"); - assert_eq!(cost.used, 40.0); - assert_eq!(cost.limit, Some(100.0)); - } - - fn win(minutes: u32, used: f64) -> RateWindow { - RateWindow::with_details(used, Some(minutes), None, None) - } - - #[test] - fn f5_normalize_array_routes_session_weekly_monthly_to_lanes() { - // 5h session + weekly + monthly → (session, weekly, monthly, None) - let windows = vec![win(300, 10.0), win(10_080, 20.0), win(43_200, 30.0)]; - let (primary, secondary, tertiary, code_review) = normalize_array_windows(windows); - assert_eq!(primary.window_minutes, Some(300)); - assert_eq!(secondary.unwrap().window_minutes, Some(10_080)); - assert_eq!(tertiary.unwrap().window_minutes, Some(43_200)); - assert!(code_review.is_none()); - } - - #[test] - fn f5_normalize_array_monthly_routes_to_tertiary_not_secondary() { - // Monthly must go to tertiary, NOT secondary — so #268's weekly math - // and "Weekly" label stay untouched. - let windows = vec![win(43_200, 50.0), win(10_080, 20.0)]; - let (primary, secondary, tertiary, _) = normalize_array_windows(windows); - assert_eq!(primary.window_minutes, Some(300)); // no session → placeholder - assert_eq!(secondary.unwrap().window_minutes, Some(10_080)); - assert_eq!(tertiary.unwrap().window_minutes, Some(43_200)); - } - - #[test] - fn f5_normalize_array_empty_returns_placeholder_primary() { - let (primary, secondary, tertiary, code_review) = normalize_array_windows(vec![]); - assert!(primary.is_informational); - assert!(secondary.is_none()); - assert!(tertiary.is_none()); - assert!(code_review.is_none()); - } - - #[test] - fn f5_normalize_array_unknown_windows_fall_to_code_review() { - // Windows with unrecognized durations (not 300/10080/43200) go to the - // remaining/code_review bucket. - let windows = vec![win(300, 10.0), win(999, 5.0)]; - let (primary, secondary, tertiary, code_review) = normalize_array_windows(windows); - assert_eq!(primary.window_minutes, Some(300)); - assert!(secondary.is_none()); - assert!(tertiary.is_none()); - assert_eq!(code_review.unwrap().window_minutes, Some(999)); - } - - // ── Upstream 0.50.1 #2944: external OAuth source gate ────────────────── - - #[test] - fn confirmation_failure_fallback_keeps_first_successful_usage_and_cost() { - let state = weekly_reset::AccountState::default(); - let first = UsageSnapshot::new(RateWindow::new(10.0)).with_secondary(RateWindow::new(0.5)); - let cost = Some(CostSnapshot::new(3.25, "USD", "Monthly")); - let (usage, kept_cost) = CodexApi::preserve_after_confirmation_failure(&state, first, cost); - assert!((usage.secondary.expect("weekly").used_percent - 0.5).abs() < f64::EPSILON); - assert_eq!(kept_cost.expect("cost").used, 3.25); - } - #[test] - fn api_key_credentials_are_not_external_oauth() { - let creds = CodexApi::parse_credentials_json(r#"{"OPENAI_API_KEY": "sk-test"}"#) - .expect("credentials"); - assert!(!creds.is_external_oauth); - assert!(creds.access_token_expires_at.is_none()); - assert!(creds.last_refresh.is_none()); - assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); - } - - #[test] - fn oauth_tokens_with_refresh_token_are_external_source() { - let creds = CodexApi::parse_credentials_json( - r#"{ - "tokens": { - "access_token": "access", - "refresh_token": "refresh", - "account_id": "acct_123" - } - }"#, - ) - .expect("credentials"); - assert!(creds.is_external_oauth); - assert!(creds.access_token_expires_at.is_none()); - assert!(creds.last_refresh.is_none()); - } - - #[test] - fn oauth_tokens_without_refresh_token_are_not_external() { - let creds = CodexApi::parse_credentials_json( - r#"{ - "tokens": { - "access_token": "access", - "account_id": "acct_123" - } - }"#, - ) - .expect("credentials"); - assert!(!creds.is_external_oauth); - } - - #[test] - fn external_oauth_gate_fails_closed_without_last_refresh() { - let creds = CodexCredentials { - access_token: "access".to_string(), - account_id: None, - is_external_oauth: true, - access_token_expires_at: None, - last_refresh: None, - }; - let err = CodexApi::enforce_external_oauth_gate(&creds) - .expect_err("external OAuth without provenance must fail closed"); - assert!(matches!(err, ProviderError::AuthRequired)); - } - - #[test] - fn external_oauth_gate_ignores_old_last_refresh_for_opaque_token() { - let old = Utc::now() - chrono::Duration::days(10); - let creds = CodexCredentials { - access_token: "access".to_string(), - account_id: None, - is_external_oauth: true, - access_token_expires_at: None, - last_refresh: Some(old), - }; - assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); - } - - #[test] - fn external_oauth_gate_allows_refresh_provenance() { - let fresh = Utc::now() - chrono::Duration::hours(1); - let creds = CodexCredentials { - access_token: "access".to_string(), - account_id: None, - is_external_oauth: true, - access_token_expires_at: None, - last_refresh: Some(fresh), - }; - assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); - } - - #[test] - fn external_oauth_gate_uses_future_jwt_expiry_over_old_last_refresh() { - let now = Utc::now(); - let future = now + chrono::Duration::hours(2); - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(format!(r#"{{"exp":{}}}"#, future.timestamp())); - let token = format!("header.{payload}.signature"); - let json = format!( - r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"2026-01-01T00:00:00Z"}}"# - ); - let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); - assert!(creds.access_token_expires_at.is_some()); - assert!(CodexApi::enforce_external_oauth_gate_at(&creds, false, now).is_ok()); - assert!(CodexApi::enforce_external_oauth_gate_at(&creds, true, now).is_ok()); - } - - #[test] - fn external_oauth_gate_rejects_expired_jwt() { - let expired = Utc::now() - chrono::Duration::minutes(1); - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(format!(r#"{{"exp":{}}}"#, expired.timestamp())); - let token = format!("header.{payload}.signature"); - let json = format!( - r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"{}"}}"#, - Utc::now().to_rfc3339() - ); - let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); - let err = CodexApi::enforce_external_oauth_gate(&creds) - .expect_err("expired native OAuth must be rejected"); - assert!(matches!(err, ProviderError::AuthRequired)); - } - - #[test] - fn external_oauth_gate_requires_cli_refresh_when_jwt_is_near_expiry() { - let soon = Utc::now() + chrono::Duration::minutes(2); - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(format!(r#"{{"exp":{}}}"#, soon.timestamp())); - let token = format!("header.{payload}.signature"); - let json = format!( - r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"{}"}}"#, - Utc::now().to_rfc3339() - ); - let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); - let err = CodexApi::enforce_external_oauth_gate(&creds) - .expect_err("near-expiry native OAuth must refresh through the CLI"); - assert!(matches!(err, ProviderError::AuthRequired)); - } - - #[test] - fn external_oauth_gate_allows_missing_last_refresh_when_opted_in() { - let creds = CodexCredentials { - access_token: "opaque-token".to_string(), - account_id: None, - is_external_oauth: true, - access_token_expires_at: None, - last_refresh: None, - }; - assert!(CodexApi::enforce_external_oauth_gate_at(&creds, true, Utc::now()).is_ok()); - } - - #[test] - fn opaque_token_uses_refresh_provenance_when_no_jwt_expiry_exists() { - let fresh = Utc::now().to_rfc3339(); - let json = format!( - r#"{{"tokens":{{"access_token":"opaque-token","refresh_token":"refresh"}},"last_refresh":"{fresh}"}}"# - ); - let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); - assert!(creds.access_token_expires_at.is_none()); - assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); - } - #[test] - fn parse_timestamp_reads_iso8601() { - assert!(parse_timestamp("2026-08-17T10:00:00Z").is_some()); - assert!(parse_timestamp("2026-08-17T10:00:00.123Z").is_some()); - assert!(parse_timestamp(" 2026-08-17T10:00:00Z ").is_some()); - assert!(parse_timestamp("").is_none()); - assert!(parse_timestamp("not-a-date").is_none()); - } -} +mod tests; diff --git a/rust/src/providers/codex/api/credential_retry_tests.rs b/rust/src/providers/codex/api/credential_retry_tests.rs index 7d9aca1163..44068ca182 100644 --- a/rust/src/providers/codex/api/credential_retry_tests.rs +++ b/rust/src/providers/codex/api/credential_retry_tests.rs @@ -6,7 +6,9 @@ //! injected reader does. The clock is paused, so the 50 ms retry delays are //! exact and the tests never race a writer thread. +use super::credentials::CREDENTIAL_READ_RETRY_DELAY; use super::*; +use base64::Engine; use std::cell::Cell; use std::path::Path; use std::time::Duration; diff --git a/rust/src/providers/codex/api/credentials.rs b/rust/src/providers/codex/api/credentials.rs new file mode 100644 index 0000000000..0b9bab8316 --- /dev/null +++ b/rust/src/providers/codex/api/credentials.rs @@ -0,0 +1,287 @@ +use super::CodexApi; +use crate::core::ProviderError; +use base64::Engine; +use chrono::{DateTime, TimeZone, Utc}; +use std::path::PathBuf; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, Instant, SystemTime}; + +const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(5); +/// Upstream 0.69.0 #4088: the Codex CLI owns `auth.json` and may be publishing a +/// replacement while we read it. A failed or stale read is repeated up to this many +/// times, `CREDENTIAL_READ_RETRY_DELAY` apart, before the error is reported. +const CREDENTIAL_READ_RETRIES: u32 = 2; +pub(super) const CREDENTIAL_READ_RETRY_DELAY: Duration = Duration::from_millis(50); +const EXTERNAL_OAUTH_REFRESH_WINDOW: chrono::TimeDelta = chrono::Duration::minutes(5); + +static CREDENTIAL_CACHE: OnceLock>> = OnceLock::new(); + +impl CodexApi { + /// Load credentials, tolerating a brief owner publication of `auth.json`. + /// + /// Upstream 0.69.0 #4088 (`CodexOAuthFetchStrategy.loadCredentials` on the + /// usage path, `retryStale: true`): every failed read is repeated. That covers + /// a missing (`NotInstalled`), unreadable (`Other`), malformed or incomplete + /// (`Parse`) file, and a credential the gate rejects as stale (`AuthRequired`, + /// such as a token inside its renewal window), because the CLI may be + /// publishing its renewal. This only rereads the file: no token is redeemed, + /// nothing is written, and the credential cache semantics are unchanged. After + /// the last read the error keeps its category, so unchanged stale credentials + /// still need their owner's renewal. + pub(super) async fn load_credentials(&self) -> Result { + Self::reread_during_owner_publication(|| self.load_credentials_once()).await + } + + /// The bounded reread behind [`Self::load_credentials`]: `read` runs up to + /// `1 + CREDENTIAL_READ_RETRIES` times, `CREDENTIAL_READ_RETRY_DELAY` apart, + /// until it succeeds, and the last result is returned unchanged. Dropping the + /// returned future cancels the pending delay and any further read (upstream + /// checks task cancellation before each read). + pub(super) async fn reread_during_owner_publication( + mut read: impl FnMut() -> Result, + ) -> Result { + let mut retries_remaining = CREDENTIAL_READ_RETRIES; + loop { + match read() { + Err(_) if retries_remaining > 0 => { + retries_remaining -= 1; + tokio::time::sleep(CREDENTIAL_READ_RETRY_DELAY).await; + } + result => return result, + } + } + } + + pub(super) fn load_credentials_once(&self) -> Result { + let auth_path = self.get_auth_path(); + + let metadata = + std::fs::metadata(&auth_path).map_err(|error| self.credential_file_error(error))?; + let modified = metadata.modified().ok(); + if let Some(cached) = Self::cached_credentials(&auth_path, modified) { + Self::enforce_external_oauth_gate(&cached)?; + return Ok(cached); + } + + let content = std::fs::read_to_string(&auth_path) + .map_err(|error| self.credential_file_error(error))?; + + let credentials = Self::parse_credentials_json(&content)?; + Self::enforce_external_oauth_gate(&credentials)?; + Self::store_cached_credentials(auth_path, modified, credentials.clone()); + Ok(credentials) + } + + fn missing_credentials_error(&self) -> ProviderError { + // Upstream 0.50.0 #2679: when the CLI targets Amazon Bedrock or + // another custom backend without ChatGPT auth, sign-in guidance + // is wrong — rate limits simply are not available there. + if self.uses_custom_backend() { + return ProviderError::NotInstalled( + "Codex uses a custom backend (chatgpt_base_url / model_provider) without \ + ChatGPT auth. ChatGPT rate limits are unavailable for this setup." + .to_string(), + ); + } + + ProviderError::NotInstalled( + "Codex auth.json not found. Run `codex login` in a terminal to sign in.".to_string(), + ) + } + + fn credential_file_error(&self, error: std::io::Error) -> ProviderError { + if error.kind() == std::io::ErrorKind::NotFound { + return self.missing_credentials_error(); + } + + ProviderError::Other(format!("Failed to read Codex credentials: {error}")) + } + + pub(super) fn parse_credentials_json(content: &str) -> Result { + let json: serde_json::Value = serde_json::from_str(content) + .map_err(|e| ProviderError::Parse(format!("Invalid Codex credentials JSON: {}", e)))?; + + // Check for OPENAI_API_KEY first + if let Some(api_key) = json.get("OPENAI_API_KEY").and_then(|v| v.as_str()) { + let trimmed = api_key.trim(); + if !trimmed.is_empty() { + return Ok(CodexCredentials { + access_token: trimmed.to_string(), + account_id: None, + is_external_oauth: false, + access_token_expires_at: None, + last_refresh: None, + }); + } + } + + // Otherwise, look for tokens object (external OAuth source) + let tokens = json.get("tokens").ok_or_else(|| { + ProviderError::Parse("Codex auth.json exists but contains no tokens.".to_string()) + })?; + + let access_token = tokens + .get("access_token") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + ProviderError::Parse("Missing access_token in Codex credentials".to_string()) + })? + .to_string(); + + let account_id = tokens + .get("account_id") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()); + + // Upstream 0.50.1 #2944: an OAuth token set with a refresh_token is an + // external (CLI-owned) OAuth source. The `last_refresh` timestamp is + // retained only as provenance for the opt-in safety gate. + let has_refresh_token = tokens + .get("refresh_token") + .and_then(|v| v.as_str()) + .is_some_and(|s| !s.trim().is_empty()); + let last_refresh = json + .get("last_refresh") + .and_then(|v| v.as_str()) + .and_then(parse_timestamp); + + let access_token_expires_at = parse_access_token_expiry(&access_token); + + Ok(CodexCredentials { + access_token, + account_id, + is_external_oauth: has_refresh_token, + access_token_expires_at, + last_refresh, + }) + } + + /// Upstream 0.50.1 #2944: when `codex_external_oauth_sources_allowed` is + /// OFF (the default), external OAuth credential files without refresh + /// provenance fail closed instead of being used silently. An external + /// OAuth source is an auth.json `tokens` object with a `refresh_token` + /// (CLI-owned OAuth, not an API key). Win-CodexBar never refreshes or + /// writes this source: the gate only decides whether the read-only usage + /// request may use it. When the access token is a JWT, its native expiry + /// is the validity authority; opaque tokens are sent to the server. + pub(super) fn enforce_external_oauth_gate( + credentials: &CodexCredentials, + ) -> Result<(), ProviderError> { + if !credentials.is_external_oauth { + return Ok(()); + } + // The opt-in only matters without refresh provenance. Skip the settings + // load otherwise: credential reads repeat while the owner publishes. + let external_sources_allowed = credentials.last_refresh.is_some() + || crate::settings::Settings::load().codex_external_oauth_sources_allowed; + Self::enforce_external_oauth_gate_at(credentials, external_sources_allowed, Utc::now()) + } + + pub(super) fn enforce_external_oauth_gate_at( + credentials: &CodexCredentials, + external_sources_allowed: bool, + now: DateTime, + ) -> Result<(), ProviderError> { + if !credentials.is_external_oauth { + return Ok(()); + } + if !external_sources_allowed && credentials.last_refresh.is_none() { + return Err(ProviderError::AuthRequired); + } + if let Some(expires_at) = credentials.access_token_expires_at + && expires_at - now <= EXTERNAL_OAUTH_REFRESH_WINDOW + { + return Err(ProviderError::AuthRequired); + } + Ok(()) + } + + fn credential_cache() -> &'static Mutex> { + CREDENTIAL_CACHE.get_or_init(|| Mutex::new(None)) + } + + fn cached_credentials( + path: &std::path::Path, + modified: Option, + ) -> Option { + let guard = Self::credential_cache().lock().ok()?; + let cached = guard.as_ref()?; + if cached.path == path + && cached.modified == modified + && cached.loaded_at.elapsed() <= CREDENTIAL_CACHE_TTL + { + return Some(cached.credentials.clone()); + } + None + } + + fn store_cached_credentials( + path: PathBuf, + modified: Option, + credentials: CodexCredentials, + ) { + if let Ok(mut guard) = Self::credential_cache().lock() { + *guard = Some(CachedCodexCredentials { + path, + modified, + loaded_at: Instant::now(), + credentials, + }); + } + } +} + +#[derive(Clone)] +pub(super) struct CodexCredentials { + pub(super) access_token: String, + pub(super) account_id: Option, + /// True when the source is an external OAuth token set (has a + /// `refresh_token`), as opposed to an `OPENAI_API_KEY`. The Codex CLI owns + /// refresh and persistence for this source; this app only reads it. The + /// `codex_external_oauth_sources_allowed` setting gates that read + /// (upstream 0.50.1 #2944). + pub(super) is_external_oauth: bool, + /// Native access-token JWT expiry. When available, this is authoritative + /// for validity; the Codex CLI still owns the refresh lifecycle. + pub(super) access_token_expires_at: Option>, + /// `last_refresh` timestamp from auth.json, when present. Its presence + /// supplies provenance when the external-source opt-in setting is OFF; + /// its age is not an access-token expiry signal. + pub(super) last_refresh: Option>, +} + +struct CachedCodexCredentials { + path: PathBuf, + modified: Option, + loaded_at: Instant, + credentials: CodexCredentials, +} + +/// Parse the native `exp` claim from an access-token JWT. Opaque or malformed +/// tokens return `None` and are handled by the read-only usage request. +fn parse_access_token_expiry(token: &str) -> Option> { + let payload = token.split('.').nth(1)?; + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(payload) + .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload)) + .ok()?; + let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; + let exp = json.get("exp")?.as_i64()?; + Utc.timestamp_opt(exp, 0).single() +} +pub(super) fn parse_timestamp(raw: &str) -> Option> { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return None; + } + DateTime::parse_from_rfc3339(trimmed) + .ok() + .map(|dt| dt.with_timezone(&Utc)) + .or_else(|| { + chrono::NaiveDateTime::parse_from_str(trimmed, "%Y-%m-%dT%H:%M:%S%.f") + .ok() + .map(|naive| DateTime::::from_naive_utc_and_offset(naive, Utc)) + }) +} diff --git a/rust/src/providers/codex/api/parse.rs b/rust/src/providers/codex/api/parse.rs new file mode 100644 index 0000000000..e3cd5efeeb --- /dev/null +++ b/rust/src/providers/codex/api/parse.rs @@ -0,0 +1,463 @@ +use super::CodexApi; +use crate::core::{ + CostSnapshot, NamedRateWindow, ProviderError, RateWindow, RateWindowCadence, UsageSnapshot, +}; +use chrono::{DateTime, TimeZone, Utc}; + +impl CodexApi { + pub(super) fn build_result_from_json( + &self, + json: &serde_json::Value, + ) -> Result<(UsageSnapshot, Option), ProviderError> { + // Extract plan type + let plan_type = json + .get("plan_type") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + // Extract rate limit info - handle multiple possible structures + let (primary, secondary, monthly, code_review, code_review_verified) = + self.extract_rate_limits(json); + + // Build login method string + let login_method = plan_type.as_deref().map(format_plan_type); + + let mut usage = UsageSnapshot::new(primary); + if let Some(sec) = secondary { + usage = usage.with_secondary(sec); + } + // F5 (upstream 0.48.0): monthly (30-day) windows go to tertiary so the + // bridge and frontend can show a monthly reset instead of swallowing it. + if let Some(mo) = monthly { + usage = usage.with_tertiary(mo); + } + if let Some(cr) = code_review { + usage = if code_review_verified { + usage.with_code_review(cr) + } else { + usage.with_model_specific(cr) + }; + } + for extra in self.extract_additional_rate_limits(json) { + usage.extra_rate_windows.push(extra); + } + if let Some(method) = login_method { + usage = usage.with_login_method(method); + } + + // Extract credits if present + let cost = self.extract_credits(json); + + Ok((usage, cost)) + } + + #[cfg(test)] + pub(crate) fn build_result_from_json_for_test( + &self, + json: &serde_json::Value, + ) -> Result<(UsageSnapshot, Option), ProviderError> { + self.build_result_from_json(json) + } + + fn extract_rate_limits( + &self, + json: &serde_json::Value, + ) -> ( + RateWindow, + Option, + Option, + Option, + bool, + ) { + // Try rate_limit object + if let Some(rate_limit) = json.get("rate_limit") { + let primary_opt = rate_limit + .get("primary_window") + .and_then(|w| self.parse_window_if_present(w)); + + let secondary_opt = rate_limit + .get("secondary_window") + .and_then(|w| self.parse_window_if_present(w)); + + let code_review = rate_limit + .get("code_review_window") + .and_then(|w| self.parse_window_if_present(w)); + + let (primary, secondary) = normalize_named_windows(primary_opt, secondary_opt); + + // F5 (upstream 0.48.0): named windows carry only session/weekly/code_review. + // Monthly is extracted separately (from array windows) — return None here. + let code_review_verified = code_review.is_some(); + return (primary, secondary, None, code_review, code_review_verified); + } + + // Try rate_limits array + if let Some(rate_limits) = json.get("rate_limits").and_then(|v| v.as_array()) { + let windows = rate_limits + .iter() + .filter_map(|window| self.parse_window_if_present(window)) + .collect::>(); + let (primary, secondary, monthly, code_review) = normalize_array_windows(windows); + // F5 (upstream 0.48.0): route monthly to its own tertiary lane. + return (primary, secondary, monthly, code_review, false); + } + + // Try direct fields + let used_percent = json + .get("used_percent") + .or_else(|| json.get("usage_percent")) + .and_then(json_f64); + let primary = RateWindow::new(used_percent.unwrap_or(0.0)) + .with_usage_known(valid_used_percent(used_percent)); + + (primary, None, None, None, false) + } + + fn parse_window(&self, window: &serde_json::Value) -> RateWindow { + let used_percent = window + .get("used_percent") + .or_else(|| window.get("usage_percent")) + .and_then(json_f64); + + let window_minutes = window + .get("limit_window_seconds") + .and_then(json_i64) + .and_then(|seconds| u32::try_from(seconds / 60).ok()); + + let reset_at = window + .get("reset_at") + .and_then(json_i64) + .and_then(|ts| Utc.timestamp_opt(ts, 0).single()); + + RateWindow::with_details( + used_percent.unwrap_or(0.0), + window_minutes, + reset_at, + format_reset_countdown(reset_at), + ) + .with_usage_known(valid_used_percent(used_percent)) + } + + fn parse_window_if_present(&self, window: &serde_json::Value) -> Option { + (!window.is_null() && !is_placeholder_window(window)).then(|| self.parse_window(window)) + } + + fn extract_additional_rate_limits(&self, json: &serde_json::Value) -> Vec { + json.get("additional_rate_limits") + .and_then(|v| v.as_array()) + .into_iter() + .flatten() + .filter_map(|entry| self.parse_additional_rate_limit(entry)) + .collect() + } + + fn parse_additional_rate_limit(&self, entry: &serde_json::Value) -> Option { + let metered_feature = entry + .get("metered_feature") + .and_then(|v| v.as_str()) + .map(str::trim) + .filter(|v| !v.is_empty()); + let limit_name = entry + .get("limit_name") + .and_then(|v| v.as_str()) + .map(str::trim) + .filter(|v| !v.is_empty()); + + let rate_limit = entry.get("rate_limit").unwrap_or(entry); + let primary = rate_limit.get("primary_window"); + let secondary = rate_limit.get("secondary_window"); + let window = primary.or(secondary)?; + if is_placeholder_window(window) { + return None; + } + + let parsed = self.parse_window(window); + let feature = metered_feature.unwrap_or_default(); + let limit = limit_name.unwrap_or_default(); + let is_spark = feature.eq_ignore_ascii_case("codex_spark") + || feature.eq_ignore_ascii_case("spark") + || limit.to_ascii_lowercase().contains("spark"); + + if is_spark { + let is_weekly = secondary.is_some() && primary.is_none() + || parsed + .window_minutes + .is_some_and(|mins| mins >= 7 * 24 * 60); + let (id, title) = if is_weekly { + ("codex-spark-weekly", "Codex Spark Weekly") + } else { + ("codex-spark", "Codex Spark 5-hour") + }; + return Some(NamedRateWindow::new(id, title, parsed)); + } + + let label = limit_name.or(metered_feature)?; + let slug = slugify(label); + if slug.is_empty() { + return None; + } + + Some(NamedRateWindow::new( + format!("codex-{slug}"), + titleize_limit_label(label), + parsed, + )) + } + + fn extract_credits(&self, json: &serde_json::Value) -> Option { + let credits = json.get("credits")?; + + let has_credits = credits + .get("has_credits") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if !has_credits { + return None; + } + + let unlimited = credits + .get("unlimited") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if unlimited { + return None; + } + + let balance = credits + .get("balance") + .and_then(|v| v.as_f64()) + .unwrap_or(0.0); + + Some(CostSnapshot::new(balance, "USD", "Credits")) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CodexWindowRole { + Session, + Weekly, + Monthly, + Unknown, +} + +fn codex_window_role(window: &RateWindow) -> CodexWindowRole { + match window + .window_minutes + .map(RateWindowCadence::from_minutes) + .unwrap_or(RateWindowCadence::Unknown) + { + RateWindowCadence::Session => CodexWindowRole::Session, + RateWindowCadence::Monthly => CodexWindowRole::Monthly, + RateWindowCadence::Weekly => CodexWindowRole::Weekly, + RateWindowCadence::Unknown => CodexWindowRole::Unknown, + } +} + +/// Normalize the named `primary_window`/`secondary_window` fields by duration. +pub(super) fn normalize_named_windows( + primary: Option, + secondary: Option, +) -> (RateWindow, Option) { + match (primary, secondary) { + (None, None) => (RateWindow::no_active_session(), None), + (Some(window), None) | (None, Some(window)) => { + if codex_window_role(&window) == CodexWindowRole::Weekly { + (RateWindow::no_active_session(), Some(window)) + } else { + (window, None) + } + } + (Some(primary), Some(secondary)) => { + match (codex_window_role(&primary), codex_window_role(&secondary)) { + (CodexWindowRole::Weekly, CodexWindowRole::Session) => (secondary, Some(primary)), + (CodexWindowRole::Weekly, CodexWindowRole::Unknown) => { + (RateWindow::no_active_session(), Some(primary)) + } + (CodexWindowRole::Unknown, CodexWindowRole::Session) => (secondary, Some(primary)), + _ => (primary, Some(secondary)), + } + } + } +} + +/// Normalize an array of Codex windows without relying on the API's ordering. +/// +/// Returns (session, weekly, monthly, code_review). F5 (upstream 0.48.0): +/// monthly (30-day) windows are routed to their own lane so surfaces can +/// display a monthly reset instead of swallowing it into the weekly label. +pub(super) fn normalize_array_windows( + windows: Vec, +) -> ( + RateWindow, + Option, + Option, + Option, +) { + if windows.is_empty() { + return (RateWindow::no_active_session(), None, None, None); + } + + // Preserve the old positional fallback when the API provides no role + // metadata at all. There is no safe way to infer session vs weekly then. + if !windows + .iter() + .any(|window| codex_window_role(window) != CodexWindowRole::Unknown) + { + let mut windows = windows.into_iter(); + return ( + windows.next().unwrap_or_else(RateWindow::no_active_session), + windows.next(), + windows.next(), + windows.next(), + ); + } + + let mut session = None; + let mut weekly = None; + let mut monthly = None; + let mut remaining = Vec::new(); + + for window in windows { + match codex_window_role(&window) { + CodexWindowRole::Session if session.is_none() => session = Some(window), + CodexWindowRole::Weekly if weekly.is_none() => weekly = Some(window), + CodexWindowRole::Monthly if monthly.is_none() => monthly = Some(window), + _ => remaining.push(window), + } + } + + ( + session.unwrap_or_else(RateWindow::no_active_session), + weekly, + monthly, + remaining.into_iter().next(), + ) +} + +pub(super) fn format_plan_type(plan_type: &str) -> String { + match plan_type { + "guest" => "Guest".to_string(), + "free" => "ChatGPT Free".to_string(), + "go" => "Codex Go".to_string(), + "plus" => "ChatGPT Plus".to_string(), + "pro" => "ChatGPT Pro".to_string(), + "pro_lite" | "prolite" | "pro-lite" => "Pro Lite".to_string(), + "team" => "ChatGPT Team".to_string(), + "business" => "ChatGPT Business".to_string(), + "enterprise" => "ChatGPT Enterprise".to_string(), + "education" | "edu" => "ChatGPT Education".to_string(), + "free_workspace" | "freeWorkspace" => "Free Workspace".to_string(), + "quorum" => "Codex Quorum".to_string(), + "k12" => "Codex K12".to_string(), + other => format!("ChatGPT {}", capitalize(other)), + } +} + +fn json_f64(value: &serde_json::Value) -> Option { + value + .as_f64() + .or_else(|| value.as_i64().map(|value| value as f64)) + .or_else(|| value.as_str()?.trim().parse::().ok()) +} + +fn json_i64(value: &serde_json::Value) -> Option { + value + .as_i64() + .or_else(|| value.as_str()?.trim().parse::().ok()) +} + +fn valid_used_percent(value: Option) -> bool { + value.is_some_and(|value| value.is_finite() && (0.0..=100.0).contains(&value)) +} + +fn is_placeholder_window(window: &serde_json::Value) -> bool { + let has_usage = window + .get("used_percent") + .or_else(|| window.get("usage_percent")) + .and_then(json_f64) + .is_some(); + let has_duration = window + .get("limit_window_seconds") + .and_then(json_i64) + .is_some(); + let has_reset = window.get("reset_at").and_then(json_i64).is_some(); + + !has_usage && !has_duration && !has_reset +} + +fn slugify(label: &str) -> String { + let mut slug = String::new(); + let mut previous_dash = false; + + for ch in label.chars() { + if ch.is_ascii_alphanumeric() { + slug.push(ch.to_ascii_lowercase()); + previous_dash = false; + } else if !previous_dash && !slug.is_empty() { + slug.push('-'); + previous_dash = true; + } + } + + while slug.ends_with('-') { + slug.pop(); + } + slug +} + +fn titleize_limit_label(label: &str) -> String { + label + .split(['_', '-', ' ']) + .filter(|part| !part.is_empty()) + .map(|part| { + let mut chars = part.chars(); + match chars.next() { + Some(first) => first + .to_uppercase() + .chain(chars.flat_map(char::to_lowercase)) + .collect(), + None => String::new(), + } + }) + .collect::>() + .join(" ") +} + +fn format_reset_countdown(reset_at: Option>) -> Option { + let dt = reset_at?; + let now = Utc::now(); + if dt <= now { + return Some("now".to_string()); + } + let diff = dt - now; + let total_mins = diff.num_minutes(); + let hours = diff.num_hours(); + let mins = total_mins % 60; + if hours >= 24 { + let days = hours / 24; + let rem_h = hours % 24; + if rem_h == 0 { + Some(format!("{}d", days)) + } else { + Some(format!("{}d {}h", days, rem_h)) + } + } else if hours > 0 { + if mins == 0 { + Some(format!("{}h", hours)) + } else { + Some(format!("{}h {}m", hours, mins)) + } + } else { + Some(format!("{}m", mins)) + } +} + +fn capitalize(s: &str) -> String { + let mut chars = s.chars(); + match chars.next() { + None => String::new(), + Some(first) => first.to_uppercase().chain(chars).collect(), + } +} diff --git a/rust/src/providers/codex/api/reset_credits.rs b/rust/src/providers/codex/api/reset_credits.rs new file mode 100644 index 0000000000..3743787db3 --- /dev/null +++ b/rust/src/providers/codex/api/reset_credits.rs @@ -0,0 +1,249 @@ +use super::super::weekly_reset; +use super::CodexApi; +use super::credentials::CodexCredentials; +use crate::core::{ProviderError, RateWindow, UsageSnapshot}; +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex as AsyncMutex; + +pub(super) const RESET_CREDITS_PATH: &str = "/wham/rate-limit-reset-credits"; +pub(super) const RESET_CREDITS_CACHE_TTL: Duration = Duration::from_secs(600); + +static RESET_CREDITS_CACHE: OnceLock>>>> = + OnceLock::new(); + +#[derive(Default)] +pub(super) struct ResetCreditsCache { + pub(super) loaded_at: Option, + value: Option, + confirmation_failure_at: Option, +} + +impl ResetCreditsCache { + fn confirmation_failed_recently(&self) -> bool { + self.confirmation_failure_at + .is_some_and(|failed| failed.elapsed() < RESET_CREDITS_CACHE_TTL) + } +} + +impl CodexApi { + pub(super) fn reset_credits_cache_slot( + &self, + creds: &CodexCredentials, + base_url: &str, + ) -> Arc> { + // The Codex home is part of the scope: two homes never share an + // observation, even when they hold the same account and token. + let auth_path = self.get_auth_path(); + let home = weekly_reset::scope_key(None, &auth_path); + let account = weekly_reset::scope_key(creds.account_id.as_deref(), &auth_path); + let token = weekly_reset::scope_key(Some(&creds.access_token), &auth_path); + let key = format!( + "{}|{home}|{account}|{token}", + base_url.trim_end_matches('/') + ); + let cache = RESET_CREDITS_CACHE.get_or_init(|| Mutex::new(HashMap::new())); + let mut cache = cache + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone( + cache + .entry(key) + .or_insert_with(|| Arc::new(AsyncMutex::new(ResetCreditsCache::default()))), + ) + } + + pub(super) async fn fetch_rate_limit_reset_credits_cached( + &self, + creds: &CodexCredentials, + base_url: &str, + ) -> Option { + let slot = self.reset_credits_cache_slot(creds, base_url); + let mut cache = slot.lock().await; + if cache + .loaded_at + .is_some_and(|loaded| loaded.elapsed() < RESET_CREDITS_CACHE_TTL) + { + return cache.value.clone(); + } + cache.value = self + .fetch_rate_limit_reset_credits(creds, base_url) + .await + .ok(); + cache.loaded_at = Some(Instant::now()); + cache.confirmation_failure_at = None; + cache.value.clone() + } + + /// Reset credits for the initial weekly-reset decision. A pending delayed + /// candidate is revalidated against the current inventory, which must be + /// observed after the candidate was stored: the ten-minute cache can still + /// hold the very observation that created it, so that case reads fresh. + pub(super) async fn initial_reset_credits( + &self, + creds: &CodexCredentials, + base_url: &str, + started: Instant, + observed: Option, + candidate_pending: bool, + ) -> Option { + if !candidate_pending { + return observed; + } + self.fresh_reset_credits_for_confirmation(creds, base_url, started) + .await + } + + pub(super) async fn fresh_reset_credits_for_confirmation( + &self, + creds: &CodexCredentials, + base_url: &str, + started: Instant, + ) -> Option { + let slot = self.reset_credits_cache_slot(creds, base_url); + let mut cache = slot.lock().await; + cache.value.as_ref()?; + if cache.confirmation_failed_recently() { + return None; + } + if cache.loaded_at.is_some_and(|loaded| loaded >= started) { + return cache.value.clone(); + } + self.refetch_reset_credits_into(&mut cache, creds, base_url) + .await + } + + pub(super) async fn fetch_rate_limit_reset_credits_fresh( + &self, + creds: &CodexCredentials, + base_url: &str, + ) -> Option { + let slot = self.reset_credits_cache_slot(creds, base_url); + let mut cache = slot.lock().await; + if cache.confirmation_failed_recently() { + return None; + } + self.refetch_reset_credits_into(&mut cache, creds, base_url) + .await + } + + /// A failed read keeps the cached inventory but blocks further + /// confirmation reads for one cache TTL. + async fn refetch_reset_credits_into( + &self, + cache: &mut ResetCreditsCache, + creds: &CodexCredentials, + base_url: &str, + ) -> Option { + let fresh = self + .fetch_rate_limit_reset_credits(creds, base_url) + .await + .ok(); + if let Some(value) = fresh.as_ref() { + cache.value = Some(value.clone()); + cache.loaded_at = Some(Instant::now()); + cache.confirmation_failure_at = None; + } else { + cache.confirmation_failure_at = Some(Instant::now()); + } + fresh + } + + async fn fetch_rate_limit_reset_credits( + &self, + creds: &CodexCredentials, + base_url: &str, + ) -> Result { + let response = self + .authed_get( + &format!("{}{}", base_url, RESET_CREDITS_PATH), + &creds.access_token, + creds.account_id.as_deref(), + ) + .send() + .await?; + if !response.status().is_success() { + return Err( + super::super::authenticated_http_error(response, "Codex reset credits").await, + ); + } + decode_reset_credits(&response.bytes().await?) + } +} + +#[derive(Debug, Clone, Deserialize)] +pub(in crate::providers::codex) struct ResetCredit { + #[serde(default)] + pub(in crate::providers::codex) id: Option, + #[serde(default, alias = "resetType")] + pub(in crate::providers::codex) reset_type: Option, + #[serde(default)] + pub(in crate::providers::codex) status: Option, + #[serde(default)] + pub(in crate::providers::codex) expires_at: Option, +} + +#[derive(Debug, Clone, Deserialize)] +pub(in crate::providers::codex) struct ResetCredits { + #[serde(default)] + pub(in crate::providers::codex) credits: Vec, + pub(in crate::providers::codex) available_count: u32, +} + +pub(super) fn decode_reset_credits(data: &[u8]) -> Result { + serde_json::from_slice(data) + .map_err(|e| ProviderError::Parse(format!("Failed to parse Codex reset credits: {e}"))) +} + +fn parse_credit_expiry(raw: &str) -> Option> { + DateTime::parse_from_rfc3339(raw) + .ok() + .map(|dt| dt.with_timezone(&Utc)) +} + +fn is_available_credit(credit: &ResetCredit) -> bool { + match credit.status.as_deref() { + None | Some("") => true, + Some(status) => status.eq_ignore_ascii_case("available"), + } +} + +pub(in crate::providers::codex) fn next_available_reset_credit_expiry( + credits: &[ResetCredit], + now: DateTime, +) -> Option> { + credits + .iter() + .filter(|credit| is_available_credit(credit)) + .filter_map(|credit| credit.expires_at.as_deref().and_then(parse_credit_expiry)) + .filter(|expires_at| *expires_at > now) + .min() +} + +pub(super) fn reset_credits_rate_window(reset: &ResetCredits, now: DateTime) -> RateWindow { + let description = format!( + "{} reset credit{} available", + reset.available_count, + if reset.available_count == 1 { "" } else { "s" } + ); + let mut window = RateWindow::informational(description); + window.resets_at = next_available_reset_credit_expiry(&reset.credits, now); + window +} + +pub(super) fn apply_reset_credits_window( + mut usage: UsageSnapshot, + reset: Option<&ResetCredits>, +) -> UsageSnapshot { + usage + .extra_rate_windows + .retain(|window| window.id != "reset-credits"); + if let Some(reset) = reset.filter(|reset| reset.available_count > 0) { + let window = reset_credits_rate_window(reset, Utc::now()); + usage = usage.with_extra_rate_window("reset-credits", "Reset credits", window); + } + usage +} diff --git a/rust/src/providers/codex/api/tests.rs b/rust/src/providers/codex/api/tests.rs new file mode 100644 index 0000000000..966e204c37 --- /dev/null +++ b/rust/src/providers/codex/api/tests.rs @@ -0,0 +1,940 @@ +use super::credentials::parse_timestamp; +use super::parse::{normalize_array_windows, normalize_named_windows}; +use super::reset_credits::{ + RESET_CREDITS_CACHE_TTL, RESET_CREDITS_PATH, decode_reset_credits, reset_credits_rate_window, +}; +use super::*; +use crate::core::RateWindow; +use crate::providers::test_support::{mock_response, mock_response_expect, mock_status_expect}; +use base64::Engine; +use chrono::DateTime; +use serde_json::json; +use std::time::Duration; + +#[test] +fn non_chatgpt_model_provider_is_detected_for_guidance() { + // Upstream 0.50.0 #2679: Bedrock and other custom backends get + // rate-limit guidance instead of login instructions. + assert!(config_uses_non_chatgpt_provider( + "model_provider = \"bedrock\"\n" + )); + assert!(config_uses_non_chatgpt_provider( + "# relay\nmodel_provider = 'ollama'" + )); + assert!(!config_uses_non_chatgpt_provider( + "model_provider = \"openai\"" + )); + assert!(!config_uses_non_chatgpt_provider( + "model = \"gpt-5\"\napproval_policy = \"never\"" + )); +} + +#[test] +fn parses_codex_credentials_without_retaining_refresh_token() { + let credentials = CodexApi::parse_credentials_json( + r#"{ + "tokens": { + "access_token": "access", + "refresh_token": "refresh", + "account_id": "acct_123" + } + }"#, + ) + .expect("credentials"); + + assert_eq!(credentials.access_token, "access"); + assert_eq!(credentials.account_id.as_deref(), Some("acct_123")); +} + +#[test] +fn decodes_reset_credits() { + let credits = decode_reset_credits( + br#"{"available_count":2,"credits":[{"id":"a","status":"available","expires_at":"2026-08-01T12:00:00Z"}]}"#, + ) + .expect("reset credits"); + assert_eq!(credits.available_count, 2); + assert_eq!(credits.credits.len(), 1); + assert_eq!(credits.credits[0].status.as_deref(), Some("available")); + assert_eq!( + credits.credits[0].expires_at.as_deref(), + Some("2026-08-01T12:00:00Z") + ); +} + +#[test] +fn missing_reset_credit_count_is_unavailable_not_zero() { + assert!(decode_reset_credits(br#"{"credits":[]}"#).is_err()); +} + +#[test] +fn next_expiry_picks_soonest_available() { + let now = utc("2026-07-01T00:00:00Z"); + let credits = vec![ + credit(Some("available"), "2026-07-10T00:00:00Z"), + credit(Some("available"), "2026-07-05T00:00:00Z"), + credit(Some("available"), "2026-07-20T00:00:00Z"), + ]; + let expiry = next_available_reset_credit_expiry(&credits, now).expect("expiry"); + assert_eq!(expiry, utc("2026-07-05T00:00:00Z")); +} + +#[test] +fn next_expiry_skips_past_and_non_available() { + let now = utc("2026-07-01T00:00:00Z"); + let credits = vec![ + credit(Some("available"), "2026-06-01T00:00:00Z"), + credit(Some("used"), "2026-07-03T00:00:00Z"), + credit(Some("AVAILABLE"), "2026-07-08T00:00:00Z"), + credit(None, "2026-07-09T00:00:00Z"), + ]; + let expiry = next_available_reset_credit_expiry(&credits, now).expect("expiry"); + assert_eq!(expiry, utc("2026-07-08T00:00:00Z")); +} + +#[test] +fn reset_credits_window_sets_informational_and_expiry() { + let now = utc("2026-07-01T00:00:00Z"); + let reset = ResetCredits { + available_count: 2, + credits: vec![ + credit(Some("available"), "2026-07-15T12:00:00Z"), + credit(Some("available"), "2026-07-10T12:00:00Z"), + ], + }; + let window = reset_credits_rate_window(&reset, now); + assert!(window.is_informational); + assert_eq!( + window.reset_description.as_deref(), + Some("2 reset credits available") + ); + assert_eq!(window.resets_at, Some(utc("2026-07-10T12:00:00Z"))); +} + +#[test] +fn reset_credits_window_count_only_without_expiry() { + let now = utc("2026-07-01T00:00:00Z"); + let reset = ResetCredits { + available_count: 1, + credits: vec![], + }; + let window = reset_credits_rate_window(&reset, now); + assert!(window.is_informational); + assert_eq!( + window.reset_description.as_deref(), + Some("1 reset credit available") + ); + assert!(window.resets_at.is_none()); +} + +const PLUS_USAGE: &str = r#"{"plan_type":"plus","rate_limit":{"primary_window":{"used_percent":10,"limit_window_seconds":18000}}}"#; + +fn utc(value: &str) -> DateTime { + DateTime::parse_from_rfc3339(value) + .unwrap() + .with_timezone(&Utc) +} + +fn credit(status: Option<&str>, expires_at: &str) -> ResetCredit { + ResetCredit { + id: None, + reset_type: None, + status: status.map(str::to_string), + expires_at: Some(expires_at.to_string()), + } +} + +/// An unsigned JWT whose payload carries only `exp` (seconds since the epoch). +fn jwt_with_exp(exp: i64) -> String { + let payload = + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(format!(r#"{{"exp":{exp}}}"#)); + format!("header.{payload}.signature") +} + +fn external_creds(access_token: &str, last_refresh: Option>) -> CodexCredentials { + CodexCredentials { + access_token: access_token.to_string(), + account_id: None, + is_external_oauth: true, + access_token_expires_at: None, + last_refresh, + } +} + +async fn mock_reset_credits( + server: &mut mockito::ServerGuard, + available_count: u32, +) -> mockito::Mock { + mock_response_expect( + server, + "GET", + RESET_CREDITS_PATH, + 200, + format!(r#"{{"available_count":{available_count},"credits":[]}}"#), + 1, + ) + .await +} + +fn write_codex_home(base_url: &str) -> tempfile::TempDir { + let dir = tempfile::tempdir().expect("temp codex home"); + std::fs::write( + dir.path().join("auth.json"), + r#"{"tokens":{"access_token":"test-token","account_id":"acct_test"}}"#, + ) + .expect("auth.json"); + std::fs::write( + dir.path().join("config.toml"), + format!("chatgpt_base_url = \"{base_url}\""), + ) + .expect("config.toml"); + dir +} + +#[tokio::test] +async fn reset_credit_cache_single_flight_and_unknown_are_ten_minute_observations() { + let mut server = mockito::Server::new_async().await; + let request = mock_status_expect(&mut server, "GET", RESET_CREDITS_PATH, 503, 1).await; + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let creds = api.load_credentials().await.unwrap(); + let base = server.url(); + let (first, second) = tokio::join!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base), + api.fetch_rate_limit_reset_credits_cached(&creds, &base), + ); + assert!(first.is_none() && second.is_none()); + assert!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .is_none() + ); + request.assert_async().await; + assert!(RESET_CREDITS_CACHE_TTL == Duration::from_secs(600)); +} + +#[tokio::test] +async fn reset_credit_cache_expires_and_token_rotation_uses_new_scope() { + let mut server = mockito::Server::new_async().await; + let first = server + .mock("GET", "/wham/rate-limit-reset-credits") + .match_header("authorization", "Bearer test-token") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"available_count":2,"credits":[]}"#) + .create_async() + .await; + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let mut creds = api.load_credentials().await.unwrap(); + let base = server.url(); + assert_eq!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .unwrap() + .available_count, + 2 + ); + assert_eq!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .unwrap() + .available_count, + 2 + ); + first.assert_async().await; + first.remove_async().await; + let second = server + .mock("GET", "/wham/rate-limit-reset-credits") + .match_header("authorization", "Bearer test-token") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"available_count":1,"credits":[]}"#) + .create_async() + .await; + api.reset_credits_cache_slot(&creds, &base) + .lock() + .await + .loaded_at = Some(Instant::now() - RESET_CREDITS_CACHE_TTL); + assert_eq!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .unwrap() + .available_count, + 1 + ); + second.assert_async().await; + second.remove_async().await; + creds.access_token = "rotated-token".into(); + let rotated = server + .mock("GET", "/wham/rate-limit-reset-credits") + .match_header("authorization", "Bearer rotated-token") + .expect(1) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"available_count":3,"credits":[]}"#) + .create_async() + .await; + assert_eq!( + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .unwrap() + .available_count, + 3 + ); + rotated.assert_async().await; +} + +#[tokio::test] +async fn reset_credit_cache_is_scoped_to_the_codex_home() { + // Same base URL, account and token, different Codex homes: each home + // makes its own observation instead of reading the other's. + let mut server = mockito::Server::new_async().await; + let base = server.url(); + let first = mock_reset_credits(&mut server, 2).await; + let first_home = write_codex_home(&base); + let first_api = CodexApi::new().with_codex_home(first_home.path()); + let first_creds = first_api.load_credentials().await.unwrap(); + let first_count = first_api + .fetch_rate_limit_reset_credits_cached(&first_creds, &base) + .await + .map(|credits| credits.available_count); + assert_eq!(first_count, Some(2)); + first.assert_async().await; + first.remove_async().await; + + let second = mock_status_expect(&mut server, "GET", RESET_CREDITS_PATH, 503, 1).await; + let second_home = write_codex_home(&base); + let second_api = CodexApi::new().with_codex_home(second_home.path()); + let second_creds = second_api.load_credentials().await.unwrap(); + assert_eq!(first_creds.access_token, second_creds.access_token); + assert_eq!(first_creds.account_id, second_creds.account_id); + assert!( + second_api + .fetch_rate_limit_reset_credits_cached(&second_creds, &base) + .await + .is_none() + ); + second.assert_async().await; +} + +#[tokio::test] +async fn suspicious_weekly_reset_uses_independent_credit_observations() { + let mut server = mockito::Server::new_async().await; + let cached_response = mock_reset_credits(&mut server, 2).await; + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let creds = api.load_credentials().await.unwrap(); + let base = server.url(); + api.fetch_rate_limit_reset_credits_cached(&creds, &base) + .await + .unwrap(); + cached_response.assert_async().await; + cached_response.remove_async().await; + + let started = Instant::now(); + let initial_response = mock_reset_credits(&mut server, 1).await; + let initial = api + .fresh_reset_credits_for_confirmation(&creds, &base, started) + .await + .unwrap(); + assert_eq!(initial.available_count, 1); + initial_response.assert_async().await; + initial_response.remove_async().await; + + let confirmation_response = mock_reset_credits(&mut server, 0).await; + let confirmation = api + .fetch_rate_limit_reset_credits_fresh(&creds, &base) + .await + .unwrap(); + assert_eq!(confirmation.available_count, 0); + confirmation_response.assert_async().await; +} + +#[tokio::test] +async fn pending_delayed_candidate_revalidates_with_a_fresh_credit_observation() { + let mut server = mockito::Server::new_async().await; + let candidate_observation = mock_reset_credits(&mut server, 2).await; + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let creds = api.load_credentials().await.unwrap(); + let base = server.url(); + let cached = api + .fetch_rate_limit_reset_credits_cached(&creds, &base) + .await; + candidate_observation.assert_async().await; + candidate_observation.remove_async().await; + + // A later refresh: the ten-minute cache still holds the observation + // that created the candidate. + let started = Instant::now(); + let changed = mock_reset_credits(&mut server, 1).await; + let without_candidate = api + .initial_reset_credits(&creds, &base, started, cached.clone(), false) + .await; + assert_eq!( + without_candidate.map(|credits| credits.available_count), + Some(2) + ); + let with_candidate = api + .initial_reset_credits(&creds, &base, started, cached, true) + .await; + assert_eq!( + with_candidate.map(|credits| credits.available_count), + Some(1) + ); + changed.assert_async().await; +} + +#[tokio::test] +async fn fetch_usage_attaches_reset_credits_from_http() { + let mut server = mockito::Server::new_async().await; + let soonest = (Utc::now() + chrono::Duration::days(5)).to_rfc3339(); + let later = (Utc::now() + chrono::Duration::days(12)).to_rfc3339(); + + let usage_mock = server + .mock("GET", "/wham/usage") + .match_header("authorization", "Bearer test-token") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(PLUS_USAGE) + .create_async() + .await; + + let reset_body = format!( + r#"{{"available_count":2,"credits":[ + {{"status":"available","expires_at":"{later}"}}, + {{"status":"available","expires_at":"{soonest}"}} + ]}}"# + ); + let reset_mock = server + .mock("GET", "/wham/rate-limit-reset-credits") + .match_header("authorization", "Bearer test-token") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(reset_body) + .create_async() + .await; + + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let (usage, _, _) = api.fetch_usage().await.expect("fetch_usage"); + + usage_mock.assert_async().await; + reset_mock.assert_async().await; + + let extra = usage + .extra_rate_windows + .iter() + .find(|w| w.id == "reset-credits") + .expect("reset-credits window attached"); + assert_eq!(extra.title, "Reset credits"); + assert!(extra.window.is_informational); + assert_eq!( + extra.window.reset_description.as_deref(), + Some("2 reset credits available") + ); + let expected = utc(&soonest); + assert_eq!(extra.window.resets_at, Some(expected)); +} + +#[tokio::test] +async fn authenticated_codex_http_distinguishes_401_from_403() { + for (status, expects_authentication) in [(401, true), (403, false)] { + let mut server = mockito::Server::new_async().await; + let mock = server + .mock("GET", "/wham/usage") + .with_status(status) + .with_body("fixture refusal") + .create_async() + .await; + + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let error = match api.fetch_usage().await { + Ok(_) => panic!("expected HTTP {status} to fail"), + Err(error) => error, + }; + + if expects_authentication { + assert!(matches!(error, ProviderError::AuthRequired)); + } else { + let message = error.to_string(); + assert!(message.contains("403")); + assert!(message.contains("fixture refusal")); + assert!(!matches!(error, ProviderError::AuthRequired)); + } + mock.assert_async().await; + } +} + +#[tokio::test] +async fn old_opaque_external_oauth_reaches_usage_request() { + let mut server = mockito::Server::new_async().await; + let usage_mock = server + .mock("GET", "/wham/usage") + .match_header("authorization", "Bearer opaque-token") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(PLUS_USAGE) + .create_async() + .await; + let reset_mock = server + .mock("GET", "/wham/rate-limit-reset-credits") + .match_header("authorization", "Bearer opaque-token") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"available_count":0,"credits":[]}"#) + .create_async() + .await; + + let creds = CodexApi::parse_credentials_json( + r#"{ + "tokens": { + "access_token": "opaque-token", + "refresh_token": "refresh", + "account_id": "acct_test" + }, + "last_refresh": "2026-01-01T00:00:00Z" + }"#, + ) + .expect("credentials"); + assert!(CodexApi::enforce_external_oauth_gate_at(&creds, false, Utc::now()).is_ok()); + + let api = CodexApi::new(); + let (usage, _, _) = api + .fetch_usage_once(&creds, &server.url()) + .await + .expect("opaque OAuth usage request"); + assert_eq!(usage.primary.used_percent, 10.0); + usage_mock.assert_async().await; + reset_mock.assert_async().await; +} + +#[tokio::test] +async fn fetch_usage_skips_reset_credits_when_available_count_zero() { + let mut server = mockito::Server::new_async().await; + + let usage_mock = mock_response(&mut server, "GET", USAGE_PATH, 200, PLUS_USAGE).await; + + let reset_mock = mock_reset_credits(&mut server, 0).await; + + let home = write_codex_home(&server.url()); + let api = CodexApi::new().with_codex_home(home.path()); + let (usage, _, _) = api.fetch_usage().await.expect("fetch_usage"); + + usage_mock.assert_async().await; + reset_mock.assert_async().await; + + assert!( + usage + .extra_rate_windows + .iter() + .all(|w| w.id != "reset-credits"), + "available_count=0 must not attach reset-credits" + ); +} + +#[test] +fn keeps_weekly_window_in_secondary_when_session_is_absent() { + let api = CodexApi::new(); + let (usage, _) = api + .build_result_from_json(&json!({ + "rate_limit": { + "secondary_window": { + "used_percent": 25, + "limit_window_seconds": 604800, + "reset_at": 1783036800 + } + } + })) + .expect("codex usage"); + + assert!(usage.primary.is_informational); + assert_eq!(usage.primary.window_minutes, Some(300)); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("No active 5h session") + ); + + let weekly = usage.secondary.expect("weekly window"); + assert!(!weekly.is_informational); + assert_eq!(weekly.used_percent, 25.0); + assert_eq!(weekly.window_minutes, Some(10080)); +} + +#[test] +fn identifies_rate_limit_array_windows_by_duration() { + let api = CodexApi::new(); + let (usage, _) = api + .build_result_from_json(&json!({ + "rate_limits": [ + { + "used_percent": 25, + "limit_window_seconds": 604800, + "reset_at": 1783036800 + }, + { + "used_percent": 10, + "limit_window_seconds": 18000, + "reset_at": 1783018800 + } + ] + })) + .expect("codex usage"); + + assert!(!usage.primary.is_informational); + assert_eq!(usage.primary.used_percent, 10.0); + assert_eq!(usage.primary.window_minutes, Some(300)); + + let weekly = usage.secondary.expect("weekly window"); + assert_eq!(weekly.used_percent, 25.0); + assert_eq!(weekly.window_minutes, Some(10080)); +} + +#[test] +fn identifies_weekly_only_rate_limit_array_without_a_session() { + let api = CodexApi::new(); + let (usage, _) = api + .build_result_from_json(&json!({ + "rate_limits": [{ + "used_percent": 25, + "limit_window_seconds": 604800, + "reset_at": 1783036800 + }] + })) + .expect("codex usage"); + + assert!(usage.primary.is_informational); + assert_eq!(usage.secondary.expect("weekly window").used_percent, 25.0); +} + +#[test] +fn maps_codex_spark_additional_rate_limits() { + let api = CodexApi::new(); + let (usage, _) = api + .build_result_from_json(&json!({ + "plan_type": "pro", + "rate_limit": { + "primary_window": { "used_percent": 20, "limit_window_seconds": 18000 }, + "secondary_window": { "used_percent": 40, "limit_window_seconds": 604800 } + }, + "additional_rate_limits": [ + { + "limit_name": "Codex Spark", + "metered_feature": "codex_spark", + "rate_limit": { + "primary_window": { "used_percent": "17", "limit_window_seconds": 18000 } + } + }, + { + "limit_name": "Codex Spark Weekly", + "metered_feature": "codex_spark", + "rate_limit": { + "secondary_window": { "used_percent": 62, "limit_window_seconds": 604800 } + } + } + ] + })) + .expect("codex usage"); + + assert_eq!(usage.extra_rate_windows.len(), 2); + assert_eq!(usage.extra_rate_windows[0].id, "codex-spark"); + assert_eq!(usage.extra_rate_windows[0].title, "Codex Spark 5-hour"); + assert_eq!(usage.extra_rate_windows[0].window.used_percent, 17.0); + assert_eq!(usage.extra_rate_windows[1].id, "codex-spark-weekly"); + assert_eq!(usage.extra_rate_windows[1].title, "Codex Spark Weekly"); + assert_eq!(usage.extra_rate_windows[1].window.used_percent, 62.0); +} + +#[test] +fn ignores_placeholder_additional_rate_limits() { + let api = CodexApi::new(); + let (usage, _) = api + .build_result_from_json(&json!({ + "rate_limit": { + "primary_window": { "used_percent": 0, "limit_window_seconds": 18000 } + }, + "additional_rate_limits": [ + { + "limit_name": "placeholder", + "metered_feature": "placeholder", + "rate_limit": { "primary_window": {} } + } + ] + })) + .expect("codex usage"); + + assert!(usage.extra_rate_windows.is_empty()); +} + +fn win(minutes: u32, used: f64) -> RateWindow { + RateWindow::with_details(used, Some(minutes), None, None) +} + +#[test] +fn normalize_named_windows_routes_by_role() { + // Primary is reported as -1.0 for the "No active 5h session" placeholder. + let s = |used| win(300, used); + let w = |used| win(10_080, used); + let u = |used| win(999, used); + let m = |used| win(43_200, used); + type Row = (Option, Option, f64, Option); + let rows: Vec = vec![ + (None, None, -1.0, None), + (Some(w(1.0)), None, -1.0, Some(1.0)), + (Some(s(1.0)), None, 1.0, None), + (Some(u(1.0)), None, 1.0, None), + (Some(m(1.0)), None, 1.0, None), + (None, Some(w(2.0)), -1.0, Some(2.0)), + (None, Some(s(2.0)), 2.0, None), + (None, Some(u(2.0)), 2.0, None), + (Some(w(1.0)), Some(s(2.0)), 2.0, Some(1.0)), + (Some(w(1.0)), Some(u(2.0)), -1.0, Some(1.0)), + (Some(u(1.0)), Some(s(2.0)), 2.0, Some(1.0)), + (Some(s(1.0)), Some(w(2.0)), 1.0, Some(2.0)), + (Some(u(1.0)), Some(w(2.0)), 1.0, Some(2.0)), + (Some(s(1.0)), Some(s(2.0)), 1.0, Some(2.0)), + (Some(w(1.0)), Some(w(2.0)), 1.0, Some(2.0)), + (Some(u(1.0)), Some(u(2.0)), 1.0, Some(2.0)), + (Some(m(1.0)), Some(w(2.0)), 1.0, Some(2.0)), + ]; + for (index, (primary, secondary, want_primary, want_secondary)) in rows.into_iter().enumerate() + { + let (got_primary, got_secondary) = normalize_named_windows(primary, secondary); + let got_primary = if got_primary.is_informational { + -1.0 + } else { + got_primary.used_percent + }; + assert_eq!(got_primary, want_primary, "row {index} primary"); + assert_eq!( + got_secondary.map(|window| window.used_percent), + want_secondary, + "row {index} secondary" + ); + } +} + +#[tokio::test] +async fn authed_get_sends_account_header_only_when_non_empty() { + let mut server = mockito::Server::new_async().await; + let with_account = server + .mock("GET", "/with") + .match_header("authorization", "Bearer tok") + .match_header("user-agent", "CodexBar") + .match_header("accept", "application/json") + .match_header("chatgpt-account-id", "acct-1") + .with_status(200) + .create_async() + .await; + let without_account = server + .mock("GET", "/without") + .match_header("authorization", "Bearer tok") + .match_header("chatgpt-account-id", mockito::Matcher::Missing) + .expect(2) + .with_status(200) + .create_async() + .await; + let api = CodexApi::new(); + for (path, account_id) in [ + ("/with", Some("acct-1")), + ("/without", Some("")), + ("/without", None), + ] { + let status = api + .authed_get(&format!("{}{path}", server.url()), "tok", account_id) + .send() + .await + .expect("send") + .status(); + assert_eq!(status.as_u16(), 200, "{path} {account_id:?}"); + } + with_account.assert_async().await; + without_account.assert_async().await; +} + +#[test] +fn f5_normalize_array_routes_session_weekly_monthly_to_lanes() { + // 5h session + weekly + monthly → (session, weekly, monthly, None) + let windows = vec![win(300, 10.0), win(10_080, 20.0), win(43_200, 30.0)]; + let (primary, secondary, tertiary, code_review) = normalize_array_windows(windows); + assert_eq!(primary.window_minutes, Some(300)); + assert_eq!(secondary.unwrap().window_minutes, Some(10_080)); + assert_eq!(tertiary.unwrap().window_minutes, Some(43_200)); + assert!(code_review.is_none()); +} + +#[test] +fn f5_normalize_array_monthly_routes_to_tertiary_not_secondary() { + // Monthly must go to tertiary, NOT secondary — so #268's weekly math + // and "Weekly" label stay untouched. + let windows = vec![win(43_200, 50.0), win(10_080, 20.0)]; + let (primary, secondary, tertiary, _) = normalize_array_windows(windows); + assert_eq!(primary.window_minutes, Some(300)); // no session → placeholder + assert_eq!(secondary.unwrap().window_minutes, Some(10_080)); + assert_eq!(tertiary.unwrap().window_minutes, Some(43_200)); +} + +#[test] +fn f5_normalize_array_empty_returns_placeholder_primary() { + let (primary, secondary, tertiary, code_review) = normalize_array_windows(vec![]); + assert!(primary.is_informational); + assert!(secondary.is_none()); + assert!(tertiary.is_none()); + assert!(code_review.is_none()); +} + +#[test] +fn f5_normalize_array_unknown_windows_fall_to_code_review() { + // Windows with unrecognized durations (not 300/10080/43200) go to the + // remaining/code_review bucket. + let windows = vec![win(300, 10.0), win(999, 5.0)]; + let (primary, secondary, tertiary, code_review) = normalize_array_windows(windows); + assert_eq!(primary.window_minutes, Some(300)); + assert!(secondary.is_none()); + assert!(tertiary.is_none()); + assert_eq!(code_review.unwrap().window_minutes, Some(999)); +} + +// ── Upstream 0.50.1 #2944: external OAuth source gate ────────────────── + +#[test] +fn confirmation_failure_fallback_keeps_first_successful_usage_and_cost() { + let state = weekly_reset::AccountState::default(); + let first = UsageSnapshot::new(RateWindow::new(10.0)).with_secondary(RateWindow::new(0.5)); + let cost = Some(CostSnapshot::new(3.25, "USD", "Monthly")); + let (usage, kept_cost) = (weekly_reset::preserve_weekly(&state, first), cost); + assert!((usage.secondary.expect("weekly").used_percent - 0.5).abs() < f64::EPSILON); + assert_eq!(kept_cost.expect("cost").used, 3.25); +} +#[test] +fn api_key_credentials_are_not_external_oauth() { + let creds = + CodexApi::parse_credentials_json(r#"{"OPENAI_API_KEY": "sk-test"}"#).expect("credentials"); + assert!(!creds.is_external_oauth); + assert!(creds.access_token_expires_at.is_none()); + assert!(creds.last_refresh.is_none()); + assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); +} + +#[test] +fn oauth_tokens_with_refresh_token_are_external_source() { + let creds = CodexApi::parse_credentials_json( + r#"{ + "tokens": { + "access_token": "access", + "refresh_token": "refresh", + "account_id": "acct_123" + } + }"#, + ) + .expect("credentials"); + assert!(creds.is_external_oauth); + assert!(creds.access_token_expires_at.is_none()); + assert!(creds.last_refresh.is_none()); +} + +#[test] +fn oauth_tokens_without_refresh_token_are_not_external() { + let creds = CodexApi::parse_credentials_json( + r#"{ + "tokens": { + "access_token": "access", + "account_id": "acct_123" + } + }"#, + ) + .expect("credentials"); + assert!(!creds.is_external_oauth); +} + +#[test] +fn external_oauth_gate_fails_closed_without_last_refresh() { + let creds = external_creds("access", None); + let err = CodexApi::enforce_external_oauth_gate(&creds) + .expect_err("external OAuth without provenance must fail closed"); + assert!(matches!(err, ProviderError::AuthRequired)); +} + +#[test] +fn external_oauth_gate_ignores_old_last_refresh_for_opaque_token() { + let old = Utc::now() - chrono::Duration::days(10); + let creds = external_creds("access", Some(old)); + assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); +} + +#[test] +fn external_oauth_gate_allows_refresh_provenance() { + let fresh = Utc::now() - chrono::Duration::hours(1); + let creds = external_creds("access", Some(fresh)); + assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); +} + +#[test] +fn external_oauth_gate_uses_future_jwt_expiry_over_old_last_refresh() { + let now = Utc::now(); + let future = now + chrono::Duration::hours(2); + let token = jwt_with_exp(future.timestamp()); + let json = format!( + r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"2026-01-01T00:00:00Z"}}"# + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + assert!(creds.access_token_expires_at.is_some()); + assert!(CodexApi::enforce_external_oauth_gate_at(&creds, false, now).is_ok()); + assert!(CodexApi::enforce_external_oauth_gate_at(&creds, true, now).is_ok()); +} + +#[test] +fn external_oauth_gate_rejects_expired_jwt() { + let expired = Utc::now() - chrono::Duration::minutes(1); + let token = jwt_with_exp(expired.timestamp()); + let json = format!( + r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"{}"}}"#, + Utc::now().to_rfc3339() + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + let err = CodexApi::enforce_external_oauth_gate(&creds) + .expect_err("expired native OAuth must be rejected"); + assert!(matches!(err, ProviderError::AuthRequired)); +} + +#[test] +fn external_oauth_gate_requires_cli_refresh_when_jwt_is_near_expiry() { + let soon = Utc::now() + chrono::Duration::minutes(2); + let token = jwt_with_exp(soon.timestamp()); + let json = format!( + r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"{}"}}"#, + Utc::now().to_rfc3339() + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + let err = CodexApi::enforce_external_oauth_gate(&creds) + .expect_err("near-expiry native OAuth must refresh through the CLI"); + assert!(matches!(err, ProviderError::AuthRequired)); +} + +#[test] +fn external_oauth_gate_allows_missing_last_refresh_when_opted_in() { + let creds = external_creds("opaque-token", None); + assert!(CodexApi::enforce_external_oauth_gate_at(&creds, true, Utc::now()).is_ok()); +} + +#[test] +fn opaque_token_uses_refresh_provenance_when_no_jwt_expiry_exists() { + let fresh = Utc::now().to_rfc3339(); + let json = format!( + r#"{{"tokens":{{"access_token":"opaque-token","refresh_token":"refresh"}},"last_refresh":"{fresh}"}}"# + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + assert!(creds.access_token_expires_at.is_none()); + assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); +} +#[test] +fn parse_timestamp_reads_iso8601() { + assert!(parse_timestamp("2026-08-17T10:00:00Z").is_some()); + assert!(parse_timestamp("2026-08-17T10:00:00.123Z").is_some()); + assert!(parse_timestamp(" 2026-08-17T10:00:00Z ").is_some()); + assert!(parse_timestamp("").is_none()); + assert!(parse_timestamp("not-a-date").is_none()); +} diff --git a/rust/src/providers/codex/reset_observations.rs b/rust/src/providers/codex/reset_observations.rs index 0902e25f1a..cab1b4e4d8 100644 --- a/rust/src/providers/codex/reset_observations.rs +++ b/rust/src/providers/codex/reset_observations.rs @@ -69,17 +69,6 @@ pub struct CodexResetObservationMergeResult { pub changed: bool, } -/// Path of the store under the shared configuration root. -pub fn default_store_path() -> Result { - let root = dirs::config_dir().ok_or_else(|| { - CodexResetObservationError::Read(std::io::Error::new( - std::io::ErrorKind::NotFound, - "configuration directory not found", - )) - })?; - Ok(root.join("CodexBar").join(STORE_RELATIVE_PATH)) -} - /// Path of the store relative to an explicit config root (tests, proof homes). pub fn store_path(config_root: &Path) -> PathBuf { config_root.join(STORE_RELATIVE_PATH) @@ -93,16 +82,20 @@ fn validate_scope(account_scope: &str) -> Result<(), CodexResetObservationError> } } -/// Read the observations recorded for `account_scope`. -pub fn load_reset_observations( - config_root: &Path, - account_scope: &str, -) -> Result, CodexResetObservationError> { - validate_scope(account_scope)?; - let path = store_path(config_root); - let raw = match secure_file::read_string(&path) { +/// A missing store reads as empty. Only the merge path also treats a blank +/// file as empty; `load_reset_observations` reports it as a decode error. +fn read_store( + path: &Path, + blank_is_empty: bool, +) -> Result { + let raw = match secure_file::read_string(path) { + Ok(raw) if blank_is_empty && raw.trim().is_empty() => { + return Ok(CodexResetObservationStore::default()); + } Ok(raw) => raw, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => { + return Ok(CodexResetObservationStore::default()); + } Err(error) => return Err(CodexResetObservationError::Read(error)), }; let store: CodexResetObservationStore = @@ -112,6 +105,16 @@ pub fn load_reset_observations( store.version, )); } + Ok(store) +} + +/// Read the observations recorded for `account_scope`. +pub fn load_reset_observations( + config_root: &Path, + account_scope: &str, +) -> Result, CodexResetObservationError> { + validate_scope(account_scope)?; + let store = read_store(&store_path(config_root), false)?; Ok(store .accounts .get(account_scope) @@ -163,19 +166,7 @@ pub fn merge_and_persist_reset_observation( } let path = store_path(config_root); - let mut store = match secure_file::read_string(&path) { - Ok(raw) if raw.trim().is_empty() => CodexResetObservationStore::default(), - Ok(raw) => serde_json::from_str(&raw).map_err(CodexResetObservationError::Deserialize)?, - Err(error) if error.kind() == std::io::ErrorKind::NotFound => { - CodexResetObservationStore::default() - } - Err(error) => return Err(CodexResetObservationError::Read(error)), - }; - if store.version != STORE_VERSION { - return Err(CodexResetObservationError::UnsupportedVersion( - store.version, - )); - } + let mut store = read_store(&path, true)?; let (changed, observations) = { let rows = store.accounts.entry(account_scope.to_string()).or_default(); let changed = merge_reset_observations( @@ -312,6 +303,51 @@ mod tests { )); } + #[test] + fn blank_store_fails_load_but_merges_as_empty() { + let root = temp_root("blank"); + let path = store_path(root.path()); + std::fs::create_dir_all(path.parent().unwrap()).expect("mkdir"); + secure_file::write_string( + &path, " +", + ) + .expect("write"); + assert!(matches!( + load_reset_observations(root.path(), CODEX_ACCOUNT_SCOPE), + Err(CodexResetObservationError::Deserialize(_)) + )); + let merged = merge_and_persist_reset_observation( + root.path(), + CODEX_ACCOUNT_SCOPE, + at(4_000), + at(1_000), + ) + .expect("merge over blank store"); + assert!(merged.changed); + assert_eq!( + load_reset_observations(root.path(), CODEX_ACCOUNT_SCOPE).expect("load"), + merged.observations + ); + } + + #[test] + fn unsupported_version_blocks_merge() { + let root = temp_root("merge-version"); + let path = store_path(root.path()); + std::fs::create_dir_all(path.parent().unwrap()).expect("mkdir"); + secure_file::write_string(&path, r#"{"version": 99, "accounts": {}}"#).expect("write"); + assert!(matches!( + merge_and_persist_reset_observation( + root.path(), + CODEX_ACCOUNT_SCOPE, + at(4_000), + at(1_000) + ), + Err(CodexResetObservationError::UnsupportedVersion(99)) + )); + } + #[test] fn persisted_file_is_secure_wrapped() { let root = temp_root("secure"); diff --git a/rust/src/providers/codex/subscription.rs b/rust/src/providers/codex/subscription.rs index 8a127ac445..175cdd90a3 100644 --- a/rust/src/providers/codex/subscription.rs +++ b/rust/src/providers/codex/subscription.rs @@ -7,6 +7,10 @@ use super::CodexApi; const SUBSCRIPTION_PATH: &str = "/subscriptions"; +/// Subscription metadata is optional enrichment. Usage remains usable when +/// the endpoint is unavailable, malformed, unauthorized, or points at a +/// custom backend. A successful empty cancellation response is the only +/// result allowed to clear dates on the fresh snapshot. pub(super) async fn enrich_subscription_metadata( api: &CodexApi, base_url: &str, @@ -17,16 +21,13 @@ pub(super) async fn enrich_subscription_metadata( if !crate::settings::Settings::load().codex_openai_web_extras() { return usage; } - match api - .fetch_subscription_metadata(base_url, access_token, account_id) - .await - { + match fetch_subscription_metadata(api, base_url, access_token, account_id).await { OpenAISubscriptionFetchResult::Success(metadata) => usage.with_subscription(metadata), OpenAISubscriptionFetchResult::Unavailable => usage, } } -pub(super) async fn fetch_subscription_metadata( +async fn fetch_subscription_metadata( api: &CodexApi, base_url: &str, access_token: &str, @@ -42,22 +43,15 @@ pub(super) async fn fetch_subscription_metadata( return OpenAISubscriptionFetchResult::Unavailable; } - let mut request = api - .client - .get(format!( - "{}{}", - base_url.trim_end_matches('/'), - SUBSCRIPTION_PATH - )) - .header("Authorization", format!("Bearer {access_token}")) - .header("User-Agent", "CodexBar") - .header("Accept", "application/json") + let request = api + .authed_get( + &format!("{}{}", base_url.trim_end_matches('/'), SUBSCRIPTION_PATH), + access_token, + account_id, + ) .header("Cache-Control", "no-cache, no-store, max-age=0") .header("Pragma", "no-cache") .timeout(Duration::from_secs(8)); - if let Some(account_id) = account_id.filter(|id| !id.is_empty()) { - request = request.header("ChatGPT-Account-Id", account_id); - } let Ok(response) = request.send().await else { return OpenAISubscriptionFetchResult::Unavailable; }; diff --git a/rust/src/providers/codex/weekly_reset.rs b/rust/src/providers/codex/weekly_reset.rs index d6acaa7292..077352c2fe 100644 --- a/rust/src/providers/codex/weekly_reset.rs +++ b/rust/src/providers/codex/weekly_reset.rs @@ -132,54 +132,43 @@ pub(super) fn scope_key(account_id: Option<&str>, auth_path: &Path) -> String { digest.iter().map(|byte| format!("{byte:02x}")).collect() } +/// An unreadable, malformed or other-version state file reads as absent. +fn read_state_file(path: &Path) -> Option { + crate::secure_file::read_string(path) + .ok() + .and_then(|raw| serde_json::from_str::(&raw).ok()) + .filter(|file| file.version == STATE_VERSION) +} + pub(super) fn load(scope: &str) -> AccountState { - let Some(path) = state_path() else { - return AccountState::default(); - }; - let Ok(raw) = crate::secure_file::read_string(&path) else { - return AccountState::default(); - }; - let Ok(file) = serde_json::from_str::(&raw) else { - return AccountState::default(); - }; - if file.version != STATE_VERSION { - return AccountState::default(); - } - file.accounts.get(scope).cloned().unwrap_or_default() + state_path() + .and_then(|path| read_state_file(&path)) + .and_then(|file| file.accounts.get(scope).cloned()) + .unwrap_or_default() } pub(super) fn save(scope: &str, state: &AccountState) { - let Some(path) = state_path() else { + let skipped = || { log_reset_diagnostic( "candidatePersistence", "skipped", ResetDiagnosticReason::StoreUnavailable, ); - return; }; - let mut file = crate::secure_file::read_string(&path) - .ok() - .and_then(|raw| serde_json::from_str::(&raw).ok()) - .filter(|file| file.version == STATE_VERSION) - .unwrap_or_else(|| StateFile { - version: STATE_VERSION, - accounts: HashMap::new(), - }); - file.accounts.insert(scope.to_string(), state.clone()); - let Some(parent) = path.parent() else { - log_reset_diagnostic( - "candidatePersistence", - "skipped", - ResetDiagnosticReason::StoreUnavailable, - ); + let Some(path) = state_path() else { + skipped(); return; }; - if std::fs::create_dir_all(parent).is_err() { - log_reset_diagnostic( - "candidatePersistence", - "skipped", - ResetDiagnosticReason::StoreUnavailable, - ); + let mut file = read_state_file(&path).unwrap_or_else(|| StateFile { + version: STATE_VERSION, + accounts: HashMap::new(), + }); + file.accounts.insert(scope.to_string(), state.clone()); + if path + .parent() + .is_none_or(|parent| std::fs::create_dir_all(parent).is_err()) + { + skipped(); return; } if let Ok(raw) = serde_json::to_string_pretty(&file) { @@ -413,91 +402,67 @@ fn maybe_store_delayed_candidate( exact_oauth: bool, observed_at: DateTime, ) { + match delayed_candidate_admission( + state, + initial, + confirmation, + confirmation_inventory, + exact_oauth, + observed_at, + ) { + Ok(candidate) => { + state.candidate = Some(candidate); + log_reset_diagnostic( + "candidateCreation", + "created", + ResetDiagnosticReason::CandidateCreated, + ); + } + Err(reason) => log_reset_diagnostic("candidateCreation", "rejected", reason), + } +} + +/// The checks run in a fixed order; the first failure is the logged reason. +fn delayed_candidate_admission( + state: &AccountState, + initial: &UsageSnapshot, + confirmation: &UsageSnapshot, + confirmation_inventory: Option<&CreditInventory>, + exact_oauth: bool, + observed_at: DateTime, +) -> Result { + use ResetDiagnosticReason as Reason; if !exact_oauth { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::SourceNotExactOAuth, - ); - return; + return Err(Reason::SourceNotExactOAuth); } if !plans_match(state.plan.as_deref(), initial, confirmation) { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::PlanMismatch, - ); - return; + return Err(Reason::PlanMismatch); } - let Some(previous_weekly) = state.published_weekly.as_ref() else { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::MissingPreviousSnapshot, - ); - return; - }; - let Some(initial_weekly) = weekly(initial) else { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::MissingWeeklyWindow, - ); - return; - }; - let Some(confirmation_weekly) = weekly(confirmation) else { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::MissingWeeklyWindow, - ); - return; - }; - let Some(previous_inventory) = state.credit_inventory.as_ref() else { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::MissingCreditInventory, - ); - return; - }; - let Some(confirmation_inventory) = confirmation_inventory else { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::MissingCreditInventory, - ); - return; - }; + let previous_weekly = state + .published_weekly + .as_ref() + .ok_or(Reason::MissingPreviousSnapshot)?; + let initial_weekly = weekly(initial).ok_or(Reason::MissingWeeklyWindow)?; + let confirmation_weekly = weekly(confirmation).ok_or(Reason::MissingWeeklyWindow)?; + let previous_inventory = state + .credit_inventory + .as_ref() + .ok_or(Reason::MissingCreditInventory)?; + let confirmation_inventory = confirmation_inventory.ok_or(Reason::MissingCreditInventory)?; if previous_inventory.available_count == 0 || previous_inventory != confirmation_inventory { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::ChangedCreditInventory, - ); - return; + return Err(Reason::ChangedCreditInventory); } if !supported_delayed_boundary(previous_weekly, initial_weekly) || !supported_delayed_boundary(previous_weekly, confirmation_weekly) { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::UnsupportedResetBoundary, - ); - return; + return Err(Reason::UnsupportedResetBoundary); } if boundary_distance_seconds(initial_weekly, confirmation_weekly).abs() >= RESET_TOLERANCE_SECONDS { - log_reset_diagnostic( - "candidateCreation", - "rejected", - ResetDiagnosticReason::InconsistentResetBoundary, - ); - return; + return Err(Reason::InconsistentResetBoundary); } - state.candidate = Some(DelayedCandidate { + Ok(DelayedCandidate { evidence_version: EVIDENCE_VERSION, first_observed_at: initial.updated_at, created_at: observed_at, @@ -505,12 +470,7 @@ fn maybe_store_delayed_candidate( weekly: confirmation_weekly.clone(), plan: confirmation.login_method.clone(), inventory: confirmation_inventory.clone(), - }); - log_reset_diagnostic( - "candidateCreation", - "created", - ResetDiagnosticReason::CandidateCreated, - ); + }) } fn delayed_candidate_decision( diff --git a/rust/src/providers/codex/weekly_reset/diagnostics.rs b/rust/src/providers/codex/weekly_reset/diagnostics.rs index c870b60a28..fec1c17999 100644 --- a/rust/src/providers/codex/weekly_reset/diagnostics.rs +++ b/rust/src/providers/codex/weekly_reset/diagnostics.rs @@ -1,52 +1,42 @@ -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(super) enum ResetDiagnosticReason { - CandidateCreated, - SourceNotExactOAuth, - MissingPreviousSnapshot, - MissingWeeklyWindow, - ResetThresholdMismatch, - InvalidResetBoundary, - InconsistentResetBoundary, - UnsupportedResetBoundary, - PlanMismatch, - PlanChanged, - MissingCreditInventory, - ChangedCreditInventory, - EvidenceVersionMismatch, - FutureCandidate, - ExpiredCandidate, - StaleObservation, - MinimumDelay, - ConfirmedObservation, - StoreUnavailable, - StoreRequested, -} +/// Each reason maps to one fixed, redacted log code. +macro_rules! reset_diagnostic_reasons { + ($($variant:ident => $code:literal,)*) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub(super) enum ResetDiagnosticReason { + $($variant,)* + } -impl ResetDiagnosticReason { - pub(super) const fn code(self) -> &'static str { - match self { - Self::CandidateCreated => "candidateCreated", - Self::SourceNotExactOAuth => "sourceNotExactOAuth", - Self::MissingPreviousSnapshot => "missingPreviousSnapshot", - Self::MissingWeeklyWindow => "missingWeeklyWindow", - Self::ResetThresholdMismatch => "resetThresholdMismatch", - Self::InvalidResetBoundary => "invalidResetBoundary", - Self::InconsistentResetBoundary => "inconsistentResetBoundary", - Self::UnsupportedResetBoundary => "unsupportedResetBoundary", - Self::PlanMismatch => "planMismatch", - Self::PlanChanged => "planChanged", - Self::MissingCreditInventory => "missingCreditInventory", - Self::ChangedCreditInventory => "changedCreditInventory", - Self::EvidenceVersionMismatch => "evidenceVersionMismatch", - Self::FutureCandidate => "futureCandidate", - Self::ExpiredCandidate => "expiredCandidate", - Self::StaleObservation => "staleObservation", - Self::MinimumDelay => "minimumDelay", - Self::ConfirmedObservation => "confirmedObservation", - Self::StoreUnavailable => "storeUnavailable", - Self::StoreRequested => "storeRequested", + impl ResetDiagnosticReason { + pub(super) const fn code(self) -> &'static str { + match self { + $(Self::$variant => $code,)* + } + } } - } + }; +} + +reset_diagnostic_reasons! { + CandidateCreated => "candidateCreated", + SourceNotExactOAuth => "sourceNotExactOAuth", + MissingPreviousSnapshot => "missingPreviousSnapshot", + MissingWeeklyWindow => "missingWeeklyWindow", + ResetThresholdMismatch => "resetThresholdMismatch", + InvalidResetBoundary => "invalidResetBoundary", + InconsistentResetBoundary => "inconsistentResetBoundary", + UnsupportedResetBoundary => "unsupportedResetBoundary", + PlanMismatch => "planMismatch", + PlanChanged => "planChanged", + MissingCreditInventory => "missingCreditInventory", + ChangedCreditInventory => "changedCreditInventory", + EvidenceVersionMismatch => "evidenceVersionMismatch", + FutureCandidate => "futureCandidate", + ExpiredCandidate => "expiredCandidate", + StaleObservation => "staleObservation", + MinimumDelay => "minimumDelay", + ConfirmedObservation => "confirmedObservation", + StoreUnavailable => "storeUnavailable", + StoreRequested => "storeRequested", } pub(super) fn log_reset_diagnostic( diff --git a/rust/src/providers/codex/weekly_reset/tests.rs b/rust/src/providers/codex/weekly_reset/tests.rs index 6848b8ff83..5058638c91 100644 --- a/rust/src/providers/codex/weekly_reset/tests.rs +++ b/rust/src/providers/codex/weekly_reset/tests.rs @@ -11,6 +11,21 @@ fn reset_diagnostic_codes_are_fixed_and_redacted() { ResetDiagnosticReason::ExpiredCandidate.code(), ResetDiagnosticReason::ChangedCreditInventory.code(), ResetDiagnosticReason::StoreRequested.code(), + ResetDiagnosticReason::MissingPreviousSnapshot.code(), + ResetDiagnosticReason::MissingWeeklyWindow.code(), + ResetDiagnosticReason::ResetThresholdMismatch.code(), + ResetDiagnosticReason::InvalidResetBoundary.code(), + ResetDiagnosticReason::InconsistentResetBoundary.code(), + ResetDiagnosticReason::UnsupportedResetBoundary.code(), + ResetDiagnosticReason::PlanMismatch.code(), + ResetDiagnosticReason::PlanChanged.code(), + ResetDiagnosticReason::MissingCreditInventory.code(), + ResetDiagnosticReason::EvidenceVersionMismatch.code(), + ResetDiagnosticReason::FutureCandidate.code(), + ResetDiagnosticReason::StaleObservation.code(), + ResetDiagnosticReason::MinimumDelay.code(), + ResetDiagnosticReason::ConfirmedObservation.code(), + ResetDiagnosticReason::StoreUnavailable.code(), ]; assert_eq!( codes, @@ -20,6 +35,21 @@ fn reset_diagnostic_codes_are_fixed_and_redacted() { "expiredCandidate", "changedCreditInventory", "storeRequested", + "missingPreviousSnapshot", + "missingWeeklyWindow", + "resetThresholdMismatch", + "invalidResetBoundary", + "inconsistentResetBoundary", + "unsupportedResetBoundary", + "planMismatch", + "planChanged", + "missingCreditInventory", + "evidenceVersionMismatch", + "futureCandidate", + "staleObservation", + "minimumDelay", + "confirmedObservation", + "storeUnavailable", ] ); assert!(codes.iter().all(|code| { @@ -45,6 +75,37 @@ fn snapshot(used: f64, reset_days: i64, captured_minutes: i64) -> UsageSnapshot snapshot } +/// Exact-OAuth confirmation at `now()` where both fetches saw the same inventory. +fn confirm( + state: &mut AccountState, + initial: &UsageSnapshot, + confirmation: &UsageSnapshot, + inventory: Option<&CreditInventory>, +) -> ConfirmationDecision { + confirmation_decision( + state, + initial, + inventory, + confirmation, + inventory, + true, + now(), + ) +} + +/// A candidate admitted at `now()` against the `credit-a` inventory. +fn pending_candidate(weekly: RateWindow, plan: &str) -> DelayedCandidate { + DelayedCandidate { + evidence_version: EVIDENCE_VERSION, + first_observed_at: now(), + created_at: now(), + snapshot_updated_at: now(), + weekly, + plan: Some(plan.to_string()), + inventory: inventory("credit-a"), + } +} + fn inventory(id: &str) -> CreditInventory { CreditInventory { available_count: 1, @@ -120,15 +181,7 @@ fn early_low_usage_requires_confirmation_without_spending_credit() { ); let confirmation = snapshot(0.0, 9, 2); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Preserve ); assert!(state.candidate.is_some()); @@ -142,15 +195,7 @@ fn delayed_candidate_publishes_after_sixty_seconds_and_expires_after_thirty_minu let confirmation = snapshot(0.0, 9, 2); let inv = inventory("credit-a"); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Preserve ); let current = snapshot(0.0, 9, 3); @@ -193,15 +238,10 @@ fn delayed_candidate_publishes_after_sixty_seconds_and_expires_after_thirty_minu #[test] fn credits_only_refresh_retains_candidate_and_account_scope_hashes_differ() { let mut state = baseline(); - state.candidate = Some(DelayedCandidate { - evidence_version: EVIDENCE_VERSION, - first_observed_at: now(), - created_at: now(), - snapshot_updated_at: now(), - weekly: snapshot(0.0, 9, 1).secondary.unwrap(), - plan: Some("ChatGPT Pro".to_string()), - inventory: inventory("credit-a"), - }); + state.candidate = Some(pending_candidate( + snapshot(0.0, 9, 1).secondary.unwrap(), + "ChatGPT Pro", + )); let mut credits_only = UsageSnapshot::new(RateWindow::new(20.0)); credits_only.updated_at = now() + chrono::Duration::minutes(1); // A credits-only refresh has no weekly window and may omit both plan and @@ -230,15 +270,10 @@ fn credits_only_refresh_retains_candidate_and_account_scope_hashes_differ() { #[test] fn credits_only_refresh_candidate_survives_state_reload_until_full_usage() { let mut state = baseline(); - state.candidate = Some(DelayedCandidate { - evidence_version: EVIDENCE_VERSION, - first_observed_at: now(), - created_at: now(), - snapshot_updated_at: now(), - weekly: snapshot(0.0, 9, 1).secondary.unwrap(), - plan: Some("ChatGPT Pro".to_string()), - inventory: inventory("credit-a"), - }); + state.candidate = Some(pending_candidate( + snapshot(0.0, 9, 1).secondary.unwrap(), + "ChatGPT Pro", + )); let candidate_before = serde_json::to_value(&state.candidate).unwrap(); let mut credits_only = UsageSnapshot::new(RateWindow::new(20.0)); credits_only.updated_at = now() + chrono::Duration::minutes(1); @@ -307,15 +342,7 @@ fn consumed_credit_allows_immediate_confirmation() { credits: Vec::new(), }; assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&consumed), - &confirmation, - Some(&consumed), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&consumed)), ConfirmationDecision::Publish ); } @@ -353,15 +380,7 @@ fn rolling_state() -> AccountState { let initial = rolling_snapshot(0.0, 7 * 24 * 60, 1, WEEK_SECONDS - 1); let confirmation = rolling_snapshot(0.0, 7 * 24 * 60, 2, WEEK_SECONDS - 2); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Preserve ); assert!(state.candidate.is_some()); diff --git a/rust/src/providers/codex/weekly_reset/tests/plan_change.rs b/rust/src/providers/codex/weekly_reset/tests/plan_change.rs index 2a22edb68e..ba8f45809f 100644 --- a/rust/src/providers/codex/weekly_reset/tests/plan_change.rs +++ b/rust/src/providers/codex/weekly_reset/tests/plan_change.rs @@ -36,15 +36,7 @@ fn plan_upgrade_starts_a_new_baseline_and_publishes_the_new_plan() { assert!(state.credit_inventory.is_none()); assert!(state.candidate.is_none()); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Publish ); } @@ -63,15 +55,7 @@ fn same_plan_near_zero_reading_keeps_the_previous_weekly_pinned() { ); assert!(state.published_weekly.is_some()); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Preserve ); } @@ -93,15 +77,7 @@ fn plan_upgrade_does_not_pin_the_previous_plan_weekly_window() { #[test] fn plan_change_discards_a_pending_candidate() { let mut state = plus_baseline(); - state.candidate = Some(DelayedCandidate { - evidence_version: EVIDENCE_VERSION, - first_observed_at: now(), - created_at: now(), - snapshot_updated_at: now(), - weekly: RateWindow::new(0.0), - plan: Some("ChatGPT Plus".to_string()), - inventory: inventory("credit-a"), - }); + state.candidate = Some(pending_candidate(RateWindow::new(0.0), "ChatGPT Plus")); let current = plan_snapshot(Some("ChatGPT Pro"), 5.0, 10); assert_eq!( initial_decision(&mut state, ¤t, None, true, now()), @@ -181,15 +157,7 @@ fn near_zero_confirmation_must_report_the_initial_plan() { let initial = plan_snapshot(Some("ChatGPT Pro"), 0.0, 10); let confirmation = plan_snapshot(confirmation_plan, 0.0, 11); assert_eq!( - confirmation_decision( - &mut state, - &initial, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, &initial, &confirmation, Some(&inv)), ConfirmationDecision::Preserve, "{confirmation_plan:?} baseline {has_baseline}" ); @@ -204,7 +172,7 @@ fn nonzero_confirmation_can_publish_its_own_plan() { let initial = plan_snapshot(Some("ChatGPT Pro"), 0.0, 10); let confirmation = plan_snapshot(Some("ChatGPT Plus"), 5.0, 11); assert_eq!( - confirmation_decision(&mut state, &initial, None, &confirmation, None, true, now()), + confirm(&mut state, &initial, &confirmation, None), ConfirmationDecision::Publish ); } @@ -253,15 +221,7 @@ fn new_plan_replaces_the_previous_plan_quota_baseline() { InitialDecision::Publish => current.clone(), InitialDecision::RequiresConfirmation => { assert_eq!( - confirmation_decision( - &mut state, - ¤t, - Some(&inv), - &confirmation, - Some(&inv), - true, - now(), - ), + confirm(&mut state, ¤t, &confirmation, Some(&inv)), ConfirmationDecision::Publish, "{used}% confirmation" ); @@ -299,15 +259,7 @@ fn same_or_unknown_plan_cannot_discard_previous_quota_evidence() { "{plan:?}" ); assert_eq!( - confirmation_decision( - &mut state, - ¤t, - Some(&inv), - ¤t, - Some(&inv), - true, - now(), - ), + confirm(&mut state, ¤t, ¤t, Some(&inv)), ConfirmationDecision::Preserve, "{plan:?}" ); diff --git a/rust/src/providers/copilot/api.rs b/rust/src/providers/copilot/api.rs index 9b97b0d7b5..dbecaddbf1 100755 --- a/rust/src/providers/copilot/api.rs +++ b/rust/src/providers/copilot/api.rs @@ -64,35 +64,17 @@ impl CopilotApi { github_host: Option<&str>, seat_credit_entitlement: Option, ) -> Result { - let api_url = copilot_usage_url(github_host); - let response = self + let request = self .client - .get(api_url) + .get(copilot_usage_url(github_host)) .header("Authorization", format!("token {}", token.trim())) .header("Accept", "application/json") .header("Editor-Version", "vscode/1.96.2") .header("Editor-Plugin-Version", "copilot-chat/0.26.7") .header("User-Agent", "GitHubCopilotChat/0.26.7") - .header("X-Github-Api-Version", "2025-04-01") - .send() - .await - .map_err(|e| ProviderError::Other(format!("Request failed: {}", e)))?; - - if response.status() == 401 || response.status() == 403 { - return Err(ProviderError::AuthRequired); - } - - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "GitHub Copilot usage endpoint returned {}", - response.status() - ))); - } - - let usage_response: CopilotUsageResponse = response - .json() - .await - .map_err(|e| ProviderError::Parse(e.to_string()))?; + .header("X-Github-Api-Version", "2025-04-01"); + let usage_response: CopilotUsageResponse = + send_json(request, "GitHub Copilot usage endpoint").await?; snapshot_from_response_with_seat_entitlement(usage_response, seat_credit_entitlement) } @@ -103,33 +85,14 @@ impl CopilotApi { token: &str, github_host: Option<&str>, ) -> Result { - let url = github_api_url(github_host, GITHUB_USER_PATH); - let response = self + let request = self .client - .get(url) + .get(github_api_url(github_host, GITHUB_USER_PATH)) .header("Authorization", format!("token {}", token.trim())) .header("Accept", "application/vnd.github+json") .header("User-Agent", "Win-CodexBar") - .header("X-GitHub-Api-Version", "2022-11-28") - .send() - .await - .map_err(|e| ProviderError::Other(format!("Request failed: {}", e)))?; - - if response.status() == 401 || response.status() == 403 { - return Err(ProviderError::AuthRequired); - } - - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "GitHub identity endpoint returned {}", - response.status() - ))); - } - - response - .json() - .await - .map_err(|e| ProviderError::Parse(e.to_string())) + .header("X-GitHub-Api-Version", "2022-11-28"); + send_json(request, "GitHub identity endpoint").await } /// Resolve the Copilot OAuth token from settings/legacy API key, GitHub @@ -237,6 +200,30 @@ impl Default for CopilotApi { } } +/// 401/403 mean the token is rejected; any other non-success names `endpoint`. +async fn send_json( + request: reqwest::RequestBuilder, + endpoint: &str, +) -> Result { + let response = request + .send() + .await + .map_err(|e| ProviderError::Other(format!("Request failed: {}", e)))?; + let status = response.status(); + if status == 401 || status == 403 { + return Err(ProviderError::AuthRequired); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "{endpoint} returned {status}" + ))); + } + response + .json() + .await + .map_err(|e| ProviderError::Parse(e.to_string())) +} + // --- API Response Types --- #[derive(Debug, Deserialize)] @@ -521,10 +508,7 @@ impl CopilotUsageResponse { fn credits_used_counter(&self) -> Option { let mut chat: Option = None; let mut first: Option = None; - for (key, value) in &self.quota_snapshots.entries { - let Ok(snapshot) = serde_json::from_value::(value.clone()) else { - continue; - }; + for (key, snapshot) in self.snapshots() { let Some(credits) = snapshot.credits_used else { continue; }; @@ -534,64 +518,50 @@ impl CopilotUsageResponse { match classify_quota_kind(key, snapshot.quota_id.as_deref().unwrap_or_default()) { CopilotQuotaKind::Premium => return Some(credits), CopilotQuotaKind::Chat => { - if chat.is_none() { - chat = Some(credits); - } + chat.get_or_insert(credits); } _ => { - if first.is_none() { - first = Some(credits); - } + first.get_or_insert(credits); } } } chat.or(first) } + /// Quota snapshot entries that decode; malformed entries are skipped. + fn snapshots(&self) -> impl Iterator { + self.quota_snapshots + .entries + .iter() + .filter_map(|(key, value)| { + serde_json::from_value::(value.clone()) + .ok() + .map(|snapshot| (key, snapshot)) + }) + } + fn usable_quotas(&self, reset: Option>) -> UsableQuotas { let mut quotas = UsableQuotas::default(); - for (key, value) in &self.quota_snapshots.entries { - let Ok(snapshot) = serde_json::from_value::(value.clone()) else { - continue; - }; + for (key, snapshot) in self.snapshots() { let Some(quota) = UsableQuota::from_snapshot(key, snapshot) else { continue; }; - - match quota.kind { - CopilotQuotaKind::Premium => { - if quotas.first.is_none() { - quotas.first = Some(quota.clone()); - } - if quotas.premium.is_none() { - quotas.premium = Some(quota); - } - } - CopilotQuotaKind::Chat => { - if quotas.first.is_none() { - quotas.first = Some(quota.clone()); - } - if quotas.chat.is_none() { - quotas.chat = Some(quota); - } - } - CopilotQuotaKind::Completions => { - if quotas.first.is_none() { - quotas.first = Some(quota.clone()); - } - if quotas.completions.is_none() { - quotas.completions = Some(quota); - } - } + let slot = match quota.kind { + CopilotQuotaKind::Premium => &mut quotas.premium, + CopilotQuotaKind::Chat => &mut quotas.chat, + CopilotQuotaKind::Completions => &mut quotas.completions, CopilotQuotaKind::Other => { quotas.extra.push(NamedRateWindow::new( quota.id.clone(), quota.title.clone(), quota.to_rate_window(reset), )); + continue; } - } + }; + slot.get_or_insert_with(|| quota.clone()); + quotas.first.get_or_insert(quota); } let completions = UsableQuota::from_limited( @@ -806,500 +776,5 @@ where } #[cfg(test)] -mod tests { - use super::*; - - fn parse_snapshot(json: &str) -> UsageSnapshot { - let response: CopilotUsageResponse = serde_json::from_str(json).unwrap(); - snapshot_from_response_with_seat_entitlement(response, None).unwrap() - } - - fn parse_snapshot_result(json: &str) -> Result { - let response: CopilotUsageResponse = serde_json::from_str(json).unwrap(); - snapshot_from_response_with_seat_entitlement(response, None) - } - - #[test] - fn paid_plan_parses_premium_and_chat_quotas() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "pro", - "quota_reset_date": "2026-06-01", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 300, - "remaining": 240, - "percent_remaining": 80, - "quota_id": "premium_interactions" - }, - "chat": { - "entitlement": 1000, - "remaining": 900, - "percent_remaining": 90, - "quota_id": "chat" - } - } - }"#, - ); - - assert_eq!(usage.login_method.as_deref(), Some("Copilot Pro")); - assert!((usage.primary.used_percent - 20.0).abs() < 0.001); - assert!((usage.secondary.unwrap().used_percent - 10.0).abs() < 0.001); - } - - #[test] - fn limited_user_quotas_parse_free_schema() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "free", - "monthly_quotas": { - "completions": 2000, - "chat": "50" - }, - "limited_user_quotas": { - "completions": "1000", - "chat": 10 - } - }"#, - ); - - assert_eq!(usage.login_method.as_deref(), Some("Copilot Free")); - assert!((usage.primary.used_percent - 50.0).abs() < 0.001); - assert!((usage.secondary.unwrap().used_percent - 80.0).abs() < 0.001); - } - - #[test] - fn derives_missing_percent_and_accepts_numeric_strings() { - let usage = parse_snapshot( - r#"{ - "quota_snapshots": { - "premium_interactions": { - "entitlement": "100", - "remaining": "25", - "quota_id": "premium_interactions" - } - } - }"#, - ); - - assert!((usage.primary.used_percent - 75.0).abs() < 0.001); - } - - #[test] - fn ignores_placeholders_and_does_not_promote_chat_to_premium() { - let usage = parse_snapshot( - r#"{ - "quota_snapshots": { - "premium_interactions": { - "percent_remaining": 0, - "quota_id": "" - }, - "chat": { - "entitlement": 100, - "remaining": 75, - "percent_remaining": 75, - "quota_id": "chat" - } - } - }"#, - ); - - assert!((usage.primary.used_percent - 25.0).abs() < 0.001); - assert!(usage.secondary.is_none()); - } - - #[test] - fn drops_business_token_billing_zero_entitlement_quotas() { - let err = parse_snapshot_result( - r#"{ - "copilot_plan": "business", - "token_based_billing": true, - "quota_snapshots": { - "premium_interactions": { - "entitlement": 0, - "remaining": 0, - "percent_remaining": 100, - "quota_id": "premium_interactions" - }, - "chat": { - "entitlement": 0, - "remaining": 0, - "percent_remaining": 100, - "quota_id": "chat" - }, - "completions": { - "entitlement": 0, - "remaining": 0, - "percent_remaining": 100, - "quota_id": "completions" - } - } - }"#, - ) - .unwrap_err(); - - assert!( - err.to_string() - .contains("token-based billing usage is unavailable") - ); - } - - #[test] - fn keeps_percent_only_quota_snapshots_available() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "business", - "quota_snapshots": { - "chat": { - "percent_remaining": 40, - "quota_id": "chat" - } - } - }"#, - ); - - assert_eq!(usage.login_method.as_deref(), Some("Copilot Business")); - assert!((usage.primary.used_percent - 60.0).abs() < 0.001); - assert!(usage.secondary.is_none()); - } - - #[test] - fn keeps_fully_consumed_positive_entitlement_quota() { - let usage = parse_snapshot( - r#"{ - "quota_snapshots": { - "premium_interactions": { - "entitlement": 500, - "remaining": 0, - "percent_remaining": 0, - "quota_id": "premium_interactions" - } - } - }"#, - ); - - assert!((usage.primary.used_percent - 100.0).abs() < 0.001); - } - - #[test] - fn keeps_additional_budget_as_extra_window() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "pro", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 500, - "remaining": 250, - "quota_id": "premium_interactions" - }, - "additional_budget": { - "entitlement": 100, - "remaining": 25, - "quota_id": "additional_budget" - } - } - }"#, - ); - - assert!((usage.primary.used_percent - 50.0).abs() < 0.001); - assert_eq!(usage.extra_rate_windows.len(), 1); - assert_eq!(usage.extra_rate_windows[0].id, "additional-budget"); - assert_eq!(usage.extra_rate_windows[0].title, "Additional Budget"); - assert!((usage.extra_rate_windows[0].window.used_percent - 75.0).abs() < 0.001); - } - - #[test] - fn preserves_over_quota_percent_remaining() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "pro", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 500, - "remaining": -75, - "percent_remaining": -15, - "quota_id": "premium_interactions" - } - } - }"#, - ); - - assert_eq!(usage.login_method.as_deref(), Some("Copilot Pro")); - assert!((usage.primary.used_percent - 115.0).abs() < 0.001); - assert_eq!( - usage.primary.reset_description.as_deref(), - Some("115% used") - ); - assert!(usage.primary.is_exhausted()); - } - - #[test] - fn derives_over_quota_percent_from_negative_remaining() { - let usage = parse_snapshot( - r#"{ - "quota_snapshots": { - "chat": { - "entitlement": 500, - "remaining": -75, - "quota_id": "chat" - } - } - }"#, - ); - - assert!((usage.primary.used_percent - 115.0).abs() < 0.001); - assert_eq!( - usage.primary.reset_description.as_deref(), - Some("115% used") - ); - } - - #[test] - fn normalizes_enterprise_hosts() { - assert_eq!( - normalized_api_host(Some("github.com")), - "api.github.com".to_string() - ); - assert_eq!( - normalized_api_host(Some("github.example.com")), - "api.github.example.com".to_string() - ); - assert_eq!( - normalized_api_host(Some("api.github.example.com")), - "api.github.example.com".to_string() - ); - } - - // ── A15: credits_used counter for token-billed seats (upstream #2613) ─── - - #[test] - fn decodes_credits_used_as_number_or_string() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "business", - "quota_reset_date": "2026-06-01", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 300, - "remaining": 240, - "percent_remaining": 80, - "quota_id": "premium_interactions", - "credits_used": "1234.56" - } - } - }"#, - ); - let extra = &usage.extra_rate_windows; - assert!( - extra.iter().any(|w| w.id == "ai-credits" - && w.window.reset_description.as_deref() == Some("1234.56 AI credits used")), - "{extra:?}" - ); - } - - #[test] - fn configured_seat_allowance_adds_a_numeric_credit_window() { - let response: CopilotUsageResponse = serde_json::from_str( - r#"{ - "copilot_plan": "business", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 300, - "remaining": 240, - "percent_remaining": 80, - "quota_id": "premium_interactions", - "credits_used": 50 - } - } - }"#, - ) - .unwrap(); - let usage = snapshot_from_response_with_seat_entitlement(response, Some(200.0)).unwrap(); - - let seat = usage - .extra_rate_windows - .iter() - .find(|window| window.id == SEAT_CREDIT_WINDOW_ID) - .expect("configured seat-credit window"); - assert!((seat.window.used_percent - 25.0).abs() < 0.001); - assert!(!seat.window.is_informational); - assert_eq!(seat.title, "Credits used"); - } - - #[test] - fn missing_primary_quota_is_informational_when_seat_credit_is_available() { - let response: CopilotUsageResponse = serde_json::from_str( - r#"{ - "copilot_plan": "business", - "quota_snapshots": { - "additional_budget": { - "credits_used": 50 - } - } - }"#, - ) - .unwrap(); - let usage = snapshot_from_response_with_seat_entitlement(response, Some(200.0)).unwrap(); - - assert!(usage.primary.is_informational); - assert!( - usage - .extra_rate_windows - .iter() - .any(|window| window.id == SEAT_CREDIT_WINDOW_ID) - ); - } - - #[test] - fn non_finite_derived_seat_credit_percentage_is_omitted() { - let response: CopilotUsageResponse = serde_json::from_str( - r#"{ - "copilot_plan": "business", - "quota_snapshots": { - "premium_interactions": { - "credits_used": 1e308 - } - } - }"#, - ) - .unwrap(); - let usage = snapshot_from_response_with_seat_entitlement(response, Some(1e-308)).unwrap(); - - assert!( - usage - .extra_rate_windows - .iter() - .all(|window| window.id != SEAT_CREDIT_WINDOW_ID) - ); - } - - #[test] - fn invalid_seat_allowance_keeps_credit_progress_unknown() { - let response: CopilotUsageResponse = serde_json::from_str( - r#"{ - "copilot_plan": "business", - "token_based_billing": true, - "quota_snapshots": { - "premium_interactions": { - "entitlement": 0, - "remaining": 0, - "credits_used": 50 - } - } - }"#, - ) - .unwrap(); - let usage = snapshot_from_response_with_seat_entitlement(response, Some(0.0)).unwrap(); - - assert!(usage.primary.is_informational); - assert!( - usage - .extra_rate_windows - .iter() - .all(|window| window.id != SEAT_CREDIT_WINDOW_ID) - ); - } - - #[test] - fn zero_entitlement_business_seat_surfaces_credits_counter() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "business", - "token_based_billing": true, - "quota_reset_date": "2026-06-01", - "quota_snapshots": { - "premium_interactions": { - "entitlement": 0, - "remaining": 0, - "percent_remaining": 100, - "quota_id": "premium_interactions", - "credits_used": 1234 - } - } - }"#, - ); - // Not an error anymore: informational counter row without a fake bar. - assert!(usage.primary.is_informational); - assert_eq!( - usage.primary.reset_description.as_deref(), - Some("1234 AI credits used") - ); - assert!(usage.primary.resets_at.is_some()); - assert_eq!(usage.login_method.as_deref(), Some("Copilot Business")); - } - - #[test] - fn placeholder_snapshot_still_carries_its_credits_counter() { - // Upstream carriesCreditsCounter: a placeholder cannot become a window, - // but its absolute counter is real consumption and must survive. - let usage = parse_snapshot( - r#"{ - "copilot_plan": "business", - "token_based_billing": true, - "quota_snapshots": { - "premium_interactions": { - "entitlement": 0, - "remaining": 0, - "percent_remaining": 0, - "quota_id": "", - "placeholder": true, - "credits_used": 42.5 - } - } - }"#, - ); - assert!(usage.primary.is_informational); - assert_eq!( - usage.primary.reset_description.as_deref(), - Some("42.50 AI credits used") - ); - } - - #[test] - fn premium_credits_counter_wins_over_chat() { - let usage = parse_snapshot( - r#"{ - "copilot_plan": "pro", - "quota_snapshots": { - "chat": { - "entitlement": 100, - "remaining": 75, - "percent_remaining": 75, - "quota_id": "chat", - "credits_used": 1 - }, - "premium_interactions": { - "entitlement": 300, - "remaining": 240, - "percent_remaining": 80, - "quota_id": "premium_interactions", - "credits_used": 7 - } - } - }"#, - ); - let credits_row = usage - .extra_rate_windows - .iter() - .find(|w| w.id == "ai-credits") - .expect("ai-credits row"); - assert_eq!( - credits_row.window.reset_description.as_deref(), - Some("7 AI credits used") - ); - // Windows still render normally next to the counter. - assert!(credits_row.window.is_informational); - assert!((usage.primary.used_percent - 20.0).abs() < 0.001); - } - - #[test] - fn business_seat_without_credits_keeps_existing_error() { - let err = parse_snapshot_result( - r#"{ - "copilot_plan": "business", - "token_based_billing": true - }"#, - ); - assert!(err.is_err()); - } -} +#[path = "api_tests.rs"] +mod tests; diff --git a/rust/src/providers/copilot/api_tests.rs b/rust/src/providers/copilot/api_tests.rs new file mode 100644 index 0000000000..88f62a8dc9 --- /dev/null +++ b/rust/src/providers/copilot/api_tests.rs @@ -0,0 +1,494 @@ +use super::*; + +fn parse_with_seat(json: &str, seat: Option) -> Result { + let response: CopilotUsageResponse = serde_json::from_str(json).unwrap(); + snapshot_from_response_with_seat_entitlement(response, seat) +} + +fn parse_snapshot(json: &str) -> UsageSnapshot { + parse_with_seat(json, None).unwrap() +} + +fn parse_snapshot_result(json: &str) -> Result { + parse_with_seat(json, None) +} + +fn seat_window(usage: &UsageSnapshot) -> Option<&NamedRateWindow> { + usage + .extra_rate_windows + .iter() + .find(|window| window.id == SEAT_CREDIT_WINDOW_ID) +} + +fn assert_used(window: &RateWindow, expected: f64) { + assert!( + (window.used_percent - expected).abs() < 0.001, + "{} != {expected}", + window.used_percent + ); +} + +#[test] +fn paid_plan_parses_premium_and_chat_quotas() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "pro", + "quota_reset_date": "2026-06-01", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 300, + "remaining": 240, + "percent_remaining": 80, + "quota_id": "premium_interactions" + }, + "chat": { + "entitlement": 1000, + "remaining": 900, + "percent_remaining": 90, + "quota_id": "chat" + } + } + }"#, + ); + + assert_eq!(usage.login_method.as_deref(), Some("Copilot Pro")); + assert_used(&usage.primary, 20.0); + assert_used(&usage.secondary.unwrap(), 10.0); +} + +#[test] +fn limited_user_quotas_parse_free_schema() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "free", + "monthly_quotas": { + "completions": 2000, + "chat": "50" + }, + "limited_user_quotas": { + "completions": "1000", + "chat": 10 + } + }"#, + ); + + assert_eq!(usage.login_method.as_deref(), Some("Copilot Free")); + assert_used(&usage.primary, 50.0); + assert_used(&usage.secondary.unwrap(), 80.0); +} + +#[test] +fn derives_missing_percent_and_accepts_numeric_strings() { + let usage = parse_snapshot( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "entitlement": "100", + "remaining": "25", + "quota_id": "premium_interactions" + } + } + }"#, + ); + + assert_used(&usage.primary, 75.0); +} + +#[test] +fn ignores_placeholders_and_does_not_promote_chat_to_premium() { + let usage = parse_snapshot( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "percent_remaining": 0, + "quota_id": "" + }, + "chat": { + "entitlement": 100, + "remaining": 75, + "percent_remaining": 75, + "quota_id": "chat" + } + } + }"#, + ); + + assert_used(&usage.primary, 25.0); + assert!(usage.secondary.is_none()); +} + +#[test] +fn drops_business_token_billing_zero_entitlement_quotas() { + let err = parse_snapshot_result( + r#"{ + "copilot_plan": "business", + "token_based_billing": true, + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 100, + "quota_id": "premium_interactions" + }, + "chat": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 100, + "quota_id": "chat" + }, + "completions": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 100, + "quota_id": "completions" + } + } + }"#, + ) + .unwrap_err(); + + assert!( + err.to_string() + .contains("token-based billing usage is unavailable") + ); +} + +#[test] +fn keeps_percent_only_quota_snapshots_available() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "business", + "quota_snapshots": { + "chat": { + "percent_remaining": 40, + "quota_id": "chat" + } + } + }"#, + ); + + assert_eq!(usage.login_method.as_deref(), Some("Copilot Business")); + assert_used(&usage.primary, 60.0); + assert!(usage.secondary.is_none()); +} + +#[test] +fn keeps_fully_consumed_positive_entitlement_quota() { + let usage = parse_snapshot( + r#"{ + "quota_snapshots": { + "premium_interactions": { + "entitlement": 500, + "remaining": 0, + "percent_remaining": 0, + "quota_id": "premium_interactions" + } + } + }"#, + ); + + assert_used(&usage.primary, 100.0); +} + +#[test] +fn keeps_additional_budget_as_extra_window() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "pro", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 500, + "remaining": 250, + "quota_id": "premium_interactions" + }, + "additional_budget": { + "entitlement": 100, + "remaining": 25, + "quota_id": "additional_budget" + } + } + }"#, + ); + + assert_used(&usage.primary, 50.0); + assert_eq!(usage.extra_rate_windows.len(), 1); + assert_eq!(usage.extra_rate_windows[0].id, "additional-budget"); + assert_eq!(usage.extra_rate_windows[0].title, "Additional Budget"); + assert_used(&usage.extra_rate_windows[0].window, 75.0); +} + +#[test] +fn preserves_over_quota_percent_remaining() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "pro", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 500, + "remaining": -75, + "percent_remaining": -15, + "quota_id": "premium_interactions" + } + } + }"#, + ); + + assert_eq!(usage.login_method.as_deref(), Some("Copilot Pro")); + assert_used(&usage.primary, 115.0); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("115% used") + ); + assert!(usage.primary.is_exhausted()); +} + +#[test] +fn derives_over_quota_percent_from_negative_remaining() { + let usage = parse_snapshot( + r#"{ + "quota_snapshots": { + "chat": { + "entitlement": 500, + "remaining": -75, + "quota_id": "chat" + } + } + }"#, + ); + + assert_used(&usage.primary, 115.0); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("115% used") + ); +} + +#[test] +fn normalizes_enterprise_hosts() { + assert_eq!( + normalized_api_host(Some("github.com")), + "api.github.com".to_string() + ); + assert_eq!( + normalized_api_host(Some("github.example.com")), + "api.github.example.com".to_string() + ); + assert_eq!( + normalized_api_host(Some("api.github.example.com")), + "api.github.example.com".to_string() + ); +} + +// ── A15: credits_used counter for token-billed seats (upstream #2613) ─── + +#[test] +fn decodes_credits_used_as_number_or_string() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "business", + "quota_reset_date": "2026-06-01", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 300, + "remaining": 240, + "percent_remaining": 80, + "quota_id": "premium_interactions", + "credits_used": "1234.56" + } + } + }"#, + ); + let extra = &usage.extra_rate_windows; + assert!( + extra.iter().any(|w| w.id == "ai-credits" + && w.window.reset_description.as_deref() == Some("1234.56 AI credits used")), + "{extra:?}" + ); +} + +#[test] +fn configured_seat_allowance_adds_a_numeric_credit_window() { + let usage = parse_with_seat( + r#"{ + "copilot_plan": "business", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 300, + "remaining": 240, + "percent_remaining": 80, + "quota_id": "premium_interactions", + "credits_used": 50 + } + } + }"#, + Some(200.0), + ) + .unwrap(); + + let seat = seat_window(&usage).expect("configured seat-credit window"); + assert_used(&seat.window, 25.0); + assert!(!seat.window.is_informational); + assert_eq!(seat.title, "Credits used"); +} + +#[test] +fn missing_primary_quota_is_informational_when_seat_credit_is_available() { + let usage = parse_with_seat( + r#"{ + "copilot_plan": "business", + "quota_snapshots": { + "additional_budget": { + "credits_used": 50 + } + } + }"#, + Some(200.0), + ) + .unwrap(); + + assert!(usage.primary.is_informational); + assert!(seat_window(&usage).is_some()); +} + +#[test] +fn non_finite_derived_seat_credit_percentage_is_omitted() { + let usage = parse_with_seat( + r#"{ + "copilot_plan": "business", + "quota_snapshots": { + "premium_interactions": { + "credits_used": 1e308 + } + } + }"#, + Some(1e-308), + ) + .unwrap(); + + assert!(seat_window(&usage).is_none()); +} + +#[test] +fn invalid_seat_allowance_keeps_credit_progress_unknown() { + let usage = parse_with_seat( + r#"{ + "copilot_plan": "business", + "token_based_billing": true, + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 0, + "credits_used": 50 + } + } + }"#, + Some(0.0), + ) + .unwrap(); + + assert!(usage.primary.is_informational); + assert!(seat_window(&usage).is_none()); +} + +#[test] +fn zero_entitlement_business_seat_surfaces_credits_counter() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "business", + "token_based_billing": true, + "quota_reset_date": "2026-06-01", + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 100, + "quota_id": "premium_interactions", + "credits_used": 1234 + } + } + }"#, + ); + // Not an error anymore: informational counter row without a fake bar. + assert!(usage.primary.is_informational); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("1234 AI credits used") + ); + assert!(usage.primary.resets_at.is_some()); + assert_eq!(usage.login_method.as_deref(), Some("Copilot Business")); +} + +#[test] +fn placeholder_snapshot_still_carries_its_credits_counter() { + // Upstream carriesCreditsCounter: a placeholder cannot become a window, + // but its absolute counter is real consumption and must survive. + let usage = parse_snapshot( + r#"{ + "copilot_plan": "business", + "token_based_billing": true, + "quota_snapshots": { + "premium_interactions": { + "entitlement": 0, + "remaining": 0, + "percent_remaining": 0, + "quota_id": "", + "placeholder": true, + "credits_used": 42.5 + } + } + }"#, + ); + assert!(usage.primary.is_informational); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("42.50 AI credits used") + ); +} + +#[test] +fn premium_credits_counter_wins_over_chat() { + let usage = parse_snapshot( + r#"{ + "copilot_plan": "pro", + "quota_snapshots": { + "chat": { + "entitlement": 100, + "remaining": 75, + "percent_remaining": 75, + "quota_id": "chat", + "credits_used": 1 + }, + "premium_interactions": { + "entitlement": 300, + "remaining": 240, + "percent_remaining": 80, + "quota_id": "premium_interactions", + "credits_used": 7 + } + } + }"#, + ); + let credits_row = usage + .extra_rate_windows + .iter() + .find(|w| w.id == "ai-credits") + .expect("ai-credits row"); + assert_eq!( + credits_row.window.reset_description.as_deref(), + Some("7 AI credits used") + ); + // Windows still render normally next to the counter. + assert!(credits_row.window.is_informational); + assert_used(&usage.primary, 20.0); +} + +#[test] +fn business_seat_without_credits_keeps_existing_error() { + let err = parse_snapshot_result( + r#"{ + "copilot_plan": "business", + "token_based_billing": true + }"#, + ); + assert!(err.is_err()); +}