diff --git a/rust/src/agent_sessions/parsers.rs b/rust/src/agent_sessions/parsers.rs index b29b4ea598..3dcec92008 100644 --- a/rust/src/agent_sessions/parsers.rs +++ b/rust/src/agent_sessions/parsers.rs @@ -52,13 +52,6 @@ impl AgentPSOutputParser { record.source } - pub fn has_codex_app_server(records: &[AgentProcessRecord]) -> bool { - records.iter().any(|record| { - record.kind == AgentProcessKind::AppServer - && record.provider == Some(AgentSessionProvider::Codex) - }) - } - fn parse_line(line: &str, seen_pids: &mut HashSet) -> Option { let mut fields = line.split_whitespace(); let pid = fields.next()?.parse::().ok()?; @@ -447,10 +440,6 @@ impl ClaudeSessionProjectMapper { }); transcripts } - - pub fn newest_transcript(cwd: &str, home_directory: &Path) -> Option { - Self::transcripts(cwd, home_directory).into_iter().next() - } } impl ClaudeTranscriptMetadataParser { diff --git a/rust/src/agent_sessions/pi_family/roots.rs b/rust/src/agent_sessions/pi_family/roots.rs index 017ba8047a..7ec0d10ce4 100644 --- a/rust/src/agent_sessions/pi_family/roots.rs +++ b/rust/src/agent_sessions/pi_family/roots.rs @@ -161,14 +161,6 @@ impl DirectoryScanBudget { retained } - pub fn visit_entry(&mut self) -> bool { - if !self.has_time_remaining() || self.entries_seen >= self.max_entry_count { - return false; - } - self.entries_seen += 1; - true - } - /// Resolve a path only while the shared scan deadline is live. /// /// The operation is gated before it starts. If it finishes after expiry, diff --git a/rust/src/agent_sessions/remote.rs b/rust/src/agent_sessions/remote.rs index e6ff1aca54..892532d5b7 100644 --- a/rust/src/agent_sessions/remote.rs +++ b/rust/src/agent_sessions/remote.rs @@ -338,10 +338,6 @@ impl RemoteSessionFetcher { })?; Ok(result) } - - pub fn failed_result(host: &str, err: impl std::fmt::Display) -> AgentSessionHostResult { - AgentSessionHostResult::failed(host.to_string(), err) - } } impl Default for RemoteSessionFetcher { diff --git a/rust/src/browser/cookie_cache.rs b/rust/src/browser/cookie_cache.rs index d943f2acbb..e47b7f5eae 100755 --- a/rust/src/browser/cookie_cache.rs +++ b/rust/src/browser/cookie_cache.rs @@ -5,13 +5,6 @@ //! are session secrets, so they are written through `secure_file` (DPAPI on //! Windows, staged and published atomically) rather than as plaintext. -// The cache API is currently unused by the crate (reserved for provider -// cookie reuse), so every public item would trip dead_code. -#![allow( - dead_code, - reason = "module reserved: CookieHeaderCache has no callers yet" -)] - use crate::core::ProviderId; use crate::secure_file; use chrono::{DateTime, Utc}; @@ -154,14 +147,6 @@ impl CookieHeaderCache { .collect::>() .join("; ") } - - /// Check if a cached entry exists and is fresh - pub fn has_fresh_cache(provider: ProviderId, max_age_secs: i64) -> bool { - match Self::load(provider) { - Some(entry) => !entry.is_stale(max_age_secs), - None => false, - } - } } /// Cookie header cache errors diff --git a/rust/src/browser/cookies.rs b/rust/src/browser/cookies.rs index 3233d702a5..13ccdd7b1b 100755 --- a/rust/src/browser/cookies.rs +++ b/rust/src/browser/cookies.rs @@ -3,11 +3,6 @@ //! Chromium browsers store cookies in an SQLite database encrypted with DPAPI. //! Firefox stores cookies in an unencrypted SQLite database. -#![allow( - dead_code, - reason = "the cookie API surface is shared across providers and cfg-split Windows/WSLS code paths, so individual helpers are reached only from their own browser/provider callers" -)] - use std::path::Path; use aes_gcm::{ @@ -813,18 +808,6 @@ pub fn get_cookie_header_for_domains(domains: &[&str]) -> Result Result { - let cookies = CookieExtractor::extract_for_domain(browser, domain)?; - if cookies.is_empty() { - return Err(CookieError::NotFound(domain.to_string())); - } - Ok(CookieExtractor::build_cookie_header(&cookies)) -} - #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/browser/detection.rs b/rust/src/browser/detection.rs index bc93219621..53cc9bd9d8 100755 --- a/rust/src/browser/detection.rs +++ b/rust/src/browser/detection.rs @@ -3,11 +3,6 @@ //! On native Windows, uses standard AppData paths. //! On WSL, resolves browser paths via /mnt/c/ to access Windows browser data. -#![allow( - dead_code, - reason = "browser detection types reserved for future cookie-based session management" -)] - use std::path::{Path, PathBuf}; use crate::wsl; diff --git a/rust/src/browser/watchdog.rs b/rust/src/browser/watchdog.rs index b6c5e8210c..cd7a905b96 100755 --- a/rust/src/browser/watchdog.rs +++ b/rust/src/browser/watchdog.rs @@ -3,11 +3,6 @@ //! Monitors and manages child web probe processes to prevent orphaned processes. //! Ensures that browser automation or scraping processes are properly cleaned up. -#![allow( - dead_code, - reason = "watchdog types reserved for future browser-session monitoring integration" -)] - use std::collections::HashMap; use std::process::Child; use std::sync::Arc; diff --git a/rust/src/cli/account.rs b/rust/src/cli/account.rs index efa3afb5f2..7b81127679 100755 --- a/rust/src/cli/account.rs +++ b/rust/src/cli/account.rs @@ -111,17 +111,7 @@ async fn list_accounts(provider_name: &str) -> anyhow::Result<()> { /// Add a new account async fn add_account(provider_name: &str, label: &str, token: &str) -> anyhow::Result<()> { - let provider = parse_provider(provider_name)?; - - if !TokenAccountSupport::is_supported(provider) { - anyhow::bail!( - "{} does not support token accounts.", - provider.display_name() - ); - } - - let store = TokenAccountStore::new(); - let mut data = store.load_provider(provider)?; + let (provider, store, mut data) = load_supported_accounts(provider_name)?; // Check for duplicate label if data @@ -142,17 +132,7 @@ async fn add_account(provider_name: &str, label: &str, token: &str) -> anyhow::R /// Remove an account async fn remove_account(provider_name: &str, account_ref: &str) -> anyhow::Result<()> { - let provider = parse_provider(provider_name)?; - - if !TokenAccountSupport::is_supported(provider) { - anyhow::bail!( - "{} does not support token accounts.", - provider.display_name() - ); - } - - let store = TokenAccountStore::new(); - let mut data = store.load_provider(provider)?; + let (provider, store, mut data) = load_supported_accounts(provider_name)?; let account = find_account(&data, account_ref)?; let label = account.label.clone(); @@ -171,17 +151,7 @@ async fn remove_account(provider_name: &str, account_ref: &str) -> anyhow::Resul /// Switch active account async fn switch_account(provider_name: &str, account_ref: &str) -> anyhow::Result<()> { - let provider = parse_provider(provider_name)?; - - if !TokenAccountSupport::is_supported(provider) { - anyhow::bail!( - "{} does not support token accounts.", - provider.display_name() - ); - } - - let store = TokenAccountStore::new(); - let mut data = store.load_provider(provider)?; + let (provider, store, mut data) = load_supported_accounts(provider_name)?; let account = find_account(&data, account_ref)?; let label = account.label.clone(); @@ -198,6 +168,22 @@ async fn switch_account(provider_name: &str, account_ref: &str) -> anyhow::Resul Ok(()) } +/// Resolve a provider that supports token accounts and load its stored accounts. +fn load_supported_accounts( + provider_name: &str, +) -> anyhow::Result<(ProviderId, TokenAccountStore, ProviderAccountData)> { + let provider = parse_provider(provider_name)?; + if !TokenAccountSupport::is_supported(provider) { + anyhow::bail!( + "{} does not support token accounts.", + provider.display_name() + ); + } + let store = TokenAccountStore::new(); + let data = store.load_provider(provider)?; + Ok((provider, store, data)) +} + /// Parse provider name to ProviderId fn parse_provider(name: &str) -> anyhow::Result { ProviderId::from_cli_name(name) diff --git a/rust/src/cli/config.rs b/rust/src/cli/config.rs index f9384d2296..85856f3013 100755 --- a/rust/src/cli/config.rs +++ b/rust/src/cli/config.rs @@ -134,13 +134,7 @@ impl ConfigOutputArgs { } fn print_json(self, value: &T) -> anyhow::Result<()> { - let json = if self.pretty { - serde_json::to_string_pretty(value)? - } else { - serde_json::to_string(value)? - }; - println!("{json}"); - Ok(()) + super::print_json(value, self.pretty) } } @@ -729,369 +723,4 @@ async fn show_paths() -> anyhow::Result<()> { } #[cfg(test)] -mod tests { - use super::{ - ClaudeCodeCredentialsConsent, ConfigCommand, ConfigFileError, ConfigOutputArgs, - ConsentAction, apply_consent_action, consent_text, provider_statuses, read_json_config, - sanitize_settings_for_dump, - }; - use crate::cli::{Cli, Commands}; - use crate::core::ProviderId; - #[cfg(windows)] - use crate::secure_file; - use crate::settings::{ManualCookies, Settings}; - use clap::Parser; - use serde_json::json; - - #[test] - fn preferences_subcommands_parse_a_path() { - use super::{ConfigArgs, ConfigCommand, PreferencesAction}; - use clap::Parser; - - let export = - ConfigArgs::try_parse_from(["config", "preferences", "export", "--file", "p.json"]) - .expect("export parses"); - assert!(matches!( - export.command, - ConfigCommand::Preferences { - action: PreferencesAction::Export { file: Some(_) } - } - )); - let stdout = ConfigArgs::try_parse_from(["config", "preferences", "export"]) - .expect("export without --file parses"); - assert!(matches!( - stdout.command, - ConfigCommand::Preferences { - action: PreferencesAction::Export { file: None } - } - )); - let import = - ConfigArgs::try_parse_from(["config", "preferences", "import", "--file", "p.json"]) - .expect("import parses"); - assert!(matches!( - import.command, - ConfigCommand::Preferences { - action: PreferencesAction::Import { .. } - } - )); - assert!(ConfigArgs::try_parse_from(["config", "preferences", "import"]).is_err()); - } - - fn parse_config(argv: &[&str]) -> Result { - let cli = Cli::try_parse_from(["codexbar", "config"].iter().chain(argv.iter()))?; - match cli.command { - Some(Commands::Config(args)) => Ok(args.command), - other => panic!("expected the config command, got {other:?}"), - } - } - - #[test] - fn config_providers_accepts_json_output_flags() { - let ConfigCommand::Providers { output } = parse_config(&["providers"]).unwrap() else { - panic!("expected providers"); - }; - assert!(!output.is_json()); - for argv in [ - &["providers", "--json"][..], - &["providers", "--format", "json"][..], - &["providers", "-f", "json", "--pretty"][..], - ] { - let ConfigCommand::Providers { output } = parse_config(argv).unwrap() else { - panic!("expected providers for {argv:?}"); - }; - assert!(output.is_json(), "{argv:?}"); - } - assert!(parse_config(&["providers", "--format", "toml"]).is_err()); - } - - #[test] - fn provider_status_json_matches_the_upstream_shape() { - let mut settings = Settings::default(); - settings.enabled_providers.clear(); - settings.enable_provider(ProviderId::Cursor); - - let statuses = provider_statuses(&settings); - - assert_eq!( - statuses.len(), - ProviderId::all() - .iter() - .filter(|p| !p.is_deprecated()) - .count() - ); - let cursor = statuses - .iter() - .find(|status| status.provider == "cursor") - .expect("cursor row"); - assert!(cursor.enabled); - let codex = statuses - .iter() - .find(|status| status.provider == "codex") - .expect("codex row"); - assert!(!codex.enabled); - - let value = serde_json::to_value(codex).unwrap(); - let mut keys: Vec<&str> = value - .as_object() - .unwrap() - .keys() - .map(String::as_str) - .collect(); - keys.sort_unstable(); - assert_eq!( - keys, - ["defaultEnabled", "displayName", "enabled", "provider"] - ); - assert_eq!(value["displayName"], ProviderId::Codex.display_name()); - assert_eq!(value["enabled"], false); - assert_eq!(value["defaultEnabled"], codex.default_enabled); - } - - #[test] - fn deprecated_providers_are_hidden_until_enabled_then_listed() { - let mut settings = Settings::default(); - settings.enabled_providers.clear(); - let statuses = provider_statuses(&settings); - for retired in [ProviderId::KimiK2, ProviderId::CrossModel] { - assert!(retired.is_deprecated()); - assert!(statuses.iter().all(|s| s.provider != retired.cli_name())); - } - - settings.enable_provider(ProviderId::KimiK2); - let statuses = provider_statuses(&settings); - let kimi = statuses - .iter() - .find(|s| s.provider == "kimik2") - .expect("enabled deprecated provider is listed"); - assert!(kimi.enabled); - assert!(statuses.iter().all(|s| s.provider != "crossmodel")); - } - - #[test] - fn provider_status_text_lines_keep_the_existing_format() { - let mut settings = Settings::default(); - settings.enabled_providers.clear(); - settings.enable_provider(ProviderId::Codex); - let statuses = provider_statuses(&settings); - let codex = statuses.iter().find(|s| s.provider == "codex").unwrap(); - let claude = statuses.iter().find(|s| s.provider == "claude").unwrap(); - - let marker = |default: bool| if default { " default" } else { "" }; - assert_eq!( - codex.text_line(), - format!("codex: enabled{} (Codex)", marker(codex.default_enabled)) - ); - assert_eq!( - claude.text_line(), - format!( - "claude: disabled{} (Claude)", - marker(claude.default_enabled) - ) - ); - } - - #[test] - fn provider_subcommands_accept_the_positional_or_provider_option() { - for argv in [ - &["enable", "cursor"][..], - &["enable", "-p", "cursor"][..], - &["enable", "--provider", "cursor", "--json"][..], - ] { - let ConfigCommand::Enable { target, .. } = parse_config(argv).unwrap() else { - panic!("expected enable for {argv:?}"); - }; - assert_eq!(target.name(), "cursor", "{argv:?}"); - } - let ConfigCommand::Disable { target, output } = - parse_config(&["disable", "--provider", "cursor", "--format", "json"]).unwrap() - else { - panic!("expected disable"); - }; - assert_eq!(target.name(), "cursor"); - assert!(output.is_json()); - let ConfigCommand::SetApiKey { target, stdin, .. } = - parse_config(&["set-api-key", "-p", "openrouter", "--stdin"]).unwrap() - else { - panic!("expected set-api-key"); - }; - assert_eq!(target.name(), "openrouter"); - assert!(stdin); - - assert!(parse_config(&["enable"]).is_err(), "provider is required"); - assert!( - parse_config(&["enable", "cursor", "--provider", "codex"]).is_err(), - "positional and --provider conflict" - ); - } - - #[test] - fn default_output_args_are_text() { - assert!(!ConfigOutputArgs::default().is_json()); - } - - #[test] - fn claude_code_credentials_parses_each_action() { - for (word, expected) in [ - ("allow", ConsentAction::Allow), - ("deny", ConsentAction::Deny), - ("status", ConsentAction::Status), - ] { - let ConfigCommand::ClaudeCodeCredentials { action, output } = - parse_config(&["claude-code-credentials", word, "--json"]).unwrap() - else { - panic!("expected claude-code-credentials for {word}"); - }; - assert_eq!(action, expected); - assert!(output.is_json()); - } - assert!(parse_config(&["claude-code-credentials"]).is_err()); - assert!(parse_config(&["claude-code-credentials", "maybe"]).is_err()); - } - - #[test] - fn consent_actions_toggle_only_the_claude_code_flag() { - let mut settings = Settings::default(); - assert!(!settings.claude_allow_reading_claude_code_credentials); - let before = serde_json::to_value(&settings).unwrap(); - - assert!(!apply_consent_action(&mut settings, ConsentAction::Status)); - assert!(!apply_consent_action(&mut settings, ConsentAction::Deny)); - assert!(apply_consent_action(&mut settings, ConsentAction::Allow)); - assert!(settings.claude_allow_reading_claude_code_credentials); - assert!(!apply_consent_action(&mut settings, ConsentAction::Allow)); - assert!(!apply_consent_action(&mut settings, ConsentAction::Status)); - assert!(settings.claude_allow_reading_claude_code_credentials); - - let mut after = serde_json::to_value(&settings).unwrap(); - after["claude_allow_reading_claude_code_credentials"] = json!(false); - assert_eq!(after, before, "no other setting changes"); - - assert!(apply_consent_action(&mut settings, ConsentAction::Deny)); - assert!(!settings.claude_allow_reading_claude_code_credentials); - } - - #[test] - fn consent_output_reports_the_choice() { - let value = serde_json::to_value(ClaudeCodeCredentialsConsent { - allowed: true, - config_path: Some("settings.json".to_string()), - }) - .unwrap(); - assert_eq!( - value, - json!({ "allowed": true, "configPath": "settings.json" }) - ); - assert!(consent_text(true).contains("allowed")); - assert!(consent_text(false).contains("not allowed")); - } - - #[cfg(windows)] - #[test] - fn validates_dpapi_protected_manual_cookies_without_plaintext() { - let dir = tempfile::tempdir().expect("create temporary directory"); - let path = dir.path().join("manual_cookies.json"); - let expected = r#"{"cookies":{}}"#; - - secure_file::write_string(&path, expected).expect("write protected cookies"); - - let raw = std::fs::read_to_string(&path).expect("read protected wrapper"); - assert!(raw.contains("\"format\": \"codexbar.secure-file\"")); - assert!(!raw.contains(expected)); - assert!( - read_json_config::(&path).is_ok(), - "valid protected cookies should validate" - ); - } - - #[test] - fn rejects_malformed_secure_wrapper_without_echoing_payload() { - let dir = tempfile::tempdir().expect("create temporary directory"); - let path = dir.path().join("manual_cookies.json"); - let payload = "AA=="; - let malformed = format!( - r#"{{"format":"codexbar.secure-file","version":1,"protection":"unsupported","payload":"{payload}"}}"# - ); - std::fs::write(&path, malformed).expect("write malformed wrapper"); - - let error = read_json_config::(&path) - .expect_err("reject malformed protected wrapper"); - let ConfigFileError::Read(error) = error else { - panic!("malformed protected wrapper must be a read error"); - }; - let message = error.to_string(); - assert!(message.contains("unsupported secure file protection")); - assert!(!message.contains(payload)); - } - - #[test] - fn sanitize_settings_for_dump_redacts_secret_fields() { - let raw = json!({ - "http_proxy_password": "proxy-secret", - "provider_configs": { - "claude": { - "manual_cookie_header": "sessionKey=abc", - "api_token": "tok-123", - "management_api_token": "management-secret" - } - }, - "api_keys": { - "keys": { - "zai": { - "api_key": "cb_test_api_key_456", - "label": "Work", - "saved_at": "2026-01-01" - } - } - }, - "manual_cookies": { - "cookies": { - "claude": { - "cookie_header": "sessionKey=secret-cookie", - "saved_at": "2026-01-01" - } - } - }, - "token_accounts": { - "factory": { - "accounts": [{ - "id": "11111111-1111-1111-1111-111111111111", - "label": "Team", - "token": "raw-token-value", - "added_at": 1 - }] - } - }, - "secret_key": "top-secret" - }); - - let redacted = sanitize_settings_for_dump(raw, false); - let text = serde_json::to_string(&redacted).expect("serialize"); - - assert!(text.contains("[REDACTED]")); - assert!(!text.contains("proxy-secret")); - assert!(!text.contains("sessionKey=abc")); - assert!(!text.contains("tok-123")); - assert!(!text.contains("management-secret")); - assert!(!text.contains("cb_test_api_key_456")); - assert!(!text.contains("secret-cookie")); - assert!(!text.contains("raw-token-value")); - assert!(!text.contains("top-secret")); - // Non-secret identity fields stay. - assert!(text.contains("Work")); - assert!(text.contains("Team")); - assert!(text.contains("11111111-1111-1111-1111-111111111111")); - } - - #[test] - fn sanitize_settings_for_dump_show_secrets_keeps_raw() { - let raw = json!({ - "api_key": "keep-me", - "nested": { "token": "also-keep", "label": "Team" } - }); - - let out = sanitize_settings_for_dump(raw.clone(), true); - assert_eq!(out, raw); - assert_eq!(out["api_key"], "keep-me"); - assert_eq!(out["nested"]["token"], "also-keep"); - } -} +mod tests; diff --git a/rust/src/cli/config/tests.rs b/rust/src/cli/config/tests.rs new file mode 100644 index 0000000000..8c9114f9b4 --- /dev/null +++ b/rust/src/cli/config/tests.rs @@ -0,0 +1,364 @@ +use super::{ + ClaudeCodeCredentialsConsent, ConfigCommand, ConfigFileError, ConfigOutputArgs, ConsentAction, + apply_consent_action, consent_text, provider_statuses, read_json_config, + sanitize_settings_for_dump, +}; +use crate::cli::{Cli, Commands}; +use crate::core::ProviderId; +#[cfg(windows)] +use crate::secure_file; +use crate::settings::{ManualCookies, Settings}; +use clap::Parser; +use serde_json::json; + +#[test] +fn preferences_subcommands_parse_a_path() { + use super::{ConfigArgs, ConfigCommand, PreferencesAction}; + use clap::Parser; + + let export = + ConfigArgs::try_parse_from(["config", "preferences", "export", "--file", "p.json"]) + .expect("export parses"); + assert!(matches!( + export.command, + ConfigCommand::Preferences { + action: PreferencesAction::Export { file: Some(_) } + } + )); + let stdout = ConfigArgs::try_parse_from(["config", "preferences", "export"]) + .expect("export without --file parses"); + assert!(matches!( + stdout.command, + ConfigCommand::Preferences { + action: PreferencesAction::Export { file: None } + } + )); + let import = + ConfigArgs::try_parse_from(["config", "preferences", "import", "--file", "p.json"]) + .expect("import parses"); + assert!(matches!( + import.command, + ConfigCommand::Preferences { + action: PreferencesAction::Import { .. } + } + )); + assert!(ConfigArgs::try_parse_from(["config", "preferences", "import"]).is_err()); +} + +fn parse_config(argv: &[&str]) -> Result { + let cli = Cli::try_parse_from(["codexbar", "config"].iter().chain(argv.iter()))?; + match cli.command { + Some(Commands::Config(args)) => Ok(args.command), + other => panic!("expected the config command, got {other:?}"), + } +} + +#[test] +fn config_providers_accepts_json_output_flags() { + let ConfigCommand::Providers { output } = parse_config(&["providers"]).unwrap() else { + panic!("expected providers"); + }; + assert!(!output.is_json()); + for argv in [ + &["providers", "--json"][..], + &["providers", "--format", "json"][..], + &["providers", "-f", "json", "--pretty"][..], + ] { + let ConfigCommand::Providers { output } = parse_config(argv).unwrap() else { + panic!("expected providers for {argv:?}"); + }; + assert!(output.is_json(), "{argv:?}"); + } + assert!(parse_config(&["providers", "--format", "toml"]).is_err()); +} + +#[test] +fn provider_status_json_matches_the_upstream_shape() { + let mut settings = Settings::default(); + settings.enabled_providers.clear(); + settings.enable_provider(ProviderId::Cursor); + + let statuses = provider_statuses(&settings); + + assert_eq!( + statuses.len(), + ProviderId::all() + .iter() + .filter(|p| !p.is_deprecated()) + .count() + ); + let cursor = statuses + .iter() + .find(|status| status.provider == "cursor") + .expect("cursor row"); + assert!(cursor.enabled); + let codex = statuses + .iter() + .find(|status| status.provider == "codex") + .expect("codex row"); + assert!(!codex.enabled); + + let value = serde_json::to_value(codex).unwrap(); + let mut keys: Vec<&str> = value + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect(); + keys.sort_unstable(); + assert_eq!( + keys, + ["defaultEnabled", "displayName", "enabled", "provider"] + ); + assert_eq!(value["displayName"], ProviderId::Codex.display_name()); + assert_eq!(value["enabled"], false); + assert_eq!(value["defaultEnabled"], codex.default_enabled); +} + +#[test] +fn deprecated_providers_are_hidden_until_enabled_then_listed() { + let mut settings = Settings::default(); + settings.enabled_providers.clear(); + let statuses = provider_statuses(&settings); + for retired in [ProviderId::KimiK2, ProviderId::CrossModel] { + assert!(retired.is_deprecated()); + assert!(statuses.iter().all(|s| s.provider != retired.cli_name())); + } + + settings.enable_provider(ProviderId::KimiK2); + let statuses = provider_statuses(&settings); + let kimi = statuses + .iter() + .find(|s| s.provider == "kimik2") + .expect("enabled deprecated provider is listed"); + assert!(kimi.enabled); + assert!(statuses.iter().all(|s| s.provider != "crossmodel")); +} + +#[test] +fn provider_status_text_lines_keep_the_existing_format() { + let mut settings = Settings::default(); + settings.enabled_providers.clear(); + settings.enable_provider(ProviderId::Codex); + let statuses = provider_statuses(&settings); + let codex = statuses.iter().find(|s| s.provider == "codex").unwrap(); + let claude = statuses.iter().find(|s| s.provider == "claude").unwrap(); + + let marker = |default: bool| if default { " default" } else { "" }; + assert_eq!( + codex.text_line(), + format!("codex: enabled{} (Codex)", marker(codex.default_enabled)) + ); + assert_eq!( + claude.text_line(), + format!( + "claude: disabled{} (Claude)", + marker(claude.default_enabled) + ) + ); +} + +#[test] +fn provider_subcommands_accept_the_positional_or_provider_option() { + for argv in [ + &["enable", "cursor"][..], + &["enable", "-p", "cursor"][..], + &["enable", "--provider", "cursor", "--json"][..], + ] { + let ConfigCommand::Enable { target, .. } = parse_config(argv).unwrap() else { + panic!("expected enable for {argv:?}"); + }; + assert_eq!(target.name(), "cursor", "{argv:?}"); + } + let ConfigCommand::Disable { target, output } = + parse_config(&["disable", "--provider", "cursor", "--format", "json"]).unwrap() + else { + panic!("expected disable"); + }; + assert_eq!(target.name(), "cursor"); + assert!(output.is_json()); + let ConfigCommand::SetApiKey { target, stdin, .. } = + parse_config(&["set-api-key", "-p", "openrouter", "--stdin"]).unwrap() + else { + panic!("expected set-api-key"); + }; + assert_eq!(target.name(), "openrouter"); + assert!(stdin); + + assert!(parse_config(&["enable"]).is_err(), "provider is required"); + assert!( + parse_config(&["enable", "cursor", "--provider", "codex"]).is_err(), + "positional and --provider conflict" + ); +} + +#[test] +fn default_output_args_are_text() { + assert!(!ConfigOutputArgs::default().is_json()); +} + +#[test] +fn claude_code_credentials_parses_each_action() { + for (word, expected) in [ + ("allow", ConsentAction::Allow), + ("deny", ConsentAction::Deny), + ("status", ConsentAction::Status), + ] { + let ConfigCommand::ClaudeCodeCredentials { action, output } = + parse_config(&["claude-code-credentials", word, "--json"]).unwrap() + else { + panic!("expected claude-code-credentials for {word}"); + }; + assert_eq!(action, expected); + assert!(output.is_json()); + } + assert!(parse_config(&["claude-code-credentials"]).is_err()); + assert!(parse_config(&["claude-code-credentials", "maybe"]).is_err()); +} + +#[test] +fn consent_actions_toggle_only_the_claude_code_flag() { + let mut settings = Settings::default(); + assert!(!settings.claude_allow_reading_claude_code_credentials); + let before = serde_json::to_value(&settings).unwrap(); + + assert!(!apply_consent_action(&mut settings, ConsentAction::Status)); + assert!(!apply_consent_action(&mut settings, ConsentAction::Deny)); + assert!(apply_consent_action(&mut settings, ConsentAction::Allow)); + assert!(settings.claude_allow_reading_claude_code_credentials); + assert!(!apply_consent_action(&mut settings, ConsentAction::Allow)); + assert!(!apply_consent_action(&mut settings, ConsentAction::Status)); + assert!(settings.claude_allow_reading_claude_code_credentials); + + let mut after = serde_json::to_value(&settings).unwrap(); + after["claude_allow_reading_claude_code_credentials"] = json!(false); + assert_eq!(after, before, "no other setting changes"); + + assert!(apply_consent_action(&mut settings, ConsentAction::Deny)); + assert!(!settings.claude_allow_reading_claude_code_credentials); +} + +#[test] +fn consent_output_reports_the_choice() { + let value = serde_json::to_value(ClaudeCodeCredentialsConsent { + allowed: true, + config_path: Some("settings.json".to_string()), + }) + .unwrap(); + assert_eq!( + value, + json!({ "allowed": true, "configPath": "settings.json" }) + ); + assert!(consent_text(true).contains("allowed")); + assert!(consent_text(false).contains("not allowed")); +} + +#[cfg(windows)] +#[test] +fn validates_dpapi_protected_manual_cookies_without_plaintext() { + let dir = tempfile::tempdir().expect("create temporary directory"); + let path = dir.path().join("manual_cookies.json"); + let expected = r#"{"cookies":{}}"#; + + secure_file::write_string(&path, expected).expect("write protected cookies"); + + let raw = std::fs::read_to_string(&path).expect("read protected wrapper"); + assert!(raw.contains("\"format\": \"codexbar.secure-file\"")); + assert!(!raw.contains(expected)); + assert!( + read_json_config::(&path).is_ok(), + "valid protected cookies should validate" + ); +} + +#[test] +fn rejects_malformed_secure_wrapper_without_echoing_payload() { + let dir = tempfile::tempdir().expect("create temporary directory"); + let path = dir.path().join("manual_cookies.json"); + let payload = "AA=="; + let malformed = format!( + r#"{{"format":"codexbar.secure-file","version":1,"protection":"unsupported","payload":"{payload}"}}"# + ); + std::fs::write(&path, malformed).expect("write malformed wrapper"); + + let error = + read_json_config::(&path).expect_err("reject malformed protected wrapper"); + let ConfigFileError::Read(error) = error else { + panic!("malformed protected wrapper must be a read error"); + }; + let message = error.to_string(); + assert!(message.contains("unsupported secure file protection")); + assert!(!message.contains(payload)); +} + +#[test] +fn sanitize_settings_for_dump_redacts_secret_fields() { + let raw = json!({ + "http_proxy_password": "proxy-secret", + "provider_configs": { + "claude": { + "manual_cookie_header": "sessionKey=abc", + "api_token": "tok-123", + "management_api_token": "management-secret" + } + }, + "api_keys": { + "keys": { + "zai": { + "api_key": "cb_test_api_key_456", + "label": "Work", + "saved_at": "2026-01-01" + } + } + }, + "manual_cookies": { + "cookies": { + "claude": { + "cookie_header": "sessionKey=secret-cookie", + "saved_at": "2026-01-01" + } + } + }, + "token_accounts": { + "factory": { + "accounts": [{ + "id": "11111111-1111-1111-1111-111111111111", + "label": "Team", + "token": "raw-token-value", + "added_at": 1 + }] + } + }, + "secret_key": "top-secret" + }); + + let redacted = sanitize_settings_for_dump(raw, false); + let text = serde_json::to_string(&redacted).expect("serialize"); + + assert!(text.contains("[REDACTED]")); + assert!(!text.contains("proxy-secret")); + assert!(!text.contains("sessionKey=abc")); + assert!(!text.contains("tok-123")); + assert!(!text.contains("management-secret")); + assert!(!text.contains("cb_test_api_key_456")); + assert!(!text.contains("secret-cookie")); + assert!(!text.contains("raw-token-value")); + assert!(!text.contains("top-secret")); + // Non-secret identity fields stay. + assert!(text.contains("Work")); + assert!(text.contains("Team")); + assert!(text.contains("11111111-1111-1111-1111-111111111111")); +} + +#[test] +fn sanitize_settings_for_dump_show_secrets_keeps_raw() { + let raw = json!({ + "api_key": "keep-me", + "nested": { "token": "also-keep", "label": "Team" } + }); + + let out = sanitize_settings_for_dump(raw.clone(), true); + assert_eq!(out, raw); + assert_eq!(out["api_key"], "keep-me"); + assert_eq!(out["nested"]["token"], "also-keep"); +} diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index efaef2dbf3..8ef8e24562 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -8,9 +8,7 @@ use super::cost_period::{ cost_totals_json, resolve_period, rolling_window_days, stamp_period, window_days, }; use super::usage::{OutputFormat, ProviderSelection}; -use crate::codex_costs::{ - CodexHostCostReport, CodexHostCostsArgs, CodexHostOutcome, run_codex_host_costs, -}; +use crate::codex_costs::{CodexHostCostsArgs, run_codex_host_costs}; use crate::codex_workspaces::short_session_id; use crate::core::{CostScanOptions, ProviderId}; use crate::cost_reporting_period::CostReportingPeriod; @@ -155,77 +153,42 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { let mut results: Vec = Vec::new(); for provider in providers.as_list() { - match provider { - ProviderId::Codex => { - let summary = scanner.scan_codex(); - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary, - supported: true, - token_history: None, - }); - } + let (summary, supported, token_history) = match provider { + ProviderId::Codex => (scanner.scan_codex(), true, None), ProviderId::Claude => { let summary = if pi_selected || args.provider_native_only { scanner.scan_claude_with_cancel_and_pi_sessions(None, false) } else { scanner.scan_claude() }; - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary, - supported: true, - token_history: None, - }); - } - ProviderId::Pi => { - let summary = scanner.scan_pi(); - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary, - supported: true, - token_history: None, - }); - } - ProviderId::Antigravity => { - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary: CostSummary::default(), - supported: true, - token_history: Some( - crate::providers::antigravity::local_sessions::summarize_with_pricing_refresh( - days, - args.refresh, - ) - .await, - ), - }); + (summary, true, None) } + ProviderId::Pi => (scanner.scan_pi(), true, None), + ProviderId::Antigravity => ( + CostSummary::default(), + true, + Some( + crate::providers::antigravity::local_sessions::summarize_with_pricing_refresh( + days, + args.refresh, + ) + .await, + ), + ), ProviderId::Muse => { let report = crate::providers::muse::local_usage::scan(days, None); - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary: CostSummary::default(), - supported: true, - token_history: Some(report.into()), - }); - } - _ => { - // Other providers don't have local logs to scan - results.push(CostResult { - provider: provider.cli_name().to_string(), - display_name: provider.display_name().to_string(), - summary: CostSummary::default(), - supported: false, - token_history: None, - }); + (CostSummary::default(), true, Some(report.into())) } - } + // Other providers don't have local logs to scan + _ => (CostSummary::default(), false, None), + }; + results.push(CostResult { + provider: provider.cli_name().to_string(), + display_name: provider.display_name().to_string(), + summary, + supported, + token_history, + }); } match format { @@ -247,60 +210,6 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { Ok(()) } -fn render_codex_host_report(report: &CodexHostCostReport) -> String { - let title = if report.source == "local" { - "This machine".to_string() - } else { - report.host.clone() - }; - let summary = match &report.outcome { - CodexHostOutcome::Success(summary) => summary, - CodexHostOutcome::Failed(error) => { - return format!("{title}: {error}"); - } - }; - - let window_line = |label: &str, window: &crate::codex_costs::CodexHostCostWindow| { - let cost = window - .cost_usd - .map(|value| format!("${value:.2}")) - .unwrap_or_else(|| "—".to_string()); - let tokens = window - .total_tokens - .map(format_number) - .unwrap_or_else(|| "—".to_string()); - let mut line = format!("{label}: {cost} · {tokens} tokens"); - if window.coverage.unpriced > 0 || window.coverage.unmetered > 0 { - line.push_str(" (some usage has no known price)"); - } - line - }; - - let history = if summary.history_days == 1 { - String::new() - } else { - format!( - "\n{}", - window_line( - &format!("Last {} days", summary.history_days), - &summary.history - ) - ) - }; - let coverage = if summary.history_coverage_is_established { - String::new() - } else { - "\nPartial history; scan is incomplete.".to_string() - }; - format!( - "{title} — Codex API-equivalent estimate (not billed)\n{}{}\nDay boundaries: {}{}", - window_line("Today", &summary.today), - history, - summary.bucket_time_zone, - coverage - ) -} - /// Cost result for a provider struct CostResult { provider: String, @@ -604,14 +513,7 @@ fn print_json_output( ) -> anyhow::Result<()> { let payloads = build_json_payloads(results, period, days, settings); - let output = if pretty { - serde_json::to_string_pretty(&payloads)? - } else { - serde_json::to_string(&payloads)? - }; - println!("{}", output); - - Ok(()) + super::print_json(&payloads, pretty) } /// Format a number with commas @@ -635,407 +537,4 @@ fn is_terminal() -> bool { } #[cfg(test)] -mod tests { - use super::*; - use crate::codex_costs::CodexCostSummary; - - #[test] - fn json_output_emits_a16_and_f18_fields() { - let summary = CostSummary { - sessions_count: 1, - history_coverage_established: true, - model_pricing_completeness: crate::cost_scanner::ModelPricingCompleteness::Partial { - unpriced_models: vec!["codex-auto-review".to_string()], - }, - ..Default::default() - }; - - let result = CostResult { - provider: "codex".to_string(), - display_name: "Codex".to_string(), - summary, - supported: true, - token_history: None, - }; - - // Capture stdout - // Build the JSON payload directly to assert field presence. - let payload = serde_json::json!({ - "provider": "codex", - "supported": true, - "days_scanned": 7, - "cost": { "total_usd": 0.0, "currency": "USD" }, - "tokens": { "input": 0, "output": 0, "cached": 0 }, - "sessions_count": 1, - "historyCoverageIsEstablished": true, - "knownZero": false, - "modelPricingCompleteness": { - "partial": { "unpriced_models": ["codex-auto-review"] } - }, - "by_model": {}, - "by_speed": {}, - "by_speed_tokens": {}, - "period": { "start": null, "end": null } - }); - - let s = serde_json::to_string(&payload).unwrap(); - assert!( - s.contains("historyCoverageIsEstablished"), - "A16 field present" - ); - assert!(s.contains("modelPricingCompleteness"), "F18 field present"); - assert!(s.contains("codex-auto-review"), "unpriced model listed"); - assert!(s.contains("\"partial\""), "partial branch emitted"); - // Verify backward-compat: original fields still present - assert!(s.contains("\"total_usd\"")); - assert!(s.contains("\"sessions_count\"")); - // drop the unused result - let _ = result; - } - - #[test] - fn json_output_a16_null_for_non_codex() { - let summary = CostSummary::default(); - let result = CostResult { - provider: "claude".to_string(), - display_name: "Claude".to_string(), - summary, - supported: true, - token_history: None, - }; - - // For non-codex, historyCoverageIsEstablished should be null. - let payload = serde_json::json!({ - "provider": result.provider, - "historyCoverageIsEstablished": serde_json::Value::Null, - }); - let s = serde_json::to_string(&payload).unwrap(); - assert!(s.contains("null"), "non-codex A16 is null"); - } - - #[test] - fn json_output_emits_incomplete_request_count_only_when_positive() { - let make = |incomplete_request_count| CostResult { - provider: "claude".to_string(), - display_name: "Claude".to_string(), - summary: CostSummary { - incomplete_request_count, - ..CostSummary::default() - }, - supported: true, - token_history: None, - }; - let payloads = build_json_payloads( - &[make(3), make(0)], - CostReportingPeriod::Rolling(30), - 30, - &Settings::default(), - ); - assert_eq!(payloads[0]["incompleteRequestCount"], 3); - assert!(payloads[1].get("incompleteRequestCount").is_none()); - } - - #[test] - fn incomplete_suffix_marks_only_positive_counts() { - assert_eq!(incomplete_suffix(0), ""); - assert_eq!(incomplete_suffix(2), " · Incomplete"); - } - - #[test] - fn antigravity_json_keeps_unknown_cost_distinct_from_zero() { - use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; - let payload = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - total_tokens: 12_345, - session_count: 2, - coverage: LocalHistoryCoverage::Complete, - cost_estimate: Default::default(), - ..Default::default() - }, - 30, - ); - assert!(payload["cost"]["total_usd"].is_null()); - assert_eq!(payload["tokens"]["total"], 12_345); - assert_eq!(payload["historyCoverage"], "complete"); - assert_eq!(payload["knownZero"], false); - - let partial = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - total_tokens: 999, - session_count: 1, - coverage: LocalHistoryCoverage::Partial, - cost_estimate: Default::default(), - ..Default::default() - }, - 30, - ); - assert!(partial["cost"]["total_usd"].is_null()); - assert!(partial["tokens"]["total"].is_null()); - assert_eq!(partial["historyCoverage"], "partial"); - } - - #[test] - fn antigravity_json_labels_public_price_estimates() { - use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; - let payload = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - total_tokens: 1_000, - session_count: 1, - coverage: LocalHistoryCoverage::Complete, - cost_estimate: crate::spend_contract::LocalCostEstimate { - known_subtotal_usd: Some(0.0125), - coverage: crate::spend_contract::CostCoverageCounts { - estimated: 1, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }, - 30, - ); - assert_eq!(payload["cost"]["total_usd"], 0.0125); - assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); - assert_eq!(payload["cost"]["currency"], "USD"); - assert!(payload["note"].as_str().unwrap().contains("not billed")); - } - - #[test] - fn antigravity_json_keeps_partial_scan_cost_as_a_subtotal() { - use crate::spend_contract::{ - CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, - }; - let payload = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - total_tokens: 1_000, - session_count: 1, - coverage: LocalHistoryCoverage::Partial, - cost_estimate: LocalCostEstimate { - known_subtotal_usd: Some(0.0125), - coverage: CostCoverageCounts { - estimated: 1, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }, - 30, - ); - assert!(payload["cost"]["total_usd"].is_null()); - assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); - assert_eq!(payload["historyCoverage"], "partial"); - assert!( - payload["note"] - .as_str() - .unwrap() - .contains("history is incomplete") - ); - } - - #[test] - fn antigravity_json_emits_zero_for_complete_empty_history() { - use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; - let payload = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - coverage: LocalHistoryCoverage::Complete, - ..Default::default() - }, - 30, - ); - assert_eq!(payload["cost"]["total_usd"], 0.0); - assert!(payload["cost"]["known_subtotal_usd"].is_null()); - assert_eq!(payload["cost"]["currency"], "USD"); - assert_eq!(payload["knownZero"], true); - } - - #[test] - fn antigravity_json_keeps_mixed_pricing_as_a_known_subtotal() { - use crate::spend_contract::{ - CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, - }; - let payload = crate::spend_contract::local_token_history_json( - "antigravity", - &LocalTokenHistorySummary { - total_tokens: 1_500, - session_count: 2, - coverage: LocalHistoryCoverage::Complete, - cost_estimate: LocalCostEstimate { - known_subtotal_usd: Some(0.0125), - coverage: CostCoverageCounts { - estimated: 1, - unpriced: 1, - ..Default::default() - }, - ..Default::default() - }, - ..Default::default() - }, - 30, - ); - assert!(payload["cost"]["total_usd"].is_null()); - assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); - assert_eq!(payload["cost"]["pricingCoverage"]["estimated"], 1); - assert_eq!(payload["cost"]["pricingCoverage"]["unpriced"], 1); - assert!(payload["note"].as_str().unwrap().contains("subtotal")); - } - - #[derive(clap::Parser)] - struct Wrapper { - #[command(flatten)] - args: CostArgs, - } - - fn parse_args(argv: &[&str]) -> CostArgs { - use clap::Parser; - Wrapper::try_parse_from(std::iter::once("cost").chain(argv.iter().copied())) - .expect("cost args parse") - .args - } - - #[test] - fn days_and_period_are_detectable_when_absent() { - let args = parse_args(&[]); - assert_eq!(args.days, None); - assert_eq!(args.period, None); - let args = parse_args(&["--period", "month-to-date", "--days", "7"]); - assert_eq!(args.days, Some(7)); - assert_eq!(args.period.as_deref(), Some("month-to-date")); - } - - #[test] - fn json_payload_carries_period_identity_and_totals() { - let summary = CostSummary { - total_cost_usd: 2.5, - input_tokens: 10, - output_tokens: 5, - sessions_count: 1, - ..Default::default() - }; - let results = [CostResult { - provider: "claude".to_string(), - display_name: "Claude".to_string(), - summary, - supported: true, - token_history: None, - }]; - let settings = Settings::default(); - let payloads = - build_json_payloads(&results, CostReportingPeriod::MonthToDate, 12, &settings); - let payload = &payloads[0]; - assert_eq!(payload["reportingPeriod"], "month-to-date"); - assert_eq!(payload["historyLabel"], "Month to date"); - assert_eq!(payload["days_scanned"], 12); - assert_eq!(payload["totals"]["totalCost"], 2.5); - assert_eq!(payload["totals"]["totalTokens"], 15); - // Existing fields keep their meaning. - assert_eq!(payload["cost"]["total_usd"], 2.5); - assert_eq!(payload["tokens"]["input"], 10); - - let rolling = - build_json_payloads(&results, CostReportingPeriod::Rolling(30), 30, &settings); - assert_eq!(rolling[0]["reportingPeriod"], "rolling:30"); - assert_eq!(rolling[0]["historyLabel"], "Last 30 days"); - } - - #[test] - fn token_history_payload_is_stamped_with_the_period() { - use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; - let results = [CostResult { - provider: "antigravity".to_string(), - display_name: "Antigravity".to_string(), - summary: CostSummary::default(), - supported: true, - token_history: Some(LocalTokenHistorySummary { - total_tokens: 3, - session_count: 1, - coverage: LocalHistoryCoverage::Complete, - ..Default::default() - }), - }]; - let settings = Settings::default(); - let payloads = build_json_payloads( - &results, - CostReportingPeriod::AllAvailable, - 20_000, - &settings, - ); - assert_eq!(payloads[0]["reportingPeriod"], "all"); - assert_eq!(payloads[0]["historyLabel"], "All"); - assert_eq!(payloads[0]["days_scanned"], 20_000); - } - - #[test] - fn provider_native_only_flag_default_false() { - // Default CostArgs has provider_native_only = false (backward compat). - let args = CostArgs::default(); - assert!(!args.provider_native_only); - } - - #[test] - fn cost_output_format_rejects_toon() { - assert!("toon".parse::().is_err()); - } - - #[test] - fn group_by_defaults_none_and_accepts_session() { - assert_eq!(CostGroupBy::from_arg(None), CostGroupBy::None); - assert_eq!(CostGroupBy::from_arg(Some("session")), CostGroupBy::Session); - } - - #[test] - fn remote_failure_retains_local_report() { - let local_summary = CodexCostSummary::from_summaries_at( - &CostSummary { - history_coverage_established: true, - known_zero: true, - ..CostSummary::default() - }, - &CostSummary { - history_coverage_established: true, - known_zero: true, - ..CostSummary::default() - }, - 30, - chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap(), - "UTC", - ); - let reports = [ - CodexHostCostReport::success("local", "local", local_summary), - CodexHostCostReport::failure( - "build-host", - "ssh", - crate::codex_costs::REMOTE_CODEX_COST_UNAVAILABLE, - ), - ]; - - assert!(reports[0].summary().is_some()); - assert!(reports[1].summary().is_none()); - assert_eq!( - reports[1].outcome, - CodexHostOutcome::Failed(crate::codex_costs::REMOTE_CODEX_COST_UNAVAILABLE.to_string()) - ); - } - - #[test] - fn host_text_preserves_unknown_values_and_separate_boundaries() { - let partial = CodexCostSummary::from_summaries_at( - &CostSummary::default(), - &CostSummary::default(), - 30, - chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap(), - "UTC", - ); - let text = - render_codex_host_report(&CodexHostCostReport::success("local", "local", partial)); - assert!(text.contains("Today: — · — tokens")); - assert!(text.contains("Last 30 days: — · — tokens")); - assert!(text.contains("Partial history")); - assert!(text.contains("Day boundaries: UTC")); - } -} +mod tests; diff --git a/rust/src/cli/cost/tests.rs b/rust/src/cli/cost/tests.rs new file mode 100644 index 0000000000..790c9ba90b --- /dev/null +++ b/rust/src/cli/cost/tests.rs @@ -0,0 +1,383 @@ +use super::*; +use crate::codex_costs::{CodexCostSummary, CodexHostCostReport, CodexHostOutcome}; + +#[test] +fn json_output_emits_a16_and_f18_fields() { + let summary = CostSummary { + sessions_count: 1, + history_coverage_established: true, + model_pricing_completeness: crate::cost_scanner::ModelPricingCompleteness::Partial { + unpriced_models: vec!["codex-auto-review".to_string()], + }, + ..Default::default() + }; + + let result = CostResult { + provider: "codex".to_string(), + display_name: "Codex".to_string(), + summary, + supported: true, + token_history: None, + }; + + // Capture stdout + // Build the JSON payload directly to assert field presence. + let payload = serde_json::json!({ + "provider": "codex", + "supported": true, + "days_scanned": 7, + "cost": { "total_usd": 0.0, "currency": "USD" }, + "tokens": { "input": 0, "output": 0, "cached": 0 }, + "sessions_count": 1, + "historyCoverageIsEstablished": true, + "knownZero": false, + "modelPricingCompleteness": { + "partial": { "unpriced_models": ["codex-auto-review"] } + }, + "by_model": {}, + "by_speed": {}, + "by_speed_tokens": {}, + "period": { "start": null, "end": null } + }); + + let s = serde_json::to_string(&payload).unwrap(); + assert!( + s.contains("historyCoverageIsEstablished"), + "A16 field present" + ); + assert!(s.contains("modelPricingCompleteness"), "F18 field present"); + assert!(s.contains("codex-auto-review"), "unpriced model listed"); + assert!(s.contains("\"partial\""), "partial branch emitted"); + // Verify backward-compat: original fields still present + assert!(s.contains("\"total_usd\"")); + assert!(s.contains("\"sessions_count\"")); + // drop the unused result + let _ = result; +} + +#[test] +fn json_output_a16_null_for_non_codex() { + let summary = CostSummary::default(); + let result = CostResult { + provider: "claude".to_string(), + display_name: "Claude".to_string(), + summary, + supported: true, + token_history: None, + }; + + // For non-codex, historyCoverageIsEstablished should be null. + let payload = serde_json::json!({ + "provider": result.provider, + "historyCoverageIsEstablished": serde_json::Value::Null, + }); + let s = serde_json::to_string(&payload).unwrap(); + assert!(s.contains("null"), "non-codex A16 is null"); +} + +#[test] +fn json_output_emits_incomplete_request_count_only_when_positive() { + let make = |incomplete_request_count| CostResult { + provider: "claude".to_string(), + display_name: "Claude".to_string(), + summary: CostSummary { + incomplete_request_count, + ..CostSummary::default() + }, + supported: true, + token_history: None, + }; + let payloads = build_json_payloads( + &[make(3), make(0)], + CostReportingPeriod::Rolling(30), + 30, + &Settings::default(), + ); + assert_eq!(payloads[0]["incompleteRequestCount"], 3); + assert!(payloads[1].get("incompleteRequestCount").is_none()); +} + +#[test] +fn incomplete_suffix_marks_only_positive_counts() { + assert_eq!(incomplete_suffix(0), ""); + assert_eq!(incomplete_suffix(2), " · Incomplete"); +} + +#[test] +fn antigravity_json_keeps_unknown_cost_distinct_from_zero() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + total_tokens: 12_345, + session_count: 2, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: Default::default(), + ..Default::default() + }, + 30, + ); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["tokens"]["total"], 12_345); + assert_eq!(payload["historyCoverage"], "complete"); + assert_eq!(payload["knownZero"], false); + + let partial = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + total_tokens: 999, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: Default::default(), + ..Default::default() + }, + 30, + ); + assert!(partial["cost"]["total_usd"].is_null()); + assert!(partial["tokens"]["total"].is_null()); + assert_eq!(partial["historyCoverage"], "partial"); +} + +#[test] +fn antigravity_json_labels_public_price_estimates() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + total_tokens: 1_000, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: crate::spend_contract::LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: crate::spend_contract::CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() + }, + 30, + ); + assert_eq!(payload["cost"]["total_usd"], 0.0125); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); + assert_eq!(payload["cost"]["currency"], "USD"); + assert!(payload["note"].as_str().unwrap().contains("not billed")); +} + +#[test] +fn antigravity_json_keeps_partial_scan_cost_as_a_subtotal() { + use crate::spend_contract::{ + CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, + }; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + total_tokens: 1_000, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() + }, + 30, + ); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); + assert_eq!(payload["historyCoverage"], "partial"); + assert!( + payload["note"] + .as_str() + .unwrap() + .contains("history is incomplete") + ); +} + +#[test] +fn antigravity_json_emits_zero_for_complete_empty_history() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }, + 30, + ); + assert_eq!(payload["cost"]["total_usd"], 0.0); + assert!(payload["cost"]["known_subtotal_usd"].is_null()); + assert_eq!(payload["cost"]["currency"], "USD"); + assert_eq!(payload["knownZero"], true); +} + +#[test] +fn antigravity_json_keeps_mixed_pricing_as_a_known_subtotal() { + use crate::spend_contract::{ + CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, + }; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + &LocalTokenHistorySummary { + total_tokens: 1_500, + session_count: 2, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + ..Default::default() + }, + ..Default::default() + }, + 30, + ); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); + assert_eq!(payload["cost"]["pricingCoverage"]["estimated"], 1); + assert_eq!(payload["cost"]["pricingCoverage"]["unpriced"], 1); + assert!(payload["note"].as_str().unwrap().contains("subtotal")); +} + +#[derive(clap::Parser)] +struct Wrapper { + #[command(flatten)] + args: CostArgs, +} + +fn parse_args(argv: &[&str]) -> CostArgs { + use clap::Parser; + Wrapper::try_parse_from(std::iter::once("cost").chain(argv.iter().copied())) + .expect("cost args parse") + .args +} + +#[test] +fn days_and_period_are_detectable_when_absent() { + let args = parse_args(&[]); + assert_eq!(args.days, None); + assert_eq!(args.period, None); + let args = parse_args(&["--period", "month-to-date", "--days", "7"]); + assert_eq!(args.days, Some(7)); + assert_eq!(args.period.as_deref(), Some("month-to-date")); +} + +#[test] +fn json_payload_carries_period_identity_and_totals() { + let summary = CostSummary { + total_cost_usd: 2.5, + input_tokens: 10, + output_tokens: 5, + sessions_count: 1, + ..Default::default() + }; + let results = [CostResult { + provider: "claude".to_string(), + display_name: "Claude".to_string(), + summary, + supported: true, + token_history: None, + }]; + let settings = Settings::default(); + let payloads = build_json_payloads(&results, CostReportingPeriod::MonthToDate, 12, &settings); + let payload = &payloads[0]; + assert_eq!(payload["reportingPeriod"], "month-to-date"); + assert_eq!(payload["historyLabel"], "Month to date"); + assert_eq!(payload["days_scanned"], 12); + assert_eq!(payload["totals"]["totalCost"], 2.5); + assert_eq!(payload["totals"]["totalTokens"], 15); + // Existing fields keep their meaning. + assert_eq!(payload["cost"]["total_usd"], 2.5); + assert_eq!(payload["tokens"]["input"], 10); + + let rolling = build_json_payloads(&results, CostReportingPeriod::Rolling(30), 30, &settings); + assert_eq!(rolling[0]["reportingPeriod"], "rolling:30"); + assert_eq!(rolling[0]["historyLabel"], "Last 30 days"); +} + +#[test] +fn token_history_payload_is_stamped_with_the_period() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let results = [CostResult { + provider: "antigravity".to_string(), + display_name: "Antigravity".to_string(), + summary: CostSummary::default(), + supported: true, + token_history: Some(LocalTokenHistorySummary { + total_tokens: 3, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }), + }]; + let settings = Settings::default(); + let payloads = build_json_payloads( + &results, + CostReportingPeriod::AllAvailable, + 20_000, + &settings, + ); + assert_eq!(payloads[0]["reportingPeriod"], "all"); + assert_eq!(payloads[0]["historyLabel"], "All"); + assert_eq!(payloads[0]["days_scanned"], 20_000); +} + +#[test] +fn provider_native_only_flag_default_false() { + // Default CostArgs has provider_native_only = false (backward compat). + let args = CostArgs::default(); + assert!(!args.provider_native_only); +} + +#[test] +fn cost_output_format_rejects_toon() { + assert!("toon".parse::().is_err()); +} + +#[test] +fn group_by_defaults_none_and_accepts_session() { + assert_eq!(CostGroupBy::from_arg(None), CostGroupBy::None); + assert_eq!(CostGroupBy::from_arg(Some("session")), CostGroupBy::Session); +} + +#[test] +fn remote_failure_retains_local_report() { + let local_summary = CodexCostSummary::from_summaries_at( + &CostSummary { + history_coverage_established: true, + known_zero: true, + ..CostSummary::default() + }, + &CostSummary { + history_coverage_established: true, + known_zero: true, + ..CostSummary::default() + }, + 30, + chrono::DateTime::from_timestamp(1_700_000_000, 0).unwrap(), + "UTC", + ); + let reports = [ + CodexHostCostReport::success("local", "local", local_summary), + CodexHostCostReport::failure( + "build-host", + "ssh", + crate::codex_costs::REMOTE_CODEX_COST_UNAVAILABLE, + ), + ]; + + assert!(reports[0].summary().is_some()); + assert!(reports[1].summary().is_none()); + assert_eq!( + reports[1].outcome, + CodexHostOutcome::Failed(crate::codex_costs::REMOTE_CODEX_COST_UNAVAILABLE.to_string()) + ); +} diff --git a/rust/src/cli/dashboard.rs b/rust/src/cli/dashboard.rs index e27498c5c4..7fe4be62e8 100644 --- a/rust/src/cli/dashboard.rs +++ b/rust/src/cli/dashboard.rs @@ -43,11 +43,7 @@ pub async fn run(args: DashboardArgs) -> anyhow::Result<()> { let producer = SnapshotProducer::new(60, Some(identity)).with_fetch_timeout(fetch_timeout); let payload = producer.collect().await.map_err(anyhow::Error::msg)?; - let body = if args.pretty { - serde_json::to_string_pretty(&payload)? - } else { - serde_json::to_string(&payload)? - }; + let body = super::to_json(&payload, args.pretty)?; match &args.output { Some(path) => { diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 623389af5c..76e9681ea3 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -132,12 +132,7 @@ pub async fn run(args: DiagnoseArgs) -> anyhow::Result<()> { diagnostics, }; - let json = if args.pretty { - serde_json::to_string_pretty(&export) - } else { - serde_json::to_string(&export) - } - .context("failed to serialize diagnostics")?; + let json = super::to_json(&export, args.pretty).context("failed to serialize diagnostics")?; println!("{json}"); Ok(()) @@ -169,16 +164,11 @@ async fn collect_provider_diagnostic( let source_mode = configured_source_mode(provider_id, settings, source_mode); let ctx = FetchContext { source_mode, - include_credits: true, web_timeout, - verbose: false, manual_cookie_header: manual_cookies .get(provider_id.cli_name()) .map(ToOwned::to_owned), - manual_cookie_missing: false, api_key: api_keys.get(provider_id.cli_name()).map(ToOwned::to_owned), - token_account_kind: None, - token_account_isolated: false, workspace_id: settings .provider_config(provider_id) .and_then(|config| config.workspace_id.clone()), @@ -189,11 +179,10 @@ async fn collect_provider_diagnostic( gateway_url: settings .provider_config(provider_id) .and_then(|config| config.gateway_url.clone()), - auto_prefer_web: false, - browser_cookie_import: false, // Diagnostics keep the short optional-join grace (upstream #2583 gate). requires_optional_usage_completeness: false, optional_details_enabled: settings.optional_details_enabled(provider_id), + ..FetchContext::default() }; let fetch_result = provider.fetch_usage(&ctx).await; diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index da5e9f138b..235583b40c 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -224,34 +224,10 @@ struct GuardResultPayload { /// Run the guard command. Returns the process exit code. pub async fn run(args: GuardArgs) -> i32 { - let window = match GuardWindow::parse(&args.window) { - Some(w) => w, - None => { - eprintln!("Error: --window must be session|weekly."); - return exit_codes::USAGE_ERROR; - } - }; - - let minimum_remaining = match parse_min_remaining(args.min_remaining) { - Ok(v) => v, - Err(msg) => { - eprintln!("Error: {}", msg); - return exit_codes::USAGE_ERROR; - } - }; - - let timeout_secs = match parse_timeout_secs(args.timeout) { - Ok(v) => v, + let (window, minimum_remaining, timeout_secs, provider_id) = match parse_guard_args(&args) { + Ok(parsed) => parsed, Err(msg) => { - eprintln!("Error: {}", msg); - return exit_codes::USAGE_ERROR; - } - }; - - let provider_id = match resolve_guard_provider(&args.provider) { - Ok(id) => id, - Err(msg) => { - eprintln!("Error: {}", msg); + eprintln!("Error: {msg}"); return exit_codes::USAGE_ERROR; } }; @@ -286,6 +262,16 @@ pub async fn run(args: GuardArgs) -> i32 { evaluation.exit_code } +/// Validate the guard flags in order; the first failure is the usage error. +fn parse_guard_args(args: &GuardArgs) -> Result<(GuardWindow, f64, f64, ProviderId), String> { + let window = GuardWindow::parse(&args.window) + .ok_or_else(|| "--window must be session|weekly.".to_string())?; + let minimum_remaining = parse_min_remaining(args.min_remaining)?; + let timeout_secs = parse_timeout_secs(args.timeout)?; + let provider_id = resolve_guard_provider(&args.provider)?; + Ok((window, minimum_remaining, timeout_secs, provider_id)) +} + async fn run_guard_fetch(timeout_secs: f64, operation: F) -> GuardFetchOutcome where F: FnOnce() -> Fut, @@ -314,21 +300,9 @@ async fn fetch_guard_outcome( source_mode, include_credits: false, web_timeout, - verbose: false, - manual_cookie_header: None, - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, - workspace_id: None, - seat_credit_entitlement: None, - api_region: None, - gateway_url: None, - auto_prefer_web: false, - browser_cookie_import: false, // Guard checks keep the short optional-join grace. requires_optional_usage_completeness: false, - optional_details_enabled: false, + ..FetchContext::default() }; match provider.fetch_usage(&ctx).await { @@ -366,12 +340,7 @@ fn emit_guard_result( .unavailable_reason .map(|r| r.as_str().to_string()), }; - let rendered = if pretty { - serde_json::to_string_pretty(&payload) - } else { - serde_json::to_string(&payload) - }; - match rendered { + match super::to_json(&payload, pretty) { Ok(s) => println!("{}", s), Err(e) => eprintln!("Error: failed to encode JSON: {}", e), } diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index ea705e022a..387d1f973c 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -1,5 +1,6 @@ //! `codexbar hooks` — list / enable / disable / test / watch external hook rules. +use super::print_json; use clap::{Args, Subcommand}; use serde::Serialize; use std::sync::Arc; @@ -24,29 +25,20 @@ pub struct HooksArgs { #[derive(Subcommand, Debug, Clone)] pub enum HooksCommand { /// Print configured hook rules - List(HooksListArgs), + List(HooksOutputArgs), /// Enable hooks in hooks.json (master switch) - Enable(HooksToggleArgs), + Enable(HooksOutputArgs), /// Disable hooks in hooks.json (master switch) - Disable(HooksToggleArgs), + Disable(HooksOutputArgs), /// Run matching rules for a sample event Test(HooksTestArgs), /// Continuously poll providers and fire hooks on real transitions Watch(HooksWatchArgs), } +/// `--json` / `--pretty`, shared by `hooks list|enable|disable|test`. #[derive(Args, Debug, Clone)] -pub struct HooksListArgs { - /// Emit JSON - #[arg(long)] - pub json: bool, - /// Pretty-print JSON - #[arg(long)] - pub pretty: bool, -} - -#[derive(Args, Debug, Clone)] -pub struct HooksToggleArgs { +pub struct HooksOutputArgs { /// Emit JSON #[arg(long)] pub json: bool, @@ -62,12 +54,8 @@ pub struct HooksTestArgs { /// Provider CLI name #[arg(long)] pub provider: String, - /// Emit JSON - #[arg(long)] - pub json: bool, - /// Pretty-print JSON - #[arg(long)] - pub pretty: bool, + #[command(flatten)] + pub output: HooksOutputArgs, } /// Default poll period (seconds). Longer than serve cache TTL — watch originates @@ -288,33 +276,24 @@ async fn hooks_watch_observation( let region = settings.api_region(provider_id); let gateway = settings.gateway_url(provider_id); - let mut ctx = FetchContext { + let ctx = FetchContext { source_mode, include_credits: false, web_timeout, verbose, - manual_cookie_header: None, - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, + api_key: ApiKeys::load() + .get(provider_id.cli_name()) + .map(|s| s.to_string()), workspace_id: (!workspace.is_empty()).then(|| workspace.to_string()), seat_credit_entitlement: settings.seat_credit_entitlement(provider_id), api_region: (!region.is_empty()).then(|| region.to_string()), gateway_url: (!gateway.is_empty()).then(|| gateway.to_string()), - auto_prefer_web: false, - browser_cookie_import: false, // Hook watches keep the short optional-join grace. requires_optional_usage_completeness: false, optional_details_enabled: settings.optional_details_enabled(provider_id), + ..FetchContext::default() }; - if ctx.api_key.is_none() { - ctx.api_key = ApiKeys::load() - .get(provider_id.cli_name()) - .map(|s| s.to_string()); - } - let provider = instantiate_provider(provider_id); match provider.fetch_usage(&ctx).await { Ok(result) => { @@ -450,7 +429,7 @@ fn report_hook_event(event: &HookEvent, json: bool, pretty: bool) -> anyhow::Res Ok(()) } -fn run_list(args: HooksListArgs) -> anyhow::Result<()> { +fn run_list(args: HooksOutputArgs) -> anyhow::Result<()> { let config = HooksConfig::load(); let settings = Settings::load(); let path = HooksConfig::path().map(|p| p.display().to_string()); @@ -519,7 +498,7 @@ fn run_list(args: HooksListArgs) -> anyhow::Result<()> { Ok(()) } -fn run_set_enabled(enabled: bool, args: HooksToggleArgs) -> anyhow::Result<()> { +fn run_set_enabled(enabled: bool, args: HooksOutputArgs) -> anyhow::Result<()> { let mut config = HooksConfig::load(); config.enabled = enabled; let path = config.save().map_err(anyhow::Error::msg)?; @@ -574,26 +553,18 @@ fn run_test(args: HooksTestArgs) -> anyhow::Result<()> { let base_env = std::env::vars().collect(); let mut results = Vec::new(); for rule in rules { - match HookRunner::run(rule, &event, &base_env) { - Ok(()) => results.push(HookTestResult { - executable: rule.executable.display().to_string(), - event: event_type.as_str().into(), - provider: provider.cli_name().into(), - ok: true, - error: None, - }), - Err(err) => results.push(HookTestResult { - executable: rule.executable.display().to_string(), - event: event_type.as_str().into(), - provider: provider.cli_name().into(), - ok: false, - error: Some(err), - }), - } + let error = HookRunner::run(rule, &event, &base_env).err(); + results.push(HookTestResult { + executable: rule.executable.display().to_string(), + event: event_type.as_str().into(), + provider: provider.cli_name().into(), + ok: error.is_none(), + error, + }); } - if args.json { - print_json(&results, args.pretty)?; + if args.output.json { + print_json(&results, args.output.pretty)?; } else { for r in &results { if r.ok { @@ -650,19 +621,78 @@ fn parse_event(raw: &str) -> anyhow::Result { } } -fn print_json(value: &T, pretty: bool) -> anyhow::Result<()> { - if pretty { - println!("{}", serde_json::to_string_pretty(value)?); - } else { - println!("{}", serde_json::to_string(value)?); - } - Ok(()) -} - #[cfg(test)] mod tests { use super::*; + /// `hooks --help` exactly as printed, so arg refactors keep the CLI text. + const EXPECTED_HOOKS_HELP: &str = r#"@@list +Print configured hook rules + +Usage: codexbar hooks list [OPTIONS] + +Options: + --json Emit JSON + -v, --verbose Enable verbose logging + --json-output Emit machine-readable logs (JSON) to stderr + --pretty Pretty-print JSON + --no-color Disable ANSI colors in output + -h, --help Print help +@@enable +Enable hooks in hooks.json (master switch) + +Usage: codexbar hooks enable [OPTIONS] + +Options: + --json Emit JSON + -v, --verbose Enable verbose logging + --json-output Emit machine-readable logs (JSON) to stderr + --pretty Pretty-print JSON + --no-color Disable ANSI colors in output + -h, --help Print help +@@disable +Disable hooks in hooks.json (master switch) + +Usage: codexbar hooks disable [OPTIONS] + +Options: + --json Emit JSON + -v, --verbose Enable verbose logging + --json-output Emit machine-readable logs (JSON) to stderr + --pretty Pretty-print JSON + --no-color Disable ANSI colors in output + -h, --help Print help +@@test +Run matching rules for a sample event + +Usage: codexbar hooks test [OPTIONS] --provider + +Arguments: + Event name (quota_low, quota_reached, quota_reset, usage_updated, provider_unavailable, provider_recovered, refresh_failed) + +Options: + --provider Provider CLI name + -v, --verbose Enable verbose logging + --json Emit JSON + --json-output Emit machine-readable logs (JSON) to stderr + --no-color Disable ANSI colors in output + --pretty Pretty-print JSON + -h, --help Print help +"#; + + #[test] + fn hooks_subcommand_help_is_pinned() { + use clap::CommandFactory; + let mut out = String::new(); + for sub in ["list", "enable", "disable", "test"] { + let err = crate::cli::Cli::command() + .try_get_matches_from(["codexbar", "hooks", sub, "--help"]) + .unwrap_err(); + out.push_str(&format!("@@{sub}\n{}", err.render())); + } + assert_eq!(out, EXPECTED_HOOKS_HELP); + } + #[test] fn browser_sign_in_failures_report_auth_required() { let error = ProviderError::BrowserSignInRequired { diff --git a/rust/src/cli/mod.rs b/rust/src/cli/mod.rs index f5d96f0501..73ea8ca424 100755 --- a/rust/src/cli/mod.rs +++ b/rust/src/cli/mod.rs @@ -6,11 +6,6 @@ //! - `codexbar cost` - print local token cost usage //! - `codexbar autostart` - manage Windows auto-start -#![allow( - dead_code, - reason = "CLI subcommand types reserved for future command expansion" -)] - pub mod account; pub mod autostart; pub mod config; @@ -111,6 +106,27 @@ pub enum Commands { Workspaces(workspaces::WorkspacesArgs), } +/// Serialize `value` as pretty or compact JSON. +pub(crate) fn to_json( + value: &T, + pretty: bool, +) -> serde_json::Result { + if pretty { + serde_json::to_string_pretty(value) + } else { + serde_json::to_string(value) + } +} + +/// Print `value` as pretty or compact JSON on stdout. +pub(crate) fn print_json( + value: &T, + pretty: bool, +) -> anyhow::Result<()> { + println!("{}", to_json(value, pretty)?); + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/rust/src/cli/serve/dashboard/coordinator/mod.rs b/rust/src/cli/serve/dashboard/coordinator/mod.rs index ee20275ae3..ee115e7c88 100644 --- a/rust/src/cli/serve/dashboard/coordinator/mod.rs +++ b/rust/src/cli/serve/dashboard/coordinator/mod.rs @@ -89,6 +89,7 @@ impl SnapshotCoordinator { /// or missing cache is refreshed in the background. Concurrent callers /// share one refresh. When no Tokio runtime is active, this remains a pure /// cache lookup rather than claiming a flight that cannot be driven. + #[cfg(test)] fn latest_or_trigger_refresh(&self) -> Option> { self.latest_cached_or_trigger_refresh() .map(|cached| cached.payload) diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index cd95719095..88abc1ebe7 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -59,290 +59,115 @@ pub fn icon_url_map() -> String { out } +/// Embeds `icons/ProviderIcon-.svg` under `ProviderIcon-`. +macro_rules! icons { + ($($name:literal),* $(,)?) => { + &[$(( + concat!("ProviderIcon-", $name), + include_bytes!(concat!("icons/ProviderIcon-", $name, ".svg")), + )),*] + }; +} + /// Sorted static table (enables binary search). Keep in ASCII order. -static ICONS: &[(&str, &[u8])] = &[ - ( - "ProviderIcon-abacus", - include_bytes!("icons/ProviderIcon-abacus.svg"), - ), - ( - "ProviderIcon-aiand", - include_bytes!("icons/ProviderIcon-aiand.svg"), - ), - ( - "ProviderIcon-alibaba", - include_bytes!("icons/ProviderIcon-alibaba.svg"), - ), - ( - "ProviderIcon-amp", - include_bytes!("icons/ProviderIcon-amp.svg"), - ), - ( - "ProviderIcon-antigravity", - include_bytes!("icons/ProviderIcon-antigravity.svg"), - ), - ( - "ProviderIcon-augment", - include_bytes!("icons/ProviderIcon-augment.svg"), - ), - ( - "ProviderIcon-bedrock", - include_bytes!("icons/ProviderIcon-bedrock.svg"), - ), - ( - "ProviderIcon-chutes", - include_bytes!("icons/ProviderIcon-chutes.svg"), - ), - ( - "ProviderIcon-claude", - include_bytes!("icons/ProviderIcon-claude.svg"), - ), - ( - "ProviderIcon-clinepass", - include_bytes!("icons/ProviderIcon-clinepass.svg"), - ), - ( - "ProviderIcon-codebuff", - include_bytes!("icons/ProviderIcon-codebuff.svg"), - ), - ( - "ProviderIcon-coderabbit", - include_bytes!("icons/ProviderIcon-coderabbit.svg"), - ), - ( - "ProviderIcon-codex", - include_bytes!("icons/ProviderIcon-codex.svg"), - ), - ( - "ProviderIcon-commandcode", - include_bytes!("icons/ProviderIcon-commandcode.svg"), - ), - ( - "ProviderIcon-copilot", - include_bytes!("icons/ProviderIcon-copilot.svg"), - ), - ( - "ProviderIcon-cursor", - include_bytes!("icons/ProviderIcon-cursor.svg"), - ), - ( - "ProviderIcon-deepgram", - include_bytes!("icons/ProviderIcon-deepgram.svg"), - ), - ( - "ProviderIcon-deepinfra", - include_bytes!("icons/ProviderIcon-deepinfra.svg"), - ), - ( - "ProviderIcon-deepseek", - include_bytes!("icons/ProviderIcon-deepseek.svg"), - ), - ( - "ProviderIcon-devin", - include_bytes!("icons/ProviderIcon-devin.svg"), - ), - ( - "ProviderIcon-devpass", - include_bytes!("icons/ProviderIcon-devpass.svg"), - ), - ( - "ProviderIcon-doubao", - include_bytes!("icons/ProviderIcon-doubao.svg"), - ), - ( - "ProviderIcon-elevenlabs", - include_bytes!("icons/ProviderIcon-elevenlabs.svg"), - ), - ( - "ProviderIcon-factory", - include_bytes!("icons/ProviderIcon-factory.svg"), - ), - ( - "ProviderIcon-gemini", - include_bytes!("icons/ProviderIcon-gemini.svg"), - ), - ( - "ProviderIcon-grok", - include_bytes!("icons/ProviderIcon-grok.svg"), - ), - ( - "ProviderIcon-groq", - include_bytes!("icons/ProviderIcon-groq.svg"), - ), - ( - "ProviderIcon-huggingface", - include_bytes!("icons/ProviderIcon-huggingface.svg"), - ), - ( - "ProviderIcon-jetbrains", - include_bytes!("icons/ProviderIcon-jetbrains.svg"), - ), - ( - "ProviderIcon-kilo", - include_bytes!("icons/ProviderIcon-kilo.svg"), - ), - ( - "ProviderIcon-kimi", - include_bytes!("icons/ProviderIcon-kimi.svg"), - ), - ( - "ProviderIcon-kiro", - include_bytes!("icons/ProviderIcon-kiro.svg"), - ), - ( - "ProviderIcon-litellm", - include_bytes!("icons/ProviderIcon-litellm.svg"), - ), - ( - "ProviderIcon-llmman", - include_bytes!("icons/ProviderIcon-llmman.svg"), - ), - ( - "ProviderIcon-llmproxy", - include_bytes!("icons/ProviderIcon-llmproxy.svg"), - ), - ( - "ProviderIcon-longcat", - include_bytes!("icons/ProviderIcon-longcat.svg"), - ), - ( - "ProviderIcon-manus", - include_bytes!("icons/ProviderIcon-manus.svg"), - ), - ( - "ProviderIcon-mimo", - include_bytes!("icons/ProviderIcon-mimo.svg"), - ), - ( - "ProviderIcon-minimax", - include_bytes!("icons/ProviderIcon-minimax.svg"), - ), - ( - "ProviderIcon-mistral", - include_bytes!("icons/ProviderIcon-mistral.svg"), - ), - ( - "ProviderIcon-muse", - include_bytes!("icons/ProviderIcon-muse.svg"), - ), - ( - "ProviderIcon-neuralwatt", - include_bytes!("icons/ProviderIcon-neuralwatt.svg"), - ), - ( - "ProviderIcon-notion", - include_bytes!("icons/ProviderIcon-notion.svg"), - ), - ( - "ProviderIcon-nous", - include_bytes!("icons/ProviderIcon-nous.svg"), - ), - ( - "ProviderIcon-ollama", - include_bytes!("icons/ProviderIcon-ollama.svg"), - ), - ( - "ProviderIcon-opencode", - include_bytes!("icons/ProviderIcon-opencode.svg"), - ), - ( - "ProviderIcon-opencodego", - include_bytes!("icons/ProviderIcon-opencodego.svg"), - ), - ( - "ProviderIcon-openrouter", - include_bytes!("icons/ProviderIcon-openrouter.svg"), - ), - ( - "ProviderIcon-perplexity", - include_bytes!("icons/ProviderIcon-perplexity.svg"), - ), - ( - "ProviderIcon-poe", - include_bytes!("icons/ProviderIcon-poe.svg"), - ), - ( - "ProviderIcon-qoder", - include_bytes!("icons/ProviderIcon-qoder.svg"), - ), - ( - "ProviderIcon-qwencloud", - include_bytes!("icons/ProviderIcon-qwencloud.svg"), - ), - ( - "ProviderIcon-raycast", - include_bytes!("icons/ProviderIcon-raycast.svg"), - ), - ( - "ProviderIcon-replicate", - include_bytes!("icons/ProviderIcon-replicate.svg"), - ), - ( - "ProviderIcon-sakana", - include_bytes!("icons/ProviderIcon-sakana.svg"), - ), - ( - "ProviderIcon-stepfun", - include_bytes!("icons/ProviderIcon-stepfun.svg"), - ), - ( - "ProviderIcon-sub2api", - include_bytes!("icons/ProviderIcon-sub2api.svg"), - ), - ( - "ProviderIcon-t3chat", - include_bytes!("icons/ProviderIcon-t3chat.svg"), - ), - ( - "ProviderIcon-venice", - include_bytes!("icons/ProviderIcon-venice.svg"), - ), - ( - "ProviderIcon-vertexai", - include_bytes!("icons/ProviderIcon-vertexai.svg"), - ), - ( - "ProviderIcon-warp", - include_bytes!("icons/ProviderIcon-warp.svg"), - ), - ( - "ProviderIcon-wayfinder", - include_bytes!("icons/ProviderIcon-wayfinder.svg"), - ), - ( - "ProviderIcon-windsurf", - include_bytes!("icons/ProviderIcon-windsurf.svg"), - ), - ( - "ProviderIcon-xai", - include_bytes!("icons/ProviderIcon-xai.svg"), - ), - ( - "ProviderIcon-xkiro", - include_bytes!("icons/ProviderIcon-xkiro.svg"), - ), - ( - "ProviderIcon-zai", - include_bytes!("icons/ProviderIcon-zai.svg"), - ), - ( - "ProviderIcon-zed", - include_bytes!("icons/ProviderIcon-zed.svg"), - ), - ( - "ProviderIcon-zenmux", - include_bytes!("icons/ProviderIcon-zenmux.svg"), - ), - ( - "ProviderIcon-zoommate", - include_bytes!("icons/ProviderIcon-zoommate.svg"), - ), +static ICONS: &[(&str, &[u8])] = icons![ + "abacus", + "aiand", + "alibaba", + "amp", + "antigravity", + "augment", + "bedrock", + "chutes", + "claude", + "clinepass", + "codebuff", + "coderabbit", + "codex", + "commandcode", + "copilot", + "cursor", + "deepgram", + "deepinfra", + "deepseek", + "devin", + "devpass", + "doubao", + "elevenlabs", + "factory", + "gemini", + "grok", + "groq", + "huggingface", + "jetbrains", + "kilo", + "kimi", + "kiro", + "litellm", + "llmman", + "llmproxy", + "longcat", + "manus", + "mimo", + "minimax", + "mistral", + "muse", + "neuralwatt", + "notion", + "nous", + "ollama", + "opencode", + "opencodego", + "openrouter", + "perplexity", + "poe", + "qoder", + "qwencloud", + "raycast", + "replicate", + "sakana", + "stepfun", + "sub2api", + "t3chat", + "venice", + "vertexai", + "warp", + "wayfinder", + "windsurf", + "xai", + "xkiro", + "zai", + "zed", + "zenmux", + "zoommate", ]; #[cfg(test)] mod tests { use super::*; + #[test] + fn icon_table_embeds_every_svg_in_the_icons_dir() { + let dir = + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/cli/serve/dashboard/icons"); + let mut on_disk: Vec = std::fs::read_dir(&dir) + .unwrap() + .map(|entry| entry.unwrap().file_name().to_string_lossy().into_owned()) + .filter(|name| name.ends_with(".svg")) + .collect(); + on_disk.sort(); + let embedded: Vec = ICONS + .iter() + .map(|(name, _)| format!("{name}.svg")) + .collect(); + assert_eq!(embedded, on_disk); + assert_eq!(ICONS.len(), 69); + for (name, bytes) in ICONS { + let expected = std::fs::read(dir.join(format!("{name}.svg"))).unwrap(); + assert_eq!(*bytes, expected.as_slice(), "{name}"); + } + } + #[test] fn icon_table_is_sorted_for_binary_search() { let mut sorted = ICONS.iter().map(|(name, _)| *name).collect::>(); diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index d44a6bb559..e6e30ac71d 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -14,7 +14,7 @@ use std::time::Duration; use chrono::Utc; use crate::cli::fetch_context::populate_api_region_from_settings; -use crate::core::{CostScanOptions, FetchContext, ProviderId, SourceMode, instantiate_provider}; +use crate::core::{CostScanOptions, FetchContext, ProviderId, instantiate_provider}; use crate::cost_scanner::{self, CostScanner}; use crate::settings::Settings; use crate::spend_contract::build_local_spend_contract_from_summary; @@ -93,7 +93,7 @@ impl SnapshotProducer { for (index, provider_id) in provider_ids.iter().enumerate() { let provider_id = *provider_id; let fetch_timeout = self.fetch_timeout; - let mut ctx = provider_fetch_context(); + let mut ctx = FetchContext::default(); populate_api_region_from_settings(provider_id, &settings, &mut ctx); set.spawn(async move { ( @@ -166,28 +166,6 @@ async fn fetch_provider_envelope( } } -fn provider_fetch_context() -> FetchContext { - FetchContext { - source_mode: SourceMode::Auto, - include_credits: true, - web_timeout: 60, - verbose: false, - manual_cookie_header: None, - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, - workspace_id: None, - seat_credit_entitlement: None, - api_region: None, - gateway_url: None, - auto_prefer_web: false, - browser_cookie_import: false, - requires_optional_usage_completeness: false, - optional_details_enabled: false, - } -} - async fn bounded_fetch( provider_id: ProviderId, ctx: FetchContext, @@ -302,23 +280,8 @@ async fn collect_claude_accounts( &account.token, ); let mut ctx = FetchContext { - source_mode: SourceMode::Auto, - include_credits: true, - web_timeout: 60, - verbose: false, manual_cookie_header: Some(header), - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, - workspace_id: None, - seat_credit_entitlement: None, - api_region: None, - gateway_url: None, - auto_prefer_web: false, - browser_cookie_import: false, - requires_optional_usage_completeness: false, - optional_details_enabled: false, + ..FetchContext::default() }; populate_api_region_from_settings(ProviderId::Claude, settings, &mut ctx); set.spawn(async move { diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index ae212bdc46..67ec6298b7 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -29,7 +29,7 @@ use crate::cli::error_kind::{ }; use crate::cli::fetch_context::populate_api_region_from_settings; use crate::cli::usage::ProviderSelection; -use crate::core::{CostScanOptions, FetchContext, ProviderId, SourceMode, instantiate_provider}; +use crate::core::{CostScanOptions, FetchContext, ProviderId, instantiate_provider}; use crate::cost_reporting_period::CostReportingPeriod; use crate::cost_scanner::{self, CostScanner}; use crate::settings::Settings; @@ -127,26 +127,12 @@ pub(super) async fn usage_response( } }; let ctx = FetchContext { - source_mode: SourceMode::Auto, - include_credits: true, web_timeout: budget.web_timeout_secs(), - verbose: false, - manual_cookie_header: None, - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, - workspace_id: None, - seat_credit_entitlement: None, - api_region: None, - gateway_url: None, - auto_prefer_web: false, - browser_cookie_import: false, // Serve `/usage` is a background poll read: keep the short optional- // join grace (upstream #2583), unlike `codexbar usage` which blocks // for the full completeness window. requires_optional_usage_completeness: false, - optional_details_enabled: false, + ..FetchContext::default() }; let rows = collect_usage_rows( diff --git a/rust/src/cli/serve/metrics/mod.rs b/rust/src/cli/serve/metrics/mod.rs index 47952ddfc8..3d5dd052e0 100644 --- a/rust/src/cli/serve/metrics/mod.rs +++ b/rust/src/cli/serve/metrics/mod.rs @@ -14,16 +14,9 @@ mod snapshot; pub(crate) use snapshot::MetricsSnapshot; -pub(super) fn metrics_response(snapshot: Option<&MetricsSnapshot>) -> String { - rendering::metrics_response(snapshot) -} - -pub(super) fn render_at( - snapshot: &MetricsSnapshot, - now: chrono::DateTime, -) -> Result { - rendering::render_at(snapshot, now) -} +pub(super) use rendering::metrics_response; +#[cfg(test)] +use rendering::render_at; #[cfg(test)] mod tests; diff --git a/rust/src/cli/serve/metrics/rendering.rs b/rust/src/cli/serve/metrics/rendering.rs index 05e33b7777..806a98e1a0 100644 --- a/rust/src/cli/serve/metrics/rendering.rs +++ b/rust/src/cli/serve/metrics/rendering.rs @@ -6,9 +6,7 @@ use super::encoding::MetricsWriter; use super::snapshot::{MetricsSnapshot, QuotaMetric}; use super::{CONTENT_TYPE, MetricsRenderError}; -pub(in crate::cli::serve::metrics) fn metrics_response( - snapshot: Option<&MetricsSnapshot>, -) -> String { +pub(in crate::cli::serve) fn metrics_response(snapshot: Option<&MetricsSnapshot>) -> String { let (status, body) = match snapshot { Some(snapshot) => match render(snapshot) { Ok(body) => (200, body), diff --git a/rust/src/cli/serve/mod.rs b/rust/src/cli/serve/mod.rs index 386fb4eed4..3e0e9011c3 100644 --- a/rust/src/cli/serve/mod.rs +++ b/rust/src/cli/serve/mod.rs @@ -561,6 +561,12 @@ async fn route_request(request: &ServeRequest, config: &ServeConfig) -> String { return json_response(404, serde_json::json!({ "error": "not found" })); }; + let authorized = || { + authorize_request( + request.authorization.as_deref(), + config.token_digest.as_ref(), + ) + }; match route { ServeRoute::DashboardHome => match &config.dashboard { Some(state) => dashboard::home_response(state), @@ -572,10 +578,7 @@ async fn route_request(request: &ServeRequest, config: &ServeConfig) -> String { serde_json::json!({ "status": "ok", "version": env!("CARGO_PKG_VERSION") }), ), ServeRoute::Usage { provider } => { - if !authorize_request( - request.authorization.as_deref(), - config.token_digest.as_ref(), - ) { + if !authorized() { return unauthorized_response(); } let budget = data::RequestBudget::start(config.request_timeout); @@ -586,10 +589,7 @@ async fn route_request(request: &ServeRequest, config: &ServeConfig) -> String { .await } ServeRoute::Cost { provider } => { - if !authorize_request( - request.authorization.as_deref(), - config.token_digest.as_ref(), - ) { + if !authorized() { return unauthorized_response(); } let budget = data::RequestBudget::start(config.request_timeout); @@ -603,10 +603,7 @@ async fn route_request(request: &ServeRequest, config: &ServeConfig) -> String { if !config.metrics_enabled { return json_response(404, serde_json::json!({ "error": "not found" })); } - if !authorize_request( - request.authorization.as_deref(), - config.token_digest.as_ref(), - ) { + if !authorized() { return unauthorized_dashboard_response(); } match &config.dashboard { @@ -621,10 +618,7 @@ async fn route_request(request: &ServeRequest, config: &ServeConfig) -> String { } } ServeRoute::DashboardSnapshot => { - if !authorize_request( - request.authorization.as_deref(), - config.token_digest.as_ref(), - ) { + if !authorized() { return unauthorized_dashboard_response(); } match &config.dashboard { diff --git a/rust/src/cli/serve/tests.rs b/rust/src/cli/serve/tests.rs index 60709dbe07..fda8e9c911 100644 --- a/rust/src/cli/serve/tests.rs +++ b/rust/src/cli/serve/tests.rs @@ -1,5 +1,8 @@ use super::*; +mod dashboard_routes; +mod head_bound; + #[test] fn rejects_non_loopback_hosts_by_default() { assert!(allowed_host("127.0.0.1:8080", "127.0.0.1")); @@ -208,877 +211,6 @@ fn rejects_empty_dashboard_token() { assert!(err.contains("empty")); } -// ── Upstream 0.48.0 #2684: whole-head bound (16 KiB cap + 10 s TOTAL deadline) ── - -use std::time::Instant; - -/// Connected (server, client) TCP pair on loopback. -async fn connected_pair() -> (TcpStream, TcpStream) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let client = TcpStream::connect(addr).await.unwrap(); - let (server, _) = listener.accept().await.unwrap(); - (server, client) -} - -fn head_test_config(budget: Duration, token: Option<&str>) -> ServeConfig { - ServeConfig { - host: "127.0.0.1".to_string(), - port: 8080, - token_digest: token.map(|t| sha256_digest(t.as_bytes())), - metrics_enabled: false, - head_read_budget: budget, - identity: Some(DashboardIdentity::Redacted), - dashboard: None, - request_timeout: None, - operations: data::DataOperations::default(), - } -} - -/// Generous budget for tests that must not trip the deadline. -fn fast_budget() -> Duration { - Duration::from_millis(2_000) -} - -/// Complete request head whose `\r\n\r\n` terminator's final byte is -/// exactly byte 16,384 — the upstream-valid boundary. -fn head_at_exact_cap() -> Vec { - let mut head = String::from("GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nX-Pad: "); - let pad = HEAD_CAP - head.len() - 4; - head.push_str(&"a".repeat(pad)); - head.push_str("\r\n\r\n"); - assert_eq!(head.len(), HEAD_CAP); - head.into_bytes() -} - -/// Send `request`, read until the server closes, return the raw response. -/// Strict outer timeouts turn a hang into a test failure, not a stalled CI. -async fn request_roundtrip(request: &[u8], budget: Duration, token: Option<&str>) -> String { - let (server, mut client) = connected_pair().await; - let config = head_test_config(budget, token); - let server_task = tokio::spawn(async move { handle_client(server, &config).await }); - client.write_all(request).await.unwrap(); - let mut response = Vec::new(); - tokio::time::timeout(Duration::from_secs(10), client.read_to_end(&mut response)) - .await - .expect("client read timed out") - .unwrap(); - // Dropping the client lets the server-side drain finish immediately. - drop(client); - server_task.await.unwrap().unwrap(); - String::from_utf8_lossy(&response).into_owned() -} - -#[test] -fn invalid_request_response_is_pinned() { - let response = invalid_request_response(); - assert!(response.starts_with("HTTP/1.1 400 Bad Request\r\n")); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains("Connection: close\r\n")); - assert!(response.ends_with(r#"{"error":"invalid request"}"#)); -} - -#[test] -fn find_header_end_offsets() { - assert_eq!(find_header_end(b"\r\n\r\n"), Some(4)); - assert_eq!(find_header_end(b"a\r\n\r\n"), Some(5)); - assert_eq!(find_header_end(b"aa\r\n\r\n"), Some(6)); - assert_eq!(find_header_end(b"a\r\n\r"), None); - assert_eq!(find_header_end(b"a\r\n\rXX"), None); - // Terminator straddling a chunk boundary. - assert_eq!(find_header_end(b"abc\r\n\r"), None); - assert_eq!(find_header_end(b"abc\r\n\r\ndef"), Some(7)); -} - -#[tokio::test] -async fn head_reader_accepts_terminator_ending_exactly_at_cap() { - // Upstream boundary: a terminator whose final byte is byte 16,384 is valid. - let (mut server, mut client) = connected_pair().await; - client.write_all(&head_at_exact_cap()).await.unwrap(); - let head = read_request_head(&mut server, fast_budget()).await.unwrap(); - assert_eq!(head.len(), HEAD_CAP); -} - -#[tokio::test] -async fn head_ending_exactly_at_cap_parses_and_routes_normally() { - let response = request_roundtrip(&head_at_exact_cap(), fast_budget(), None).await; - assert!( - response.starts_with("HTTP/1.1 200"), - "exact-cap head must route to /health, got: {response}" - ); -} - -#[tokio::test] -async fn head_reader_rejects_at_cap_without_terminator() { - let (mut server, mut client) = connected_pair().await; - client.write_all(&[b'x'; HEAD_CAP]).await.unwrap(); - let result = read_request_head(&mut server, fast_budget()).await; - assert_eq!(result, Err(HeadReadError::Oversize)); -} - -#[tokio::test] -async fn head_reader_maps_incomplete_eof() { - let (mut server, mut client) = connected_pair().await; - client - .write_all(b"GET /health HTTP/1.1\r\nHost: 127.") - .await - .unwrap(); - client.shutdown().await.unwrap(); - let result = read_request_head(&mut server, fast_budget()).await; - assert_eq!(result, Err(HeadReadError::UnexpectedEof)); -} - -#[tokio::test] -async fn head_reader_maps_total_deadline_on_silent_client() { - let (mut server, _client) = connected_pair().await; - let result = read_request_head(&mut server, Duration::from_millis(150)).await; - assert_eq!(result, Err(HeadReadError::Deadline)); -} - -#[tokio::test] -async fn oversized_head_rejected_before_auth_or_routing() { - // A complete-looking authenticated request line drowned past the cap with - // no terminator: must be rejected before any bearer evaluation. - let mut junk = String::from( - "GET /usage HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\nX-Pad: ", - ); - junk.push_str(&"a".repeat(HEAD_CAP)); - assert!(junk.len() > HEAD_CAP); - let response = request_roundtrip(junk.as_bytes(), fast_budget(), Some("s3cret")).await; - assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); - // Proof the bearer gate / routing never ran: not 401, not the usage payload. - assert!(!response.starts_with("HTTP/1.1 401")); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains("Connection: close\r\n")); - assert!(response.contains(r#""error":"invalid request""#)); -} - -#[tokio::test] -async fn incomplete_head_eof_gets_pinned_400() { - let (server, mut client) = connected_pair().await; - let config = head_test_config(fast_budget(), None); - let server_task = tokio::spawn(async move { handle_client(server, &config).await }); - client - .write_all(b"GET /health HTTP/1.1\r\nHost: 127.") - .await - .unwrap(); - client.shutdown().await.unwrap(); - let mut response = Vec::new(); - client.read_to_end(&mut response).await.unwrap(); - drop(client); - server_task.await.unwrap().unwrap(); - let response = String::from_utf8_lossy(&response); - assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains(r#""error":"invalid request""#)); -} - -#[tokio::test] -async fn silent_client_is_closed_at_total_deadline() { - let budget = Duration::from_millis(250); - let (server, mut client) = connected_pair().await; - let config = head_test_config(budget, None); - let server_task = tokio::spawn(async move { handle_client(server, &config).await }); - let started = Instant::now(); - let mut response = Vec::new(); - client.read_to_end(&mut response).await.unwrap(); - let elapsed = started.elapsed(); - drop(client); - server_task.await.unwrap().unwrap(); - assert!( - elapsed >= budget, - "deadline fired early: {elapsed:?} < {budget:?}" - ); - assert!( - elapsed < Duration::from_secs(2), - "silent client outlived the total deadline: {elapsed:?}" - ); - let response = String::from_utf8_lossy(&response); - assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); -} - -#[tokio::test] -async fn trickling_bytes_do_not_reset_total_head_deadline() { - // One byte every 60 ms: under a per-read timeout this client would hold its - // connection for the full 3 s loop; the 400 ms TOTAL budget must kill it. - // (Red→green mirrored from upstream CLIServeRequestDeadlineLinuxTests.) - let budget = Duration::from_millis(400); - let (server, mut client) = connected_pair().await; - let config = head_test_config(budget, None); - let server_task = tokio::spawn(async move { handle_client(server, &config).await }); - - let started = Instant::now(); - for _ in 0..50 { - tokio::time::sleep(Duration::from_millis(60)).await; - if client.write_all(b"a").await.is_err() { - break; - } - // Stop trickling the moment the server answers or closes. - // peek() does NOT consume bytes — the full response stays readable. - let mut peek = [0_u8; 1]; - if tokio::time::timeout(Duration::from_millis(10), client.peek(&mut peek)) - .await - .is_ok() - { - break; - } - } - let mut response = Vec::new(); - // Best-effort drain to unblock the server; the deadline assertions below are the real check. - let _drained = client.read_to_end(&mut response).await; - let elapsed = started.elapsed(); - drop(client); - server_task.await.unwrap().unwrap(); - - assert!( - elapsed >= budget, - "deadline fired early: {elapsed:?} < {budget:?}" - ); - // Upper ceiling 2.5 s: under a per-read-reset design this client would - // hold the connection for the whole 50-byte loop (~3.5 s incl. peeks), - // so this still fails red — while tolerating full-suite scheduling lag. - assert!( - elapsed < Duration::from_millis(2_500), - "trickling bytes extended the overall deadline: {elapsed:?}" - ); - let response = String::from_utf8_lossy(&response); - assert!( - response.starts_with("HTTP/1.1 400"), - "trickling client must get the pinned 400, got: {response}" - ); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains(r#""error":"invalid request""#)); -} - -#[tokio::test] -async fn authenticated_request_succeeds_and_bad_tokens_stay_401() { - // Deterministic 200: /cost with a provider the local scanner reports as - // unsupported — full auth pass, zero network/disk access. - let ok = request_roundtrip( - b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", - fast_budget(), - Some("s3cret"), - ) - .await; - assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); - assert!(ok.contains("\"supported\":false")); - - let wrong = request_roundtrip( - b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer nope\r\n\r\n", - fast_budget(), - Some("s3cret"), - ) - .await; - assert!(wrong.starts_with("HTTP/1.1 401"), "got: {wrong}"); - - let missing = request_roundtrip( - b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - fast_budget(), - Some("s3cret"), - ) - .await; - assert!(missing.starts_with("HTTP/1.1 401"), "got: {missing}"); -} - -#[tokio::test] -async fn host_gate_unchanged_on_hardened_path() { - let forbidden = request_roundtrip( - b"GET /health HTTP/1.1\r\nHost: example.com\r\n\r\n", - fast_budget(), - None, - ) - .await; - assert!(forbidden.starts_with("HTTP/1.1 403"), "got: {forbidden}"); - assert!(forbidden.contains(r#""error":"forbidden host""#)); - - let ok = request_roundtrip( - b"GET /health HTTP/1.1\r\nHost: localhost:9999\r\n\r\n", - fast_budget(), - None, - ) - .await; - assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); -} - -#[tokio::test] -async fn over_cap_connection_closes_immediately_without_response() { - // Upstream 0.48.0 parity: maximumConnections = 16; slot 17 is closed at - // once, no response bytes, and a freed slot is usable again. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let config = Arc::new(head_test_config(Duration::from_secs(60), None)); - let server_task = tokio::spawn(serve_listener(listener, config, MAX_CONNECTIONS)); - - // Fill every permit with trickling clients that never complete a head. - let mut tricklers = Vec::new(); - for _ in 0..MAX_CONNECTIONS { - let mut client = TcpStream::connect(addr).await.unwrap(); - tricklers.push(tokio::spawn(async move { - loop { - tokio::time::sleep(Duration::from_millis(100)).await; - if client.write_all(b"a").await.is_err() { - break; - } - } - })); - } - - // Probe until the gate is provably full: an over-cap connection gets an - // immediate EOF with zero response bytes. - let mut rejected_seen = false; - for _ in 0..40 { - let mut probe = TcpStream::connect(addr).await.unwrap(); - let mut buf = [0_u8; 16]; - match tokio::time::timeout(Duration::from_millis(300), probe.read(&mut buf)).await { - Ok(Ok(0)) => { - rejected_seen = true; - break; - } - // Probe landed in a still-filling slot; free it and retry. - _ => drop(probe), - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - assert!( - rejected_seen, - "over-cap connection never got the immediate close" - ); - - // Ending the tricklers releases their permits via EOF; a normal client - // must then be served (strict outer timeout). Permit release races the - // server's graceful close-drain window, so a single fixed wait can see a - // connection reset; retry within a bounded budget instead. - for task in &tricklers { - task.abort(); - } - let request = b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n"; - let mut served: Option = None; - let retry = tokio::time::Instant::now(); - while retry.elapsed() < Duration::from_secs(2) { - tokio::time::sleep(Duration::from_millis(100)).await; - let Ok(mut good) = TcpStream::connect(addr).await else { - continue; - }; - if good.write_all(request).await.is_err() { - continue; - } - let mut response = Vec::new(); - match tokio::time::timeout(Duration::from_secs(5), good.read_to_end(&mut response)).await { - // A reset mid-handshake is the drain race; retry. - Ok(Err(_)) | Err(_) => continue, - Ok(Ok(_)) => {} - } - if String::from_utf8_lossy(&response).starts_with("HTTP/1.1 200") { - served = Some(String::from_utf8_lossy(&response).into_owned()); - break; - } - } - let served = served.expect("no freed slot served a normal request within retry budget"); - assert!( - served.starts_with("HTTP/1.1 200"), - "freed slot must serve a normal request, got: {served}" - ); - server_task.abort(); -} - -#[tokio::test] -async fn deadline_driven_release_frees_gated_slot() { - // Regression (review follow-up): a semaphore permit MUST be owned for the - // whole handle_client future and released when the SERVER's total head - // deadline completes its 400/close path — not by client EOF/manual drop. - // The holder client stays connected the entire test. - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let holder_budget = Duration::from_millis(500); - let config = Arc::new(head_test_config(holder_budget, None)); - let server_task = tokio::spawn(serve_listener(listener, config, 1)); - - // Fill the single permit with a holder that never sends a single byte. - let mut holder = TcpStream::connect(addr).await.unwrap(); - - // Synchronize until that permit is provably held: an over-cap probe gets - // an immediate close with zero response bytes. - let mut rejected = false; - for _ in 0..40 { - let mut probe = TcpStream::connect(addr).await.unwrap(); - let mut buf = [0_u8; 16]; - match tokio::time::timeout(Duration::from_millis(300), probe.read(&mut buf)).await { - Ok(Ok(0)) => { - rejected = true; - break; - } - // Probe landed while the holder was still being accepted; retry. - _ => drop(probe), - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - assert!(rejected, "over-cap probe was never closed immediately"); - - // Causality phase: the holder is NOT dropped/aborted/shut down. The - // server's injected total head deadline expires on its own, completing - // the pinned 400/close path. read_to_end returns at the server's FIN; - // the holder socket itself stays OPEN. - let mut holder_response = Vec::new(); - tokio::time::timeout( - Duration::from_secs(3), - holder.read_to_end(&mut holder_response), - ) - .await - .expect("server never drove its deadline/close on the held slot") - .unwrap(); - let holder_response = String::from_utf8_lossy(&holder_response); - assert!( - holder_response.starts_with("HTTP/1.1 400"), - "deadline path must answer the holder with the pinned 400, got: {holder_response}" - ); - assert!(holder_response.contains("Cache-Control: no-store\r\n")); - assert!(holder_response.contains(r#""error":"invalid request""#)); - - // The permit frees only when the server task finishes — after the - // deadline AND the bounded (~1 s) graceful-drain that runs while the - // still-connected holder stays silent. Retry a normal client until the - // freed slot serves it; early retries may still be over-cap closed. - let started = Instant::now(); - let health = loop { - let attempt = async { - let mut good = TcpStream::connect(addr).await.ok()?; - good.write_all(b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n") - .await - .ok()?; - let mut response = Vec::new(); - tokio::time::timeout(Duration::from_millis(800), good.read_to_end(&mut response)) - .await - .ok()? - .ok()?; - Some(String::from_utf8_lossy(&response).into_owned()) - }; - if let Some(text) = attempt.await - && text.starts_with("HTTP/1.1 200") - { - break text; - } - assert!( - started.elapsed() < Duration::from_secs(8), - "permit was never released after the server deadline + graceful drain" - ); - tokio::time::sleep(Duration::from_millis(120)).await; - }; - assert!(health.contains("\"status\":\"ok\""), "got: {health}"); - // Holder is still connected throughout everything above; cleanup only - // after all success assertions. - server_task.abort(); - drop(holder); -} - -// ── Upstream 0.48.0 A1–A5: dashboard routes ─────────────────────────── - -use crate::cli::serve::collection::SnapshotCollection; -use dashboard::coordinator::{SnapshotArtifacts, SnapshotArtifactsBuildFn, SnapshotBuildFn}; -use dashboard::snapshot::{ - AccountFetchEnvelope, ClaudeAccountsInput, DashboardIdentity as DashboardIdMode, - ProviderFetchEnvelope, SnapshotInput, build_snapshot, -}; - -fn stub_build(identity: DashboardIdMode, with_accounts: bool, delay: Duration) -> SnapshotBuildFn { - std::sync::Arc::new(move || { - Box::pin(async move { - if delay > Duration::ZERO { - tokio::time::sleep(delay).await; - } - let mut usage = crate::core::UsageSnapshot::new(crate::core::RateWindow::new(11.0)); - usage.account_email = Some("dev@example.com".to_string()); - usage.login_method = Some("Claude Max".to_string()); - let claude_accounts = with_accounts.then(|| ClaudeAccountsInput { - accounts: Ok(vec![AccountFetchEnvelope { - id: "u-1".to_string(), - label: "Work".to_string(), - active: true, - fetch: Ok(crate::core::ProviderFetchResult::new(usage.clone(), "test")), - }]), - }); - Ok(build_snapshot(&SnapshotInput { - collection: SnapshotCollection { - providers: vec![ProviderFetchEnvelope { - id: "claude".to_string(), - display_name: "Claude".to_string(), - session_label: "Session".to_string(), - weekly_label: "Weekly".to_string(), - fetch: Ok(crate::core::ProviderFetchResult::new(usage, "test")), - }], - costs: std::collections::HashMap::new(), - claude_accounts, - generated_at: chrono::Utc::now(), - refresh_seconds: 60, - order: vec![], - enabled: std::collections::BTreeSet::new(), - }, - identity, - version: Some("test".to_string()), - usage_bars_show_used: None, - })) - }) - }) -} - -fn stub_state_ok() -> dashboard::DashboardState { - let build: SnapshotArtifactsBuildFn = - std::sync::Arc::new(|| { - Box::pin(async move { - let mut usage = crate::core::UsageSnapshot::new( - crate::core::RateWindow::with_details(11.0, Some(300), None, None), - ); - usage.updated_at = chrono::Utc::now(); - let input = SnapshotInput { - collection: SnapshotCollection { - providers: vec![ProviderFetchEnvelope { - id: "codex".to_string(), - display_name: "Codex".to_string(), - session_label: "Session".to_string(), - weekly_label: "Weekly".to_string(), - fetch: Ok(crate::core::ProviderFetchResult::new(usage, "test")), - }], - costs: std::collections::HashMap::new(), - claude_accounts: None, - generated_at: chrono::Utc::now(), - refresh_seconds: 60, - order: vec!["codex".to_string()], - enabled: ["codex".to_string()].into_iter().collect(), - }, - identity: DashboardIdMode::Redacted, - version: Some("test".to_string()), - usage_bars_show_used: None, - }; - Ok(SnapshotArtifacts { - dashboard: build_snapshot(&input), - sidecar: Some(metrics::MetricsSnapshot::from_collection(&input.collection)), - }) - }) - }); - dashboard::DashboardState::stub_with_artifacts(build, 3600, Some(DashboardIdMode::Redacted)) -} - -fn dashboard_test_config( - token: Option<&str>, - state: Option, -) -> ServeConfig { - let mut config = head_test_config(fast_budget(), token); - config.dashboard = state; - config -} - -#[test] -fn resolve_route_maps_paths() { - let req = |path: &str, query: &[(&str, &str)]| { - let mut request = - parse_request(&format!("GET {path} HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n")).unwrap(); - for (k, v) in query { - request.query.insert(k.to_string(), v.to_string()); - } - request - }; - assert_eq!( - resolve_route(&req("/", &[])), - Some(ServeRoute::DashboardHome) - ); - assert_eq!( - resolve_route(&req("/health", &[])), - Some(ServeRoute::Health) - ); - assert_eq!( - resolve_route(&req("/usage?provider=codex", &[])), - Some(ServeRoute::Usage { - provider: Some("codex".to_string()) - }) - ); - assert_eq!( - resolve_route(&req("/metrics", &[])), - Some(ServeRoute::Metrics) - ); - assert_eq!( - resolve_route(&req("/dashboard/v1/snapshot", &[])), - Some(ServeRoute::DashboardSnapshot) - ); - assert_eq!( - resolve_route(&req("/icons/ProviderIcon-codex.svg", &[])), - Some(ServeRoute::ProviderIcon { - name: "ProviderIcon-codex".to_string() - }) - ); - assert_eq!(resolve_route(&req("/icons/../x.svg", &[])), None); - assert_eq!(resolve_route(&req("/icons/.svg", &[])), None); - assert_eq!(resolve_route(&req("/dashboard/v1/other", &[])), None); - assert_eq!(resolve_route(&req("/usage.json", &[])), None); -} - -#[tokio::test] -async fn metrics_route_is_not_found_until_enabled() { - let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - let response = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", - config, - ) - .await; - assert!(response.starts_with("HTTP/1.1 404"), "got: {response}"); - assert!(response.contains(r#""error":"not found""#)); -} - -#[tokio::test] -async fn metrics_route_uses_bearer_gate_and_prometheus_content_type() { - let mut missing_config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - missing_config.metrics_enabled = true; - let missing = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - missing_config, - ) - .await; - assert!(missing.starts_with("HTTP/1.1 401"), "got: {missing}"); - assert!(missing.contains("WWW-Authenticate: Bearer\r\n")); - - let mut wrong_config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - wrong_config.metrics_enabled = true; - let wrong = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer nope\r\n\r\n", - wrong_config, - ) - .await; - assert!(wrong.starts_with("HTTP/1.1 401"), "got: {wrong}"); - - let ready_state = stub_state_ok(); - ready_state.coordinator.get().await.unwrap(); - let mut ok_config = dashboard_test_config(Some("s3cret"), Some(ready_state)); - ok_config.metrics_enabled = true; - let ok = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", - ok_config, - ) - .await; - assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); - assert!(ok.contains("Content-Type: text/plain; version=0.0.4; charset=utf-8\r\n")); - assert!(ok.contains("Cache-Control: no-store\r\n")); - assert!( - ok.contains("codexbar_snapshot_schema_version 1\n"), - "got: {ok}" - ); - assert!( - ok.contains("codexbar_quota_session_used_ratio{provider=\"codex\"}"), - "got: {ok}" - ); -} - -#[tokio::test] -async fn metrics_route_returns_500_when_exporter_state_is_missing() { - let mut config = dashboard_test_config(Some("s3cret"), None); - config.metrics_enabled = true; - let response = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", - config, - ) - .await; - - assert!(response.starts_with("HTTP/1.1 500"), "got: {response}"); - assert!(response.contains(r#""error":"dashboard not configured""#)); -} - -#[tokio::test] -async fn metrics_route_keeps_the_host_allowlist() { - let mut config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - config.metrics_enabled = true; - config.host = "192.0.2.10".to_string(); - let forbidden = request_roundtrip_dashboard( - b"GET /metrics HTTP/1.1\r\nHost: 192.0.2.11:8080\r\nAuthorization: Bearer s3cret\r\n\r\n", - config, - ) - .await; - assert!(forbidden.starts_with("HTTP/1.1 403"), "got: {forbidden}"); - assert!(forbidden.contains(r#""error":"forbidden host""#)); -} - -#[tokio::test] -async fn dashboard_home_serves_html_no_store() { - let config = dashboard_test_config(None, Some(stub_state_ok())); - let response = - request_roundtrip_dashboard(b"GET / HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", config).await; - assert!( - response.starts_with("HTTP/1.1 200"), - "got: {}", - &response[..80.min(response.len())] - ); - assert!(response.contains("Content-Type: text/html; charset=utf-8\r\n")); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains("/dashboard/v1/snapshot")); - assert!(response.contains("ProviderIcon-codex.svg")); -} - -#[tokio::test] -async fn icon_route_serves_svg_immutable_and_404s_unknown() { - let config = dashboard_test_config(None, None); - let response = request_roundtrip_dashboard( - b"GET /icons/ProviderIcon-codex.svg HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!( - response.starts_with("HTTP/1.1 200"), - "got: {}", - &response[..80.min(response.len())] - ); - assert!(response.contains("Content-Type: image/svg+xml\r\n")); - assert!(response.contains("Cache-Control: public, max-age=86400, immutable\r\n")); - assert!(response.contains(" 401 + WWW-Authenticate (upstream dashboard-rule parity). - let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - let unauthorized = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!( - unauthorized.starts_with("HTTP/1.1 401"), - "got: {unauthorized}" - ); - assert!(unauthorized.contains("WWW-Authenticate: Bearer\r\n")); - - // Valid token -> 200 + schema v1 payload + no-store. - let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - let response = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", - config, - ) - .await; - assert!(response.starts_with("HTTP/1.1 200"), "got: {response}"); - assert!(response.contains("Cache-Control: no-store\r\n")); - assert!(response.contains("\"schemaVersion\": 1")); - assert!(response.contains("\"providers\"")); -} - -#[tokio::test] -async fn dashboard_home_and_icons_are_public_when_token_configured() { - let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - let response = - request_roundtrip_dashboard(b"GET / HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", config).await; - assert!( - response.starts_with("HTTP/1.1 200"), - "shell stays public: {response}" - ); - let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); - let icon = request_roundtrip_dashboard( - b"GET /icons/ProviderIcon-claude.svg HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!( - icon.starts_with("HTTP/1.1 200"), - "icons stay public: {icon}" - ); - assert!(!icon.contains("WWW-Authenticate")); -} - -#[tokio::test] -async fn snapshot_identity_modes_redact_or_expose() { - // Redacted (default): `redacted@domain`, raw address never leaks. - let state = dashboard::DashboardState::stub( - stub_build(DashboardIdMode::Redacted, false, Duration::ZERO), - 3600, - Some(DashboardIdMode::Redacted), - ); - let config = dashboard_test_config(None, Some(state)); - let redacted = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!(redacted.contains("redacted@example.com"), "got: {redacted}"); - assert!( - !redacted.contains("dev@example.com"), - "raw email leaked: {redacted}" - ); - - // Full opt-in: real account email exposed. - let state = dashboard::DashboardState::stub( - stub_build(DashboardIdMode::Full, false, Duration::ZERO), - 3600, - Some(DashboardIdMode::Full), - ); - let config = dashboard_test_config(None, Some(state)); - let full = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!(full.contains("dev@example.com"), "got: {full}"); -} - -#[tokio::test] -async fn snapshot_claude_accounts_nest_under_claude_row() { - let state = dashboard::DashboardState::stub( - stub_build(DashboardIdMode::Redacted, true, Duration::ZERO), - 3600, - Some(DashboardIdMode::Redacted), - ); - let config = dashboard_test_config(None, Some(state)); - let response = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - assert!(response.contains("\"accounts\""), "got: {response}"); - assert!(response.contains("\"label\": \"Work\""), "got: {response}"); - assert!(response.contains("\"active\": true"), "got: {response}"); - assert!(response.contains("redacted@example.com")); -} - -#[tokio::test] -async fn snapshot_late_build_is_delivered_not_discarded() { - // F9/2717 parity: a slow snapshot build completes and the response carries - // the finished result — never a discarded-build error. - let state = dashboard::DashboardState::stub( - stub_build(DashboardIdMode::Redacted, false, Duration::from_millis(250)), - 3600, - Some(DashboardIdMode::Redacted), - ); - let config = dashboard_test_config(None, Some(state)); - let started = std::time::Instant::now(); - let response = request_roundtrip_dashboard( - b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", - config, - ) - .await; - let elapsed = started.elapsed(); - assert!(response.starts_with("HTTP/1.1 200"), "got: {response}"); - assert!(response.contains("\"schemaVersion\": 1")); - assert!( - elapsed >= Duration::from_millis(250), - "response arrived before the build finished: {elapsed:?}" - ); -} - -/// Roundtrip helper for route-level tests (separate from head-level helper). -async fn request_roundtrip_dashboard(request: &[u8], config: ServeConfig) -> String { - let (server, mut client) = connected_pair().await; - let server_task = tokio::spawn(async move { handle_client(server, &config).await }); - client.write_all(request).await.unwrap(); - let mut response = Vec::new(); - tokio::time::timeout(Duration::from_secs(10), client.read_to_end(&mut response)) - .await - .expect("client read timed out") - .unwrap(); - drop(client); - server_task.await.unwrap().unwrap(); - String::from_utf8_lossy(&response).into_owned() -} - #[test] fn usage_route_without_provider_follows_enabled_providers() { use crate::cli::usage::ProviderSelection; diff --git a/rust/src/cli/serve/tests/dashboard_routes.rs b/rust/src/cli/serve/tests/dashboard_routes.rs new file mode 100644 index 0000000000..40edaa6835 --- /dev/null +++ b/rust/src/cli/serve/tests/dashboard_routes.rs @@ -0,0 +1,410 @@ +//! Upstream 0.48.0 A1–A5: dashboard routes. + +use super::head_bound::{connected_pair, fast_budget, head_test_config}; +use super::*; +use crate::cli::serve::collection::SnapshotCollection; +use dashboard::coordinator::{SnapshotArtifacts, SnapshotArtifactsBuildFn, SnapshotBuildFn}; +use dashboard::snapshot::{ + AccountFetchEnvelope, ClaudeAccountsInput, DashboardIdentity as DashboardIdMode, + ProviderFetchEnvelope, SnapshotInput, build_snapshot, +}; + +fn stub_build(identity: DashboardIdMode, with_accounts: bool, delay: Duration) -> SnapshotBuildFn { + std::sync::Arc::new(move || { + Box::pin(async move { + if delay > Duration::ZERO { + tokio::time::sleep(delay).await; + } + let mut usage = crate::core::UsageSnapshot::new(crate::core::RateWindow::new(11.0)); + usage.account_email = Some("dev@example.com".to_string()); + usage.login_method = Some("Claude Max".to_string()); + let claude_accounts = with_accounts.then(|| ClaudeAccountsInput { + accounts: Ok(vec![AccountFetchEnvelope { + id: "u-1".to_string(), + label: "Work".to_string(), + active: true, + fetch: Ok(crate::core::ProviderFetchResult::new(usage.clone(), "test")), + }]), + }); + Ok(build_snapshot(&SnapshotInput { + collection: SnapshotCollection { + providers: vec![ProviderFetchEnvelope { + id: "claude".to_string(), + display_name: "Claude".to_string(), + session_label: "Session".to_string(), + weekly_label: "Weekly".to_string(), + fetch: Ok(crate::core::ProviderFetchResult::new(usage, "test")), + }], + costs: std::collections::HashMap::new(), + claude_accounts, + generated_at: chrono::Utc::now(), + refresh_seconds: 60, + order: vec![], + enabled: std::collections::BTreeSet::new(), + }, + identity, + version: Some("test".to_string()), + usage_bars_show_used: None, + })) + }) + }) +} + +fn stub_state_ok() -> dashboard::DashboardState { + let build: SnapshotArtifactsBuildFn = + std::sync::Arc::new(|| { + Box::pin(async move { + let mut usage = crate::core::UsageSnapshot::new( + crate::core::RateWindow::with_details(11.0, Some(300), None, None), + ); + usage.updated_at = chrono::Utc::now(); + let input = SnapshotInput { + collection: SnapshotCollection { + providers: vec![ProviderFetchEnvelope { + id: "codex".to_string(), + display_name: "Codex".to_string(), + session_label: "Session".to_string(), + weekly_label: "Weekly".to_string(), + fetch: Ok(crate::core::ProviderFetchResult::new(usage, "test")), + }], + costs: std::collections::HashMap::new(), + claude_accounts: None, + generated_at: chrono::Utc::now(), + refresh_seconds: 60, + order: vec!["codex".to_string()], + enabled: ["codex".to_string()].into_iter().collect(), + }, + identity: DashboardIdMode::Redacted, + version: Some("test".to_string()), + usage_bars_show_used: None, + }; + Ok(SnapshotArtifacts { + dashboard: build_snapshot(&input), + sidecar: Some(metrics::MetricsSnapshot::from_collection(&input.collection)), + }) + }) + }); + dashboard::DashboardState::stub_with_artifacts(build, 3600, Some(DashboardIdMode::Redacted)) +} + +fn dashboard_test_config( + token: Option<&str>, + state: Option, +) -> ServeConfig { + let mut config = head_test_config(fast_budget(), token); + config.dashboard = state; + config +} + +#[test] +fn resolve_route_maps_paths() { + let req = |path: &str, query: &[(&str, &str)]| { + let mut request = + parse_request(&format!("GET {path} HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n")).unwrap(); + for (k, v) in query { + request.query.insert(k.to_string(), v.to_string()); + } + request + }; + assert_eq!( + resolve_route(&req("/", &[])), + Some(ServeRoute::DashboardHome) + ); + assert_eq!( + resolve_route(&req("/health", &[])), + Some(ServeRoute::Health) + ); + assert_eq!( + resolve_route(&req("/usage?provider=codex", &[])), + Some(ServeRoute::Usage { + provider: Some("codex".to_string()) + }) + ); + assert_eq!( + resolve_route(&req("/metrics", &[])), + Some(ServeRoute::Metrics) + ); + assert_eq!( + resolve_route(&req("/dashboard/v1/snapshot", &[])), + Some(ServeRoute::DashboardSnapshot) + ); + assert_eq!( + resolve_route(&req("/icons/ProviderIcon-codex.svg", &[])), + Some(ServeRoute::ProviderIcon { + name: "ProviderIcon-codex".to_string() + }) + ); + assert_eq!(resolve_route(&req("/icons/../x.svg", &[])), None); + assert_eq!(resolve_route(&req("/icons/.svg", &[])), None); + assert_eq!(resolve_route(&req("/dashboard/v1/other", &[])), None); + assert_eq!(resolve_route(&req("/usage.json", &[])), None); +} + +#[tokio::test] +async fn metrics_route_is_not_found_until_enabled() { + let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + let response = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", + config, + ) + .await; + assert!(response.starts_with("HTTP/1.1 404"), "got: {response}"); + assert!(response.contains(r#""error":"not found""#)); +} + +#[tokio::test] +async fn metrics_route_uses_bearer_gate_and_prometheus_content_type() { + let mut missing_config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + missing_config.metrics_enabled = true; + let missing = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + missing_config, + ) + .await; + assert!(missing.starts_with("HTTP/1.1 401"), "got: {missing}"); + assert!(missing.contains("WWW-Authenticate: Bearer\r\n")); + + let mut wrong_config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + wrong_config.metrics_enabled = true; + let wrong = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer nope\r\n\r\n", + wrong_config, + ) + .await; + assert!(wrong.starts_with("HTTP/1.1 401"), "got: {wrong}"); + + let ready_state = stub_state_ok(); + ready_state.coordinator.get().await.unwrap(); + let mut ok_config = dashboard_test_config(Some("s3cret"), Some(ready_state)); + ok_config.metrics_enabled = true; + let ok = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", + ok_config, + ) + .await; + assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); + assert!(ok.contains("Content-Type: text/plain; version=0.0.4; charset=utf-8\r\n")); + assert!(ok.contains("Cache-Control: no-store\r\n")); + assert!( + ok.contains("codexbar_snapshot_schema_version 1\n"), + "got: {ok}" + ); + assert!( + ok.contains("codexbar_quota_session_used_ratio{provider=\"codex\"}"), + "got: {ok}" + ); +} + +#[tokio::test] +async fn metrics_route_returns_500_when_exporter_state_is_missing() { + let mut config = dashboard_test_config(Some("s3cret"), None); + config.metrics_enabled = true; + let response = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", + config, + ) + .await; + + assert!(response.starts_with("HTTP/1.1 500"), "got: {response}"); + assert!(response.contains(r#""error":"dashboard not configured""#)); +} + +#[tokio::test] +async fn metrics_route_keeps_the_host_allowlist() { + let mut config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + config.metrics_enabled = true; + config.host = "192.0.2.10".to_string(); + let forbidden = request_roundtrip_dashboard( + b"GET /metrics HTTP/1.1\r\nHost: 192.0.2.11:8080\r\nAuthorization: Bearer s3cret\r\n\r\n", + config, + ) + .await; + assert!(forbidden.starts_with("HTTP/1.1 403"), "got: {forbidden}"); + assert!(forbidden.contains(r#""error":"forbidden host""#)); +} + +#[tokio::test] +async fn dashboard_home_serves_html_no_store() { + let config = dashboard_test_config(None, Some(stub_state_ok())); + let response = + request_roundtrip_dashboard(b"GET / HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", config).await; + assert!( + response.starts_with("HTTP/1.1 200"), + "got: {}", + &response[..80.min(response.len())] + ); + assert!(response.contains("Content-Type: text/html; charset=utf-8\r\n")); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains("/dashboard/v1/snapshot")); + assert!(response.contains("ProviderIcon-codex.svg")); +} + +#[tokio::test] +async fn icon_route_serves_svg_immutable_and_404s_unknown() { + let config = dashboard_test_config(None, None); + let response = request_roundtrip_dashboard( + b"GET /icons/ProviderIcon-codex.svg HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!( + response.starts_with("HTTP/1.1 200"), + "got: {}", + &response[..80.min(response.len())] + ); + assert!(response.contains("Content-Type: image/svg+xml\r\n")); + assert!(response.contains("Cache-Control: public, max-age=86400, immutable\r\n")); + assert!(response.contains(" 401 + WWW-Authenticate (upstream dashboard-rule parity). + let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + let unauthorized = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!( + unauthorized.starts_with("HTTP/1.1 401"), + "got: {unauthorized}" + ); + assert!(unauthorized.contains("WWW-Authenticate: Bearer\r\n")); + + // Valid token -> 200 + schema v1 payload + no-store. + let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + let response = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", + config, + ) + .await; + assert!(response.starts_with("HTTP/1.1 200"), "got: {response}"); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains("\"schemaVersion\": 1")); + assert!(response.contains("\"providers\"")); +} + +#[tokio::test] +async fn dashboard_home_and_icons_are_public_when_token_configured() { + let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + let response = + request_roundtrip_dashboard(b"GET / HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", config).await; + assert!( + response.starts_with("HTTP/1.1 200"), + "shell stays public: {response}" + ); + let config = dashboard_test_config(Some("s3cret"), Some(stub_state_ok())); + let icon = request_roundtrip_dashboard( + b"GET /icons/ProviderIcon-claude.svg HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!( + icon.starts_with("HTTP/1.1 200"), + "icons stay public: {icon}" + ); + assert!(!icon.contains("WWW-Authenticate")); +} + +#[tokio::test] +async fn snapshot_identity_modes_redact_or_expose() { + // Redacted (default): `redacted@domain`, raw address never leaks. + let state = dashboard::DashboardState::stub( + stub_build(DashboardIdMode::Redacted, false, Duration::ZERO), + 3600, + Some(DashboardIdMode::Redacted), + ); + let config = dashboard_test_config(None, Some(state)); + let redacted = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!(redacted.contains("redacted@example.com"), "got: {redacted}"); + assert!( + !redacted.contains("dev@example.com"), + "raw email leaked: {redacted}" + ); + + // Full opt-in: real account email exposed. + let state = dashboard::DashboardState::stub( + stub_build(DashboardIdMode::Full, false, Duration::ZERO), + 3600, + Some(DashboardIdMode::Full), + ); + let config = dashboard_test_config(None, Some(state)); + let full = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!(full.contains("dev@example.com"), "got: {full}"); +} + +#[tokio::test] +async fn snapshot_claude_accounts_nest_under_claude_row() { + let state = dashboard::DashboardState::stub( + stub_build(DashboardIdMode::Redacted, true, Duration::ZERO), + 3600, + Some(DashboardIdMode::Redacted), + ); + let config = dashboard_test_config(None, Some(state)); + let response = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + assert!(response.contains("\"accounts\""), "got: {response}"); + assert!(response.contains("\"label\": \"Work\""), "got: {response}"); + assert!(response.contains("\"active\": true"), "got: {response}"); + assert!(response.contains("redacted@example.com")); +} + +#[tokio::test] +async fn snapshot_late_build_is_delivered_not_discarded() { + // F9/2717 parity: a slow snapshot build completes and the response carries + // the finished result — never a discarded-build error. + let state = dashboard::DashboardState::stub( + stub_build(DashboardIdMode::Redacted, false, Duration::from_millis(250)), + 3600, + Some(DashboardIdMode::Redacted), + ); + let config = dashboard_test_config(None, Some(state)); + let started = std::time::Instant::now(); + let response = request_roundtrip_dashboard( + b"GET /dashboard/v1/snapshot HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + config, + ) + .await; + let elapsed = started.elapsed(); + assert!(response.starts_with("HTTP/1.1 200"), "got: {response}"); + assert!(response.contains("\"schemaVersion\": 1")); + assert!( + elapsed >= Duration::from_millis(250), + "response arrived before the build finished: {elapsed:?}" + ); +} + +/// Roundtrip helper for route-level tests (separate from head-level helper). +async fn request_roundtrip_dashboard(request: &[u8], config: ServeConfig) -> String { + let (server, mut client) = connected_pair().await; + let server_task = tokio::spawn(async move { handle_client(server, &config).await }); + client.write_all(request).await.unwrap(); + let mut response = Vec::new(); + tokio::time::timeout(Duration::from_secs(10), client.read_to_end(&mut response)) + .await + .expect("client read timed out") + .unwrap(); + drop(client); + server_task.await.unwrap().unwrap(); + String::from_utf8_lossy(&response).into_owned() +} diff --git a/rust/src/cli/serve/tests/head_bound.rs b/rust/src/cli/serve/tests/head_bound.rs new file mode 100644 index 0000000000..23b89925c7 --- /dev/null +++ b/rust/src/cli/serve/tests/head_bound.rs @@ -0,0 +1,462 @@ +//! Upstream 0.48.0 #2684: whole-head bound (16 KiB cap + 10 s TOTAL deadline). + +use super::*; +use std::time::Instant; + +/// Connected (server, client) TCP pair on loopback. +pub(super) async fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let client = TcpStream::connect(addr).await.unwrap(); + let (server, _) = listener.accept().await.unwrap(); + (server, client) +} + +pub(super) fn head_test_config(budget: Duration, token: Option<&str>) -> ServeConfig { + ServeConfig { + host: "127.0.0.1".to_string(), + port: 8080, + token_digest: token.map(|t| sha256_digest(t.as_bytes())), + metrics_enabled: false, + head_read_budget: budget, + identity: Some(DashboardIdentity::Redacted), + dashboard: None, + request_timeout: None, + operations: data::DataOperations::default(), + } +} + +/// Generous budget for tests that must not trip the deadline. +pub(super) fn fast_budget() -> Duration { + Duration::from_millis(2_000) +} + +/// Complete request head whose `\r\n\r\n` terminator's final byte is +/// exactly byte 16,384 — the upstream-valid boundary. +fn head_at_exact_cap() -> Vec { + let mut head = String::from("GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nX-Pad: "); + let pad = HEAD_CAP - head.len() - 4; + head.push_str(&"a".repeat(pad)); + head.push_str("\r\n\r\n"); + assert_eq!(head.len(), HEAD_CAP); + head.into_bytes() +} + +/// Send `request`, read until the server closes, return the raw response. +/// Strict outer timeouts turn a hang into a test failure, not a stalled CI. +async fn request_roundtrip(request: &[u8], budget: Duration, token: Option<&str>) -> String { + let (server, mut client) = connected_pair().await; + let config = head_test_config(budget, token); + let server_task = tokio::spawn(async move { handle_client(server, &config).await }); + client.write_all(request).await.unwrap(); + let mut response = Vec::new(); + tokio::time::timeout(Duration::from_secs(10), client.read_to_end(&mut response)) + .await + .expect("client read timed out") + .unwrap(); + // Dropping the client lets the server-side drain finish immediately. + drop(client); + server_task.await.unwrap().unwrap(); + String::from_utf8_lossy(&response).into_owned() +} + +#[test] +fn invalid_request_response_is_pinned() { + let response = invalid_request_response(); + assert!(response.starts_with("HTTP/1.1 400 Bad Request\r\n")); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains("Connection: close\r\n")); + assert!(response.ends_with(r#"{"error":"invalid request"}"#)); +} + +#[test] +fn find_header_end_offsets() { + assert_eq!(find_header_end(b"\r\n\r\n"), Some(4)); + assert_eq!(find_header_end(b"a\r\n\r\n"), Some(5)); + assert_eq!(find_header_end(b"aa\r\n\r\n"), Some(6)); + assert_eq!(find_header_end(b"a\r\n\r"), None); + assert_eq!(find_header_end(b"a\r\n\rXX"), None); + // Terminator straddling a chunk boundary. + assert_eq!(find_header_end(b"abc\r\n\r"), None); + assert_eq!(find_header_end(b"abc\r\n\r\ndef"), Some(7)); +} + +#[tokio::test] +async fn head_reader_accepts_terminator_ending_exactly_at_cap() { + // Upstream boundary: a terminator whose final byte is byte 16,384 is valid. + let (mut server, mut client) = connected_pair().await; + client.write_all(&head_at_exact_cap()).await.unwrap(); + let head = read_request_head(&mut server, fast_budget()).await.unwrap(); + assert_eq!(head.len(), HEAD_CAP); +} + +#[tokio::test] +async fn head_ending_exactly_at_cap_parses_and_routes_normally() { + let response = request_roundtrip(&head_at_exact_cap(), fast_budget(), None).await; + assert!( + response.starts_with("HTTP/1.1 200"), + "exact-cap head must route to /health, got: {response}" + ); +} + +#[tokio::test] +async fn head_reader_rejects_at_cap_without_terminator() { + let (mut server, mut client) = connected_pair().await; + client.write_all(&[b'x'; HEAD_CAP]).await.unwrap(); + let result = read_request_head(&mut server, fast_budget()).await; + assert_eq!(result, Err(HeadReadError::Oversize)); +} + +#[tokio::test] +async fn head_reader_maps_incomplete_eof() { + let (mut server, mut client) = connected_pair().await; + client + .write_all(b"GET /health HTTP/1.1\r\nHost: 127.") + .await + .unwrap(); + client.shutdown().await.unwrap(); + let result = read_request_head(&mut server, fast_budget()).await; + assert_eq!(result, Err(HeadReadError::UnexpectedEof)); +} + +#[tokio::test] +async fn head_reader_maps_total_deadline_on_silent_client() { + let (mut server, _client) = connected_pair().await; + let result = read_request_head(&mut server, Duration::from_millis(150)).await; + assert_eq!(result, Err(HeadReadError::Deadline)); +} + +#[tokio::test] +async fn oversized_head_rejected_before_auth_or_routing() { + // A complete-looking authenticated request line drowned past the cap with + // no terminator: must be rejected before any bearer evaluation. + let mut junk = String::from( + "GET /usage HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\nX-Pad: ", + ); + junk.push_str(&"a".repeat(HEAD_CAP)); + assert!(junk.len() > HEAD_CAP); + let response = request_roundtrip(junk.as_bytes(), fast_budget(), Some("s3cret")).await; + assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); + // Proof the bearer gate / routing never ran: not 401, not the usage payload. + assert!(!response.starts_with("HTTP/1.1 401")); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains("Connection: close\r\n")); + assert!(response.contains(r#""error":"invalid request""#)); +} + +#[tokio::test] +async fn incomplete_head_eof_gets_pinned_400() { + let (server, mut client) = connected_pair().await; + let config = head_test_config(fast_budget(), None); + let server_task = tokio::spawn(async move { handle_client(server, &config).await }); + client + .write_all(b"GET /health HTTP/1.1\r\nHost: 127.") + .await + .unwrap(); + client.shutdown().await.unwrap(); + let mut response = Vec::new(); + client.read_to_end(&mut response).await.unwrap(); + drop(client); + server_task.await.unwrap().unwrap(); + let response = String::from_utf8_lossy(&response); + assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains(r#""error":"invalid request""#)); +} + +#[tokio::test] +async fn silent_client_is_closed_at_total_deadline() { + let budget = Duration::from_millis(250); + let (server, mut client) = connected_pair().await; + let config = head_test_config(budget, None); + let server_task = tokio::spawn(async move { handle_client(server, &config).await }); + let started = Instant::now(); + let mut response = Vec::new(); + client.read_to_end(&mut response).await.unwrap(); + let elapsed = started.elapsed(); + drop(client); + server_task.await.unwrap().unwrap(); + assert!( + elapsed >= budget, + "deadline fired early: {elapsed:?} < {budget:?}" + ); + assert!( + elapsed < Duration::from_secs(2), + "silent client outlived the total deadline: {elapsed:?}" + ); + let response = String::from_utf8_lossy(&response); + assert!(response.starts_with("HTTP/1.1 400"), "got: {response}"); +} + +#[tokio::test] +async fn trickling_bytes_do_not_reset_total_head_deadline() { + // One byte every 60 ms: under a per-read timeout this client would hold its + // connection for the full 3 s loop; the 400 ms TOTAL budget must kill it. + // (Red→green mirrored from upstream CLIServeRequestDeadlineLinuxTests.) + let budget = Duration::from_millis(400); + let (server, mut client) = connected_pair().await; + let config = head_test_config(budget, None); + let server_task = tokio::spawn(async move { handle_client(server, &config).await }); + + let started = Instant::now(); + for _ in 0..50 { + tokio::time::sleep(Duration::from_millis(60)).await; + if client.write_all(b"a").await.is_err() { + break; + } + // Stop trickling the moment the server answers or closes. + // peek() does NOT consume bytes — the full response stays readable. + let mut peek = [0_u8; 1]; + if tokio::time::timeout(Duration::from_millis(10), client.peek(&mut peek)) + .await + .is_ok() + { + break; + } + } + let mut response = Vec::new(); + // Best-effort drain to unblock the server; the deadline assertions below are the real check. + let _drained = client.read_to_end(&mut response).await; + let elapsed = started.elapsed(); + drop(client); + server_task.await.unwrap().unwrap(); + + assert!( + elapsed >= budget, + "deadline fired early: {elapsed:?} < {budget:?}" + ); + // Upper ceiling 2.5 s: under a per-read-reset design this client would + // hold the connection for the whole 50-byte loop (~3.5 s incl. peeks), + // so this still fails red — while tolerating full-suite scheduling lag. + assert!( + elapsed < Duration::from_millis(2_500), + "trickling bytes extended the overall deadline: {elapsed:?}" + ); + let response = String::from_utf8_lossy(&response); + assert!( + response.starts_with("HTTP/1.1 400"), + "trickling client must get the pinned 400, got: {response}" + ); + assert!(response.contains("Cache-Control: no-store\r\n")); + assert!(response.contains(r#""error":"invalid request""#)); +} + +#[tokio::test] +async fn authenticated_request_succeeds_and_bad_tokens_stay_401() { + // Deterministic 200: /cost with a provider the local scanner reports as + // unsupported — full auth pass, zero network/disk access. + let ok = request_roundtrip( + b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer s3cret\r\n\r\n", + fast_budget(), + Some("s3cret"), + ) + .await; + assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); + assert!(ok.contains("\"supported\":false")); + + let wrong = request_roundtrip( + b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Bearer nope\r\n\r\n", + fast_budget(), + Some("s3cret"), + ) + .await; + assert!(wrong.starts_with("HTTP/1.1 401"), "got: {wrong}"); + + let missing = request_roundtrip( + b"GET /cost?provider=gemini HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n", + fast_budget(), + Some("s3cret"), + ) + .await; + assert!(missing.starts_with("HTTP/1.1 401"), "got: {missing}"); +} + +#[tokio::test] +async fn host_gate_unchanged_on_hardened_path() { + let forbidden = request_roundtrip( + b"GET /health HTTP/1.1\r\nHost: example.com\r\n\r\n", + fast_budget(), + None, + ) + .await; + assert!(forbidden.starts_with("HTTP/1.1 403"), "got: {forbidden}"); + assert!(forbidden.contains(r#""error":"forbidden host""#)); + + let ok = request_roundtrip( + b"GET /health HTTP/1.1\r\nHost: localhost:9999\r\n\r\n", + fast_budget(), + None, + ) + .await; + assert!(ok.starts_with("HTTP/1.1 200"), "got: {ok}"); +} + +#[tokio::test] +async fn over_cap_connection_closes_immediately_without_response() { + // Upstream 0.48.0 parity: maximumConnections = 16; slot 17 is closed at + // once, no response bytes, and a freed slot is usable again. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let config = Arc::new(head_test_config(Duration::from_secs(60), None)); + let server_task = tokio::spawn(serve_listener(listener, config, MAX_CONNECTIONS)); + + // Fill every permit with trickling clients that never complete a head. + let mut tricklers = Vec::new(); + for _ in 0..MAX_CONNECTIONS { + let mut client = TcpStream::connect(addr).await.unwrap(); + tricklers.push(tokio::spawn(async move { + loop { + tokio::time::sleep(Duration::from_millis(100)).await; + if client.write_all(b"a").await.is_err() { + break; + } + } + })); + } + + // Probe until the gate is provably full: an over-cap connection gets an + // immediate EOF with zero response bytes. + let mut rejected_seen = false; + for _ in 0..40 { + let mut probe = TcpStream::connect(addr).await.unwrap(); + let mut buf = [0_u8; 16]; + match tokio::time::timeout(Duration::from_millis(300), probe.read(&mut buf)).await { + Ok(Ok(0)) => { + rejected_seen = true; + break; + } + // Probe landed in a still-filling slot; free it and retry. + _ => drop(probe), + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!( + rejected_seen, + "over-cap connection never got the immediate close" + ); + + // Ending the tricklers releases their permits via EOF; a normal client + // must then be served (strict outer timeout). Permit release races the + // server's graceful close-drain window, so a single fixed wait can see a + // connection reset; retry within a bounded budget instead. + for task in &tricklers { + task.abort(); + } + let request = b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n"; + let mut served: Option = None; + let retry = tokio::time::Instant::now(); + while retry.elapsed() < Duration::from_secs(2) { + tokio::time::sleep(Duration::from_millis(100)).await; + let Ok(mut good) = TcpStream::connect(addr).await else { + continue; + }; + if good.write_all(request).await.is_err() { + continue; + } + let mut response = Vec::new(); + match tokio::time::timeout(Duration::from_secs(5), good.read_to_end(&mut response)).await { + // A reset mid-handshake is the drain race; retry. + Ok(Err(_)) | Err(_) => continue, + Ok(Ok(_)) => {} + } + if String::from_utf8_lossy(&response).starts_with("HTTP/1.1 200") { + served = Some(String::from_utf8_lossy(&response).into_owned()); + break; + } + } + let served = served.expect("no freed slot served a normal request within retry budget"); + assert!( + served.starts_with("HTTP/1.1 200"), + "freed slot must serve a normal request, got: {served}" + ); + server_task.abort(); +} + +#[tokio::test] +async fn deadline_driven_release_frees_gated_slot() { + // Regression (review follow-up): a semaphore permit MUST be owned for the + // whole handle_client future and released when the SERVER's total head + // deadline completes its 400/close path — not by client EOF/manual drop. + // The holder client stays connected the entire test. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let holder_budget = Duration::from_millis(500); + let config = Arc::new(head_test_config(holder_budget, None)); + let server_task = tokio::spawn(serve_listener(listener, config, 1)); + + // Fill the single permit with a holder that never sends a single byte. + let mut holder = TcpStream::connect(addr).await.unwrap(); + + // Synchronize until that permit is provably held: an over-cap probe gets + // an immediate close with zero response bytes. + let mut rejected = false; + for _ in 0..40 { + let mut probe = TcpStream::connect(addr).await.unwrap(); + let mut buf = [0_u8; 16]; + match tokio::time::timeout(Duration::from_millis(300), probe.read(&mut buf)).await { + Ok(Ok(0)) => { + rejected = true; + break; + } + // Probe landed while the holder was still being accepted; retry. + _ => drop(probe), + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!(rejected, "over-cap probe was never closed immediately"); + + // Causality phase: the holder is NOT dropped/aborted/shut down. The + // server's injected total head deadline expires on its own, completing + // the pinned 400/close path. read_to_end returns at the server's FIN; + // the holder socket itself stays OPEN. + let mut holder_response = Vec::new(); + tokio::time::timeout( + Duration::from_secs(3), + holder.read_to_end(&mut holder_response), + ) + .await + .expect("server never drove its deadline/close on the held slot") + .unwrap(); + let holder_response = String::from_utf8_lossy(&holder_response); + assert!( + holder_response.starts_with("HTTP/1.1 400"), + "deadline path must answer the holder with the pinned 400, got: {holder_response}" + ); + assert!(holder_response.contains("Cache-Control: no-store\r\n")); + assert!(holder_response.contains(r#""error":"invalid request""#)); + + // The permit frees only when the server task finishes — after the + // deadline AND the bounded (~1 s) graceful-drain that runs while the + // still-connected holder stays silent. Retry a normal client until the + // freed slot serves it; early retries may still be over-cap closed. + let started = Instant::now(); + let health = loop { + let attempt = async { + let mut good = TcpStream::connect(addr).await.ok()?; + good.write_all(b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\n\r\n") + .await + .ok()?; + let mut response = Vec::new(); + tokio::time::timeout(Duration::from_millis(800), good.read_to_end(&mut response)) + .await + .ok()? + .ok()?; + Some(String::from_utf8_lossy(&response).into_owned()) + }; + if let Some(text) = attempt.await + && text.starts_with("HTTP/1.1 200") + { + break text; + } + assert!( + started.elapsed() < Duration::from_secs(8), + "permit was never released after the server deadline + graceful drain" + ); + tokio::time::sleep(Duration::from_millis(120)).await; + }; + assert!(health.contains("\"status\":\"ok\""), "got: {health}"); + // Holder is still connected throughout everything above; cleanup only + // after all success assertions. + server_task.abort(); + drop(holder); +} diff --git a/rust/src/cli/sessions.rs b/rust/src/cli/sessions.rs index 790b71508b..42f0399f9c 100644 --- a/rust/src/cli/sessions.rs +++ b/rust/src/cli/sessions.rs @@ -147,14 +147,7 @@ pub async fn run(args: SessionsArgs) -> anyhow::Result<()> { .map(focus_session) .unwrap_or_else(|| SessionFocusResult::failed("Session was not found.")); if args.json || args.json_v2 { - println!( - "{}", - if args.pretty { - serde_json::to_string_pretty(&outcome)? - } else { - serde_json::to_string(&outcome)? - } - ); + super::print_json(&outcome, args.pretty)?; } else { print_focus_result(id, &outcome); } @@ -166,14 +159,7 @@ pub async fn run(args: SessionsArgs) -> anyhow::Result<()> { sessions: sessions_for_json(&sessions, args.json_v2), errors, }; - println!( - "{}", - if args.pretty { - serde_json::to_string_pretty(&output)? - } else { - serde_json::to_string(&output)? - } - ); + super::print_json(&output, args.pretty)?; } else if args.brief { for session in &sessions { println!("{}", render_brief(session)); diff --git a/rust/src/cli/tty_runner.rs b/rust/src/cli/tty_runner.rs index b460ccb3d6..8e917a7575 100755 --- a/rust/src/cli/tty_runner.rs +++ b/rust/src/cli/tty_runner.rs @@ -3,11 +3,6 @@ //! Executes interactive CLI commands using the platform pseudo-console. //! Provides PTY-like functionality for capturing output from interactive TUI programs. -#![allow( - dead_code, - reason = "TTY runner types reserved for future interactive session management" -)] - use super::tty_responder::{ResponderState, ScreenResponder}; use crate::process_environment::ProcessEnvironment; use regex_lite::Regex; @@ -40,171 +35,8 @@ impl TtyCommandResult { } } -/// Options for running TTY commands -#[derive(Debug, Clone)] -pub struct TtyCommandOptions { - /// Terminal rows (default: 50) - pub rows: u16, - /// Terminal columns (default: 160) - pub cols: u16, - /// Overall timeout in seconds (default: 20) - pub timeout_secs: f64, - /// Idle timeout - stop if no output for this duration (optional) - pub idle_timeout_secs: Option, - /// Working directory - pub working_directory: Option, - /// Extra arguments to pass to the command - pub extra_args: Vec, - /// Initial delay before sending script (default: 0.4s) - pub initial_delay_secs: f64, - /// Delay between script characters (default: 0s) - pub script_char_delay_secs: f64, - /// Delay between script lines (default: 0s) - pub script_line_delay_secs: f64, - /// Send enter/return every N seconds (optional) - pub send_enter_every_secs: Option, - /// Map of substrings to keys to send when detected - pub send_on_substrings: HashMap, - /// Screen-aware responder, consulted on all output since its last answer - pub screen_responder: Option, - /// Stop early when a URL is detected - pub stop_on_url: bool, - /// Stop early when any of these substrings are detected - pub stop_on_substrings: Vec, - /// Settle time after stopping (default: 0.25s) - pub settle_after_stop_secs: f64, - /// Re-send the script at these offsets (seconds since launch) while none - /// of `script_done_substrings` has appeared in the output yet. Interactive - /// CLIs such as Claude Code drop keystrokes that arrive before their input - /// widget is mounted, and that readiness delay varies between machines. - pub script_retry_delays_secs: Vec, - /// Case-insensitive markers that show the script was accepted (stops retries). - pub script_done_substrings: Vec, - /// Case-insensitive markers that show the script text was at least echoed - /// into the input widget; a retry is skipped while one is visible so the - /// same text is not typed twice into a half-processed line. - pub script_echo_substrings: Vec, - /// Shorter idle timeout used once a done marker is visible: the answer - /// is on screen, so only trailing output is awaited (optional). - pub idle_timeout_after_done_secs: Option, - /// Environment variables to set (`Debug` renders only the entry count) - pub env: ProcessEnvironment>, -} - -impl Default for TtyCommandOptions { - fn default() -> Self { - Self { - rows: 50, - cols: 160, - timeout_secs: 20.0, - idle_timeout_secs: None, - working_directory: None, - extra_args: Vec::new(), - initial_delay_secs: 0.4, - script_char_delay_secs: 0.0, - script_line_delay_secs: 0.0, - send_enter_every_secs: None, - send_on_substrings: HashMap::new(), - screen_responder: None, - stop_on_url: false, - stop_on_substrings: Vec::new(), - settle_after_stop_secs: 0.25, - script_retry_delays_secs: Vec::new(), - script_done_substrings: Vec::new(), - script_echo_substrings: Vec::new(), - idle_timeout_after_done_secs: None, - env: ProcessEnvironment::default(), - } - } -} - -impl TtyCommandOptions { - pub fn new() -> Self { - Self::default() - } - - pub fn with_timeout(mut self, secs: f64) -> Self { - self.timeout_secs = secs; - self - } - - pub fn with_idle_timeout(mut self, secs: f64) -> Self { - self.idle_timeout_secs = Some(secs); - self - } - - pub fn with_initial_delay(mut self, secs: f64) -> Self { - self.initial_delay_secs = secs; - self - } - - pub fn with_script_char_delay(mut self, secs: f64) -> Self { - self.script_char_delay_secs = secs; - self - } - - pub fn with_script_line_delay(mut self, secs: f64) -> Self { - self.script_line_delay_secs = secs; - self - } - - pub fn with_working_directory(mut self, dir: PathBuf) -> Self { - self.working_directory = Some(dir); - self - } - - pub fn with_extra_args(mut self, args: Vec) -> Self { - self.extra_args = args; - self - } - - pub fn with_stop_on_url(mut self, stop: bool) -> Self { - self.stop_on_url = stop; - self - } - - pub fn with_stop_on_substring(mut self, substring: impl Into) -> Self { - self.stop_on_substrings.push(substring.into()); - self - } - - pub fn with_script_retries( - mut self, - delays_secs: Vec, - done_substrings: Vec, - echo_substrings: Vec, - ) -> Self { - self.script_retry_delays_secs = delays_secs; - self.script_done_substrings = done_substrings - .into_iter() - .map(|marker| marker.to_lowercase()) - .collect(); - self.script_echo_substrings = echo_substrings - .into_iter() - .map(|marker| marker.to_lowercase()) - .collect(); - self - } - - pub fn with_idle_timeout_after_done(mut self, secs: f64) -> Self { - self.idle_timeout_after_done_secs = Some(secs); - self - } - - pub fn with_screen_responder(mut self, responder: ScreenResponder) -> Self { - self.screen_responder = Some(responder); - self - } - - pub fn with_send_on_substring( - mut self, - trigger: impl Into, - keys: impl Into, - ) -> Self { - self.send_on_substrings.insert(trigger.into(), keys.into()); - self - } -} +mod options; +pub use options::TtyCommandOptions; /// Errors from TTY command execution #[derive(Debug, Error)] @@ -415,22 +247,19 @@ impl TtyCommandRunner { match rx.recv_timeout(remaining.min(Duration::from_millis(50))) { Ok(chunk) => { tracing::trace!( - elapsed_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), + elapsed_ms = elapsed_ms(start), bytes = chunk.len(), "tty session: startup chunk" ); buffer.push_str(&chunk); - if chunk.contains("\x1b[6n") { - let _cursor_reply = write!(writer, "\x1b[1;1R"); - let _cursor_flushed = writer.flush(); - } + answer_cursor_query(&chunk, &mut writer); } Err(mpsc::RecvTimeoutError::Disconnected) => break, Err(mpsc::RecvTimeoutError::Timeout) => {} } } tracing::trace!( - elapsed_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), + elapsed_ms = elapsed_ms(start), buffered = buffer.len(), "tty session: sending script" ); @@ -522,22 +351,13 @@ impl TtyCommandRunner { while let Ok(chunk) = rx.try_recv() { received_output = true; tracing::trace!( - elapsed_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), + elapsed_ms = elapsed_ms(start), bytes = chunk.len(), "tty session: output chunk" ); buffer.push_str(&chunk); last_output_time = Instant::now(); - - // Some Windows ConPTY-backed shells issue an ANSI Device - // Status Report request and wait for a terminal cursor - // position response before processing scripted input. - if chunk.contains("\x1b[6n") { - // Best-effort cursor-position reply; a dead PTY ignores it. - let _cursor_reply = write!(writer, "\x1b[1;1R"); - // Best-effort flush of the cursor reply. - let _cursor_flushed = writer.flush(); - } + answer_cursor_query(&chunk, &mut writer); // Check for URLs if let Some(ref regex) = url_regex { @@ -636,7 +456,7 @@ impl TtyCommandRunner { if child.try_wait().ok().flatten().is_none() { tracing::trace!( - elapsed_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), + elapsed_ms = elapsed_ms(start), "tty session: killing child tree" ); // On Windows several CLIs (claude.exe, npm shims) are launchers @@ -649,10 +469,7 @@ impl TtyCommandRunner { let _killed = child.kill(); // Best-effort reap; the exit status is intentionally discarded. let _reaped = child.wait(); - tracing::trace!( - elapsed_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), - "tty session: child reaped" - ); + tracing::trace!(elapsed_ms = elapsed_ms(start), "tty session: child reaped"); } if buffer.is_empty() && !stopped_early { @@ -723,44 +540,18 @@ impl Default for TtyCommandRunner { } } -/// Rolling buffer for pattern matching across chunks -#[derive(Debug)] -pub struct RollingBuffer { - max_needle: usize, - tail: String, +/// Milliseconds since `start`, saturating, for trace fields. +fn elapsed_ms(start: Instant) -> u64 { + u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX) } -impl RollingBuffer { - pub fn new(max_needle: usize) -> Self { - Self { - max_needle: max_needle.max(1), - tail: String::new(), - } - } - - /// Append new data and return combined data for scanning - pub fn append(&mut self, data: &str) -> String { - if data.is_empty() { - return String::new(); - } - - let mut combined = String::with_capacity(self.tail.len() + data.len()); - combined.push_str(&self.tail); - combined.push_str(data); - - // Keep only the tail portion for next scan - if combined.len() >= self.max_needle - 1 { - let start = combined.len() - (self.max_needle - 1); - self.tail = combined[start..].to_string(); - } else { - self.tail = combined.clone(); - } - - combined - } - - pub fn reset(&mut self) { - self.tail.clear(); +/// Some Windows ConPTY-backed shells issue an ANSI Device Status Report +/// request and wait for a cursor position response before processing input; +/// answer it best-effort (a dead PTY ignores the write). +fn answer_cursor_query(chunk: &str, writer: &mut impl Write) { + if chunk.contains("\x1b[6n") { + let _cursor_reply = write!(writer, "\x1b[1;1R"); + let _cursor_flushed = writer.flush(); } } @@ -886,265 +677,4 @@ fn contains_any_marker(buffer: &str, markers: &[String]) -> bool { } #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn script_accepted_matches_markers_case_insensitively_through_ansi() { - let opts = TtyCommandOptions::new().with_script_retries( - vec![1.0], - vec!["Current session".to_string()], - vec!["/usage".to_string()], - ); - assert!(script_accepted( - "\x1b[1mCURRENT\x1b[0m session 12% used", - &opts - )); - assert!(!script_accepted("Try \"fix lint errors\"", &opts)); - assert!(!script_accepted( - "Current session", - &TtyCommandOptions::new() - )); - assert!(script_echoed("❯ /usa\x1b[0mge", &opts)); - assert!(!script_echoed("❯ Try \"fix lint errors\"", &opts)); - assert!(script_echoed("❯ /usage", &opts)); - } - - /// The trigger text is printed at startup, inside the initial delay, and - /// nothing is printed afterwards: the key must still be sent. - #[cfg(windows)] - #[test] - fn trigger_printed_inside_initial_delay_is_still_sent() { - let opts = TtyCommandOptions::new() - .with_timeout(10.0) - .with_idle_timeout(3.0) - .with_initial_delay(1.5) - .with_send_on_substring("Microsoft Windows", "echo LATE_%OS%\nexit\n"); - let result = TtyCommandRunner::new() - .run("cmd", "", opts) - .expect("pty command should run"); - assert!(result.text.contains("LATE_Windows_NT"), "{}", result.text); - } - - /// Same startup timing for the screen responder: the banner is drawn inside - /// the initial delay and nothing is printed afterwards. - #[cfg(windows)] - #[test] - fn screen_responder_fires_on_output_buffered_during_initial_delay() { - use crate::cli::tty_responder::ScreenReading; - - fn on_banner(screen: &str) -> ScreenReading { - if screen.contains("Microsoft Windows") { - ScreenReading::Answer(vec!["echo LATE_%OS%\r", "exit\r"]) - } else { - ScreenReading::Absent - } - } - let opts = TtyCommandOptions::new() - .with_timeout(10.0) - .with_idle_timeout(3.0) - .with_initial_delay(1.5) - .with_screen_responder(ScreenResponder { - read: on_banner, - after_dialog: &[], - }); - let result = TtyCommandRunner::new() - .run("cmd", "", opts) - .expect("pty command should run"); - assert!(result.text.contains("LATE_Windows_NT"), "{}", result.text); - } - - #[test] - fn substring_triggers_fire_once_on_already_buffered_output() { - let options = TtyCommandOptions::new().with_send_on_substring("Enter", "go\n"); - let mut triggered = std::collections::HashSet::new(); - let mut sent = Vec::new(); - fire_substring_triggers(&options, "Enter to confirm", &mut triggered, &mut sent); - fire_substring_triggers(&options, "Enter to confirm", &mut triggered, &mut sent); - assert_eq!(sent, b"go\r\n"); - } - - #[test] - fn test_tty_options_builder() { - let opts = TtyCommandOptions::new() - .with_timeout(30.0) - .with_idle_timeout(5.0) - .with_stop_on_url(true) - .with_stop_on_substring("error"); - - assert_eq!(opts.timeout_secs, 30.0); - assert_eq!(opts.idle_timeout_secs, Some(5.0)); - assert!(opts.stop_on_url); - assert!(opts.stop_on_substrings.contains(&"error".to_string())); - } - - #[test] - fn test_rolling_buffer() { - let mut buf = RollingBuffer::new(10); - - let result1 = buf.append("hello"); - assert_eq!(result1, "hello"); - - let result2 = buf.append(" world"); - assert!(result2.contains("hello")); - assert!(result2.contains(" world")); - } - - #[test] - fn test_tty_result_first_url() { - let result = TtyCommandResult { - text: "Visit https://example.com for more info".to_string(), - stopped_early: false, - detected_urls: vec!["https://example.com".to_string()], - }; - - assert_eq!(result.first_url(), Some("https://example.com")); - } - - #[test] - fn test_run_sends_script_through_pty() { - let runner = TtyCommandRunner::new(); - let opts = TtyCommandOptions::new() - .with_timeout(15.0) - .with_idle_timeout(6.0) - .with_initial_delay(1.0) - .with_script_line_delay(0.1); - - #[cfg(windows)] - let result = runner.run("cmd", "echo CODEXBAR_PTY_OK\nexit", opts); - #[cfg(not(windows))] - let result = runner.run("sh", "echo CODEXBAR_PTY_OK\nexit", opts); - - let result = result.expect("pty command should run"); - assert!(result.text.contains("CODEXBAR_PTY_OK"), "{}", result.text); - } - - /// A retry is input like the first attempt: a child that stays silent - /// (here `ping` with its output discarded, which never reads or echoes - /// the typed text) must not hit the idle timeout before the last retry. - #[cfg(windows)] - #[test] - fn script_retry_restarts_the_idle_window() { - let runner = TtyCommandRunner::new(); - let opts = TtyCommandOptions::new() - .with_timeout(8.0) - .with_idle_timeout(1.5) - .with_initial_delay(0.3) - .with_extra_args( - ["/d", "/c", "ping -n 10 127.0.0.1 >nul"] - .map(String::from) - .to_vec(), - ) - .with_script_retries( - vec![0.8, 1.4], - vec!["marker that never appears".to_string()], - Vec::new(), - ); - - let started = Instant::now(); - let result = runner.run("cmd", "x", opts); - let elapsed = started.elapsed(); - - let result = result.expect("silent pty command should still finish"); - assert!(result.stopped_early, "the idle timeout should end the run"); - // Without the restart the window closes 1.5 s after the first - // attempt (about 1.8 s); with it, no earlier than 1.4 s + 1.5 s. - assert!( - elapsed >= Duration::from_secs_f64(2.8), - "idle timeout fired {elapsed:?} after launch, inside the last retry's idle window" - ); - } - - /// Ending the session ends every process the PTY child started (as with a - /// `claude.cmd` launcher and its `node` child), not only the child. The - /// grandchild runs on its own hidden console, so closing the pseudoconsole - /// does not end it; only the session's job does. It holds an exclusive - /// handle on a file, so its exit is observed through that handle instead - /// of a PID that could be reused. - #[cfg(windows)] - #[test] - fn session_end_terminates_processes_started_by_the_child() { - use base64::Engine as _; - use std::os::windows::fs::OpenOptionsExt as _; - - fn encoded(command: &str) -> String { - let utf16: Vec = command - .encode_utf16() - .flat_map(|unit| unit.to_le_bytes()) - .collect(); - base64::engine::general_purpose::STANDARD.encode(utf16) - } - fn quoted(path: &std::path::Path) -> String { - format!("'{}'", path.display().to_string().replace('\'', "''")) - } - - let dir = tempfile::tempdir().expect("temp dir"); - let held = dir.path().join("held-by-grandchild.lock"); - let ready = dir.path().join("grandchild-ready"); - let grandchild = format!( - "$f = [IO.File]::Open({held}, 'OpenOrCreate', 'ReadWrite', 'None'); \ - [IO.File]::WriteAllText({ready}, 'ready'); Start-Sleep -Seconds 30", - held = quoted(&held), - ready = quoted(&ready), - ); - let child = format!( - "$psi = New-Object System.Diagnostics.ProcessStartInfo -ArgumentList \ - 'powershell.exe', '-NoLogo -NoProfile -NonInteractive -EncodedCommand {command}'; \ - $psi.UseShellExecute = $false; $psi.CreateNoWindow = $true; \ - [void][System.Diagnostics.Process]::Start($psi); \ - while (-not (Test-Path -LiteralPath {ready})) {{ Start-Sleep -Milliseconds 50 }}; \ - Write-Output 'grandchild-holds-lock'; Start-Sleep -Seconds 30", - command = encoded(&grandchild), - ready = quoted(&ready), - ); - let opts = TtyCommandOptions::new() - .with_timeout(30.0) - .with_initial_delay(0.2) - .with_stop_on_substring("grandchild-holds-lock") - .with_extra_args( - [ - "-NoLogo", - "-NoProfile", - "-NonInteractive", - "-EncodedCommand", - &encoded(&child), - ] - .map(String::from) - .to_vec(), - ); - - let result = TtyCommandRunner::new() - .run("powershell", "", opts) - .expect("pty command should run"); - assert!( - result.text.contains("grandchild-holds-lock"), - "grandchild never started: {}", - result.text - ); - - let deadline = Instant::now() + Duration::from_secs(5); - loop { - let exclusive = std::fs::OpenOptions::new() - .read(true) - .write(true) - .share_mode(0) - .open(&held); - if exclusive.is_ok() { - break; - } - assert!( - Instant::now() < deadline, - "the grandchild still holds its file after the session ended" - ); - std::thread::sleep(Duration::from_millis(50)); - } - } - - #[test] - fn test_enriched_path() { - let path = TtyCommandRunner::enriched_path(); - assert!(!path.is_empty()); - // Should contain path separator or at least a non-empty path string. - assert!(path.contains(';') || !path.is_empty()); - } -} +mod tests; diff --git a/rust/src/cli/tty_runner/options.rs b/rust/src/cli/tty_runner/options.rs new file mode 100644 index 0000000000..440a1d10fd --- /dev/null +++ b/rust/src/cli/tty_runner/options.rs @@ -0,0 +1,169 @@ +//! `TtyCommandOptions`: terminal size, timeouts, scripted input and stop markers. + +use super::*; + +/// Options for running TTY commands +#[derive(Debug, Clone)] +pub struct TtyCommandOptions { + /// Terminal rows (default: 50) + pub rows: u16, + /// Terminal columns (default: 160) + pub cols: u16, + /// Overall timeout in seconds (default: 20) + pub timeout_secs: f64, + /// Idle timeout - stop if no output for this duration (optional) + pub idle_timeout_secs: Option, + /// Working directory + pub working_directory: Option, + /// Extra arguments to pass to the command + pub extra_args: Vec, + /// Initial delay before sending script (default: 0.4s) + pub initial_delay_secs: f64, + /// Delay between script characters (default: 0s) + pub script_char_delay_secs: f64, + /// Delay between script lines (default: 0s) + pub script_line_delay_secs: f64, + /// Send enter/return every N seconds (optional) + pub send_enter_every_secs: Option, + /// Map of substrings to keys to send when detected + pub send_on_substrings: HashMap, + /// Screen-aware responder, consulted on all output since its last answer + pub screen_responder: Option, + /// Stop early when a URL is detected + pub stop_on_url: bool, + /// Stop early when any of these substrings are detected + pub stop_on_substrings: Vec, + /// Settle time after stopping (default: 0.25s) + pub settle_after_stop_secs: f64, + /// Re-send the script at these offsets (seconds since launch) while none + /// of `script_done_substrings` has appeared in the output yet. Interactive + /// CLIs such as Claude Code drop keystrokes that arrive before their input + /// widget is mounted, and that readiness delay varies between machines. + pub script_retry_delays_secs: Vec, + /// Case-insensitive markers that show the script was accepted (stops retries). + pub script_done_substrings: Vec, + /// Case-insensitive markers that show the script text was at least echoed + /// into the input widget; a retry is skipped while one is visible so the + /// same text is not typed twice into a half-processed line. + pub script_echo_substrings: Vec, + /// Shorter idle timeout used once a done marker is visible: the answer + /// is on screen, so only trailing output is awaited (optional). + pub idle_timeout_after_done_secs: Option, + /// Environment variables to set (`Debug` renders only the entry count) + pub env: ProcessEnvironment>, +} + +impl Default for TtyCommandOptions { + fn default() -> Self { + Self { + rows: 50, + cols: 160, + timeout_secs: 20.0, + idle_timeout_secs: None, + working_directory: None, + extra_args: Vec::new(), + initial_delay_secs: 0.4, + script_char_delay_secs: 0.0, + script_line_delay_secs: 0.0, + send_enter_every_secs: None, + send_on_substrings: HashMap::new(), + screen_responder: None, + stop_on_url: false, + stop_on_substrings: Vec::new(), + settle_after_stop_secs: 0.25, + script_retry_delays_secs: Vec::new(), + script_done_substrings: Vec::new(), + script_echo_substrings: Vec::new(), + idle_timeout_after_done_secs: None, + env: ProcessEnvironment::default(), + } + } +} + +impl TtyCommandOptions { + pub fn new() -> Self { + Self::default() + } + + pub fn with_timeout(mut self, secs: f64) -> Self { + self.timeout_secs = secs; + self + } + + pub fn with_idle_timeout(mut self, secs: f64) -> Self { + self.idle_timeout_secs = Some(secs); + self + } + + pub fn with_initial_delay(mut self, secs: f64) -> Self { + self.initial_delay_secs = secs; + self + } + + pub fn with_script_char_delay(mut self, secs: f64) -> Self { + self.script_char_delay_secs = secs; + self + } + + pub fn with_script_line_delay(mut self, secs: f64) -> Self { + self.script_line_delay_secs = secs; + self + } + + pub fn with_working_directory(mut self, dir: PathBuf) -> Self { + self.working_directory = Some(dir); + self + } + + pub fn with_extra_args(mut self, args: Vec) -> Self { + self.extra_args = args; + self + } + + pub fn with_stop_on_url(mut self, stop: bool) -> Self { + self.stop_on_url = stop; + self + } + + pub fn with_stop_on_substring(mut self, substring: impl Into) -> Self { + self.stop_on_substrings.push(substring.into()); + self + } + + pub fn with_script_retries( + mut self, + delays_secs: Vec, + done_substrings: Vec, + echo_substrings: Vec, + ) -> Self { + self.script_retry_delays_secs = delays_secs; + self.script_done_substrings = done_substrings + .into_iter() + .map(|marker| marker.to_lowercase()) + .collect(); + self.script_echo_substrings = echo_substrings + .into_iter() + .map(|marker| marker.to_lowercase()) + .collect(); + self + } + + pub fn with_idle_timeout_after_done(mut self, secs: f64) -> Self { + self.idle_timeout_after_done_secs = Some(secs); + self + } + + pub fn with_screen_responder(mut self, responder: ScreenResponder) -> Self { + self.screen_responder = Some(responder); + self + } + + pub fn with_send_on_substring( + mut self, + trigger: impl Into, + keys: impl Into, + ) -> Self { + self.send_on_substrings.insert(trigger.into(), keys.into()); + self + } +} diff --git a/rust/src/cli/tty_runner/tests.rs b/rust/src/cli/tty_runner/tests.rs new file mode 100644 index 0000000000..ad98945699 --- /dev/null +++ b/rust/src/cli/tty_runner/tests.rs @@ -0,0 +1,248 @@ +use super::*; + +#[test] +fn script_accepted_matches_markers_case_insensitively_through_ansi() { + let opts = TtyCommandOptions::new().with_script_retries( + vec![1.0], + vec!["Current session".to_string()], + vec!["/usage".to_string()], + ); + assert!(script_accepted( + "\x1b[1mCURRENT\x1b[0m session 12% used", + &opts + )); + assert!(!script_accepted("Try \"fix lint errors\"", &opts)); + assert!(!script_accepted( + "Current session", + &TtyCommandOptions::new() + )); + assert!(script_echoed("❯ /usa\x1b[0mge", &opts)); + assert!(!script_echoed("❯ Try \"fix lint errors\"", &opts)); + assert!(script_echoed("❯ /usage", &opts)); +} + +/// The trigger text is printed at startup, inside the initial delay, and +/// nothing is printed afterwards: the key must still be sent. +#[cfg(windows)] +#[test] +fn trigger_printed_inside_initial_delay_is_still_sent() { + let opts = TtyCommandOptions::new() + .with_timeout(10.0) + .with_idle_timeout(3.0) + .with_initial_delay(1.5) + .with_send_on_substring("Microsoft Windows", "echo LATE_%OS%\nexit\n"); + let result = TtyCommandRunner::new() + .run("cmd", "", opts) + .expect("pty command should run"); + assert!(result.text.contains("LATE_Windows_NT"), "{}", result.text); +} + +/// Same startup timing for the screen responder: the banner is drawn inside +/// the initial delay and nothing is printed afterwards. +#[cfg(windows)] +#[test] +fn screen_responder_fires_on_output_buffered_during_initial_delay() { + use crate::cli::tty_responder::ScreenReading; + + fn on_banner(screen: &str) -> ScreenReading { + if screen.contains("Microsoft Windows") { + ScreenReading::Answer(vec!["echo LATE_%OS%\r", "exit\r"]) + } else { + ScreenReading::Absent + } + } + let opts = TtyCommandOptions::new() + .with_timeout(10.0) + .with_idle_timeout(3.0) + .with_initial_delay(1.5) + .with_screen_responder(ScreenResponder { + read: on_banner, + after_dialog: &[], + }); + let result = TtyCommandRunner::new() + .run("cmd", "", opts) + .expect("pty command should run"); + assert!(result.text.contains("LATE_Windows_NT"), "{}", result.text); +} + +#[test] +fn substring_triggers_fire_once_on_already_buffered_output() { + let options = TtyCommandOptions::new().with_send_on_substring("Enter", "go\n"); + let mut triggered = std::collections::HashSet::new(); + let mut sent = Vec::new(); + fire_substring_triggers(&options, "Enter to confirm", &mut triggered, &mut sent); + fire_substring_triggers(&options, "Enter to confirm", &mut triggered, &mut sent); + assert_eq!(sent, b"go\r\n"); +} + +#[test] +fn test_tty_options_builder() { + let opts = TtyCommandOptions::new() + .with_timeout(30.0) + .with_idle_timeout(5.0) + .with_stop_on_url(true) + .with_stop_on_substring("error"); + + assert_eq!(opts.timeout_secs, 30.0); + assert_eq!(opts.idle_timeout_secs, Some(5.0)); + assert!(opts.stop_on_url); + assert!(opts.stop_on_substrings.contains(&"error".to_string())); +} + +#[test] +fn test_tty_result_first_url() { + let result = TtyCommandResult { + text: "Visit https://example.com for more info".to_string(), + stopped_early: false, + detected_urls: vec!["https://example.com".to_string()], + }; + + assert_eq!(result.first_url(), Some("https://example.com")); +} + +#[test] +fn test_run_sends_script_through_pty() { + let runner = TtyCommandRunner::new(); + let opts = TtyCommandOptions::new() + .with_timeout(15.0) + .with_idle_timeout(6.0) + .with_initial_delay(1.0) + .with_script_line_delay(0.1); + + #[cfg(windows)] + let result = runner.run("cmd", "echo CODEXBAR_PTY_OK\nexit", opts); + #[cfg(not(windows))] + let result = runner.run("sh", "echo CODEXBAR_PTY_OK\nexit", opts); + + let result = result.expect("pty command should run"); + assert!(result.text.contains("CODEXBAR_PTY_OK"), "{}", result.text); +} + +/// A retry is input like the first attempt: a child that stays silent +/// (here `ping` with its output discarded, which never reads or echoes +/// the typed text) must not hit the idle timeout before the last retry. +#[cfg(windows)] +#[test] +fn script_retry_restarts_the_idle_window() { + let runner = TtyCommandRunner::new(); + let opts = TtyCommandOptions::new() + .with_timeout(8.0) + .with_idle_timeout(1.5) + .with_initial_delay(0.3) + .with_extra_args( + ["/d", "/c", "ping -n 10 127.0.0.1 >nul"] + .map(String::from) + .to_vec(), + ) + .with_script_retries( + vec![0.8, 1.4], + vec!["marker that never appears".to_string()], + Vec::new(), + ); + + let started = Instant::now(); + let result = runner.run("cmd", "x", opts); + let elapsed = started.elapsed(); + + let result = result.expect("silent pty command should still finish"); + assert!(result.stopped_early, "the idle timeout should end the run"); + // Without the restart the window closes 1.5 s after the first + // attempt (about 1.8 s); with it, no earlier than 1.4 s + 1.5 s. + assert!( + elapsed >= Duration::from_secs_f64(2.8), + "idle timeout fired {elapsed:?} after launch, inside the last retry's idle window" + ); +} + +/// Ending the session ends every process the PTY child started (as with a +/// `claude.cmd` launcher and its `node` child), not only the child. The +/// grandchild runs on its own hidden console, so closing the pseudoconsole +/// does not end it; only the session's job does. It holds an exclusive +/// handle on a file, so its exit is observed through that handle instead +/// of a PID that could be reused. +#[cfg(windows)] +#[test] +fn session_end_terminates_processes_started_by_the_child() { + use base64::Engine as _; + use std::os::windows::fs::OpenOptionsExt as _; + + fn encoded(command: &str) -> String { + let utf16: Vec = command + .encode_utf16() + .flat_map(|unit| unit.to_le_bytes()) + .collect(); + base64::engine::general_purpose::STANDARD.encode(utf16) + } + fn quoted(path: &std::path::Path) -> String { + format!("'{}'", path.display().to_string().replace('\'', "''")) + } + + let dir = tempfile::tempdir().expect("temp dir"); + let held = dir.path().join("held-by-grandchild.lock"); + let ready = dir.path().join("grandchild-ready"); + let grandchild = format!( + "$f = [IO.File]::Open({held}, 'OpenOrCreate', 'ReadWrite', 'None'); \ + [IO.File]::WriteAllText({ready}, 'ready'); Start-Sleep -Seconds 30", + held = quoted(&held), + ready = quoted(&ready), + ); + let child = format!( + "$psi = New-Object System.Diagnostics.ProcessStartInfo -ArgumentList \ + 'powershell.exe', '-NoLogo -NoProfile -NonInteractive -EncodedCommand {command}'; \ + $psi.UseShellExecute = $false; $psi.CreateNoWindow = $true; \ + [void][System.Diagnostics.Process]::Start($psi); \ + while (-not (Test-Path -LiteralPath {ready})) {{ Start-Sleep -Milliseconds 50 }}; \ + Write-Output 'grandchild-holds-lock'; Start-Sleep -Seconds 30", + command = encoded(&grandchild), + ready = quoted(&ready), + ); + let opts = TtyCommandOptions::new() + .with_timeout(30.0) + .with_initial_delay(0.2) + .with_stop_on_substring("grandchild-holds-lock") + .with_extra_args( + [ + "-NoLogo", + "-NoProfile", + "-NonInteractive", + "-EncodedCommand", + &encoded(&child), + ] + .map(String::from) + .to_vec(), + ); + + let result = TtyCommandRunner::new() + .run("powershell", "", opts) + .expect("pty command should run"); + assert!( + result.text.contains("grandchild-holds-lock"), + "grandchild never started: {}", + result.text + ); + + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let exclusive = std::fs::OpenOptions::new() + .read(true) + .write(true) + .share_mode(0) + .open(&held); + if exclusive.is_ok() { + break; + } + assert!( + Instant::now() < deadline, + "the grandchild still holds its file after the session ended" + ); + std::thread::sleep(Duration::from_millis(50)); + } +} + +#[test] +fn test_enriched_path() { + let path = TtyCommandRunner::enriched_path(); + assert!(!path.is_empty()); + // Should contain path separator or at least a non-empty path string. + assert!(path.contains(';') || !path.is_empty()); +} diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 25bdd89d47..8e8863da0c 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -1,9 +1,8 @@ //! Usage command implementation use clap::Args; -use serde::Serialize; -use crate::core::{FetchContext, ProviderFetchResult, ProviderId, SourceMode}; +use crate::core::{FetchContext, ProviderId, SourceMode}; use crate::settings::Settings; mod claude_swap; @@ -11,11 +10,14 @@ mod fetch_helpers; mod render; use fetch_helpers::{fetch_provider_json_output, fetch_provider_text_output}; +#[cfg(test)] +use render::render_text; +use render::status_json; pub(super) use render::{ append_status_line, format_percent, render_status_indicator, render_text_error, }; use render::{is_terminal, print_usage_output}; -pub use render::{render_brief_text, render_text, render_text_with_status}; +pub use render::{render_brief_text, render_text_with_status}; pub(super) enum UsageOutput { Text(Vec), @@ -192,24 +194,6 @@ impl ProviderSelection { } } -/// JSON output payload -#[derive(Debug, Serialize)] -pub struct ProviderPayload { - pub provider: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub version: Option, - pub source: String, - #[serde(flatten)] - pub result: ProviderFetchResult, -} - -/// Error payload for JSON output -#[derive(Debug, Serialize)] -struct ErrorPayload { - provider: String, - error: String, -} - /// Run the usage command pub async fn run(args: UsageArgs) -> anyhow::Result<()> { let command = UsageCommand::from_args(args)?; @@ -301,24 +285,13 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch source_mode, include_credits: !args.no_credits, web_timeout: args.web_timeout, - verbose: false, - manual_cookie_header: None, - manual_cookie_missing: false, - api_key: None, - token_account_kind: None, - token_account_isolated: false, - workspace_id: None, - seat_credit_entitlement: None, - api_region: None, - gateway_url: None, - auto_prefer_web: false, - browser_cookie_import: false, // `codexbar usage` is a foreground read: optional enrichment (e.g. the // OpenCode Go Zen balance) is worth its full bounded wait (#2583). requires_optional_usage_completeness: true, // Per-provider opt-ins live in settings, which this shared context // does not load; the optional breakdowns stay off here. optional_details_enabled: false, + ..FetchContext::default() } } diff --git a/rust/src/cli/usage/claude_swap.rs b/rust/src/cli/usage/claude_swap.rs index 4ce16fd822..aa8599ec95 100644 --- a/rust/src/cli/usage/claude_swap.rs +++ b/rust/src/cli/usage/claude_swap.rs @@ -3,6 +3,7 @@ use serde::Serialize; use super::{ UsageCommand, UsageOutput, UsageOutputFormat, append_status_line, fetch_provider_json_output, fetch_provider_text_output, format_percent, render_status_indicator, render_text_error, + status_json, }; use crate::core::ProviderId; use crate::providers::claude::claude_swap::{ @@ -97,23 +98,33 @@ pub(super) fn claude_swap_json_payload( "account": ClaudeSwapCliAccount::from(account), }); if let Some(status) = status { - payload["status"] = serde_json::json!({ - "level": format!("{:?}", status.level).to_lowercase(), - "description": status.description, - }); + payload["status"] = status_json(status); } payload } fn claude_swap_windows(account: &ClaudeSwapAccount) -> Vec { + window_labels( + account.five_hour.as_ref(), + account.seven_day.as_ref(), + &account.scoped, + ) +} + +/// `Session N% | Weekly N% | N%` parts, skipping absent windows. +fn window_labels( + five_hour: Option<&ClaudeSwapUsageWindowDto>, + seven_day: Option<&ClaudeSwapUsageWindowDto>, + scoped: &[ClaudeSwapScopedWindowDto], +) -> Vec { let mut windows = Vec::new(); - if let Some(window) = &account.five_hour { + if let Some(window) = five_hour { windows.push(format!("Session {}", format_percent(window.used_percent))); } - if let Some(window) = &account.seven_day { + if let Some(window) = seven_day { windows.push(format!("Weekly {}", format_percent(window.used_percent))); } - for window in &account.scoped { + for window in scoped { windows.push(format!( "{} {}", window.name, @@ -137,24 +148,12 @@ fn claude_swap_spend_line( }) } -fn claude_swap_historical_windows( - historical: &crate::providers::claude::claude_swap::ClaudeSwapHistoricalUsageDto, -) -> Vec { - let mut windows = Vec::new(); - if let Some(window) = &historical.five_hour { - windows.push(format!("Session {}", format_percent(window.used_percent))); - } - if let Some(window) = &historical.seven_day { - windows.push(format!("Weekly {}", format_percent(window.used_percent))); - } - for window in &historical.scoped { - windows.push(format!( - "{} {}", - window.name, - format_percent(window.used_percent) - )); - } - windows +fn claude_swap_historical_windows(historical: &ClaudeSwapHistoricalUsageDto) -> Vec { + window_labels( + historical.five_hour.as_ref(), + historical.seven_day.as_ref(), + &historical.scoped, + ) } pub(super) fn render_claude_swap_text( @@ -289,10 +288,7 @@ fn claude_swap_error_payload(error: &str, status: Option<&StatusInfo>) -> serde_ "errorKind": crate::cli::error_kind::ERROR_KIND_UNKNOWN, }); if let Some(status) = status { - payload["status"] = serde_json::json!({ - "level": format!("{:?}", status.level).to_lowercase(), - "description": status.description, - }); + payload["status"] = status_json(status); } payload } diff --git a/rust/src/cli/usage/render.rs b/rust/src/cli/usage/render.rs index de41a08697..d4948e3c21 100644 --- a/rust/src/cli/usage/render.rs +++ b/rust/src/cli/usage/render.rs @@ -96,15 +96,20 @@ pub fn render_json_result( } if let Some(s) = status { - json_result["status"] = serde_json::json!({ - "level": format!("{:?}", s.level).to_lowercase(), - "description": s.description, - }); + json_result["status"] = status_json(s); } json_result } +/// `{level, description}` status object shared by every usage JSON row. +pub(super) fn status_json(status: &StatusInfo) -> serde_json::Value { + serde_json::json!({ + "level": format!("{:?}", status.level).to_lowercase(), + "description": status.description, + }) +} + /// Serialize a [`UsagePace`] into a compact JSON object for the `--json` output. fn pace_json(pace: UsagePace) -> serde_json::Value { serde_json::json!({ @@ -121,12 +126,7 @@ pub(super) fn print_usage_output(output: UsageOutput) -> anyhow::Result<()> { println!("{}", sections.join("\n\n")); } UsageOutput::Json { results, pretty } => { - let output = if pretty { - serde_json::to_string_pretty(&results)? - } else { - serde_json::to_string(&results)? - }; - println!("{}", output); + crate::cli::print_json(&results, pretty)?; } UsageOutput::Toon(results) => { println!( @@ -526,7 +526,12 @@ fn format_token_count(tokens: u64) -> String { } /// Render usage as text (backwards compatible version) -pub fn render_text(provider: ProviderId, result: &ProviderFetchResult, use_color: bool) -> String { +#[cfg(test)] +pub(super) fn render_text( + provider: ProviderId, + result: &ProviderFetchResult, + use_color: bool, +) -> String { render_text_with_status(provider, result, None, use_color) } diff --git a/rust/src/cli/usage_tests.rs b/rust/src/cli/usage_tests.rs index 72ec65099d..098a2b3e30 100644 --- a/rust/src/cli/usage_tests.rs +++ b/rust/src/cli/usage_tests.rs @@ -3,8 +3,8 @@ use super::*; use crate::core::{ CostSnapshot, FetchContext, NamedRateWindow, ProviderAccountData, ProviderDisplayDetail, - ProviderId, ProviderInventoryItem, RateWindow, SourceMode, TokenAccount, TokenAccountKind, - TokenAccountSupport, UsageSnapshot, + ProviderFetchResult, ProviderId, ProviderInventoryItem, RateWindow, SourceMode, TokenAccount, + TokenAccountKind, TokenAccountSupport, UsageSnapshot, }; use crate::providers::claude::claude_swap::ClaudeSwapAccount; use crate::status::{ProviderStatus as StatusInfo, StatusLevel}; diff --git a/rust/src/cli/workspaces.rs b/rust/src/cli/workspaces.rs index e270487f3c..dc702b3036 100644 --- a/rust/src/cli/workspaces.rs +++ b/rust/src/cli/workspaces.rs @@ -69,11 +69,7 @@ pub async fn run(args: WorkspacesArgs) -> anyhow::Result<()> { })?; if args.json { - if args.pretty { - println!("{}", serde_json::to_string_pretty(&snapshot)?); - } else { - println!("{}", serde_json::to_string(&snapshot)?); - } + super::print_json(&snapshot, args.pretty)?; return Ok(()); } diff --git a/rust/src/codex_accounts/account_manager.rs b/rust/src/codex_accounts/account_manager.rs index 9079eb3e94..c5e9d01236 100644 --- a/rust/src/codex_accounts/account_manager.rs +++ b/rust/src/codex_accounts/account_manager.rs @@ -181,15 +181,6 @@ impl CodexAccountManager { )) } - /// Identity of the currently active (ambient) account, if any. - pub fn load_active_identity(&self) -> Option { - let auth_path = ambient_codex_home().join("auth.json"); - if !auth_path.exists() { - return None; - } - load_identity(&ambient_codex_home()).ok() - } - /// Switch the ambient identity to `target`, materializing the previous /// ambient account as managed and preserving the desktop session. pub fn switch_active_account( diff --git a/rust/src/codex_accounts/api.rs b/rust/src/codex_accounts/api.rs index f19835cffa..9feabd60b3 100644 --- a/rust/src/codex_accounts/api.rs +++ b/rust/src/codex_accounts/api.rs @@ -691,328 +691,4 @@ fn credits_equivalent( } #[cfg(test)] -mod tests { - use super::*; - use base64::Engine; - - #[test] - fn credential_error_states_exclude_permission_and_transport_failures() { - assert_eq!( - CodexApiError::AuthenticationRequired("missing credentials".to_string()).state_kind(), - ProviderStateKind::NeedsAuthentication - ); - assert_eq!( - CodexApiError::SessionExpired("unauthorized".to_string()).state_kind(), - ProviderStateKind::ExpiredSession - ); - assert_eq!( - CodexApiError::PermissionDenied("forbidden".to_string()).state_kind(), - ProviderStateKind::Unknown - ); - assert_eq!( - CodexApiError::Network("offline".to_string()).state_kind(), - ProviderStateKind::Unknown - ); - } - - #[tokio::test] - async fn active_fetches_use_and_sync_ambient_credentials_even_when_usage_fails() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::TcpListener; - use tokio::time::{Duration, timeout}; - - for (target_id, managed_newer, expected_token) in [ - ("active", false, "ambient-token"), - ("active", true, "managed-token"), - ("other", true, "managed-token"), - ] { - let ambient = tempfile::tempdir().unwrap(); - let managed = tempfile::tempdir().unwrap(); - let credentials = |account: &str, token: &str| AuthCredentials { - access_token: token.into(), - refresh_token: format!("refresh-{token}"), - id_token: None, - account_id: Some(account.into()), - last_refresh: Some(Utc::now()), - }; - save_credentials(ambient.path(), &credentials("active", "ambient-token")).unwrap(); - save_credentials(managed.path(), &credentials(target_id, "managed-token")).unwrap(); - let now = Utc::now(); - for (home, refreshed) in [ - (ambient.path(), now - chrono::TimeDelta::hours(1)), - ( - managed.path(), - if managed_newer { - now - } else { - now - chrono::TimeDelta::hours(2) - }, - ), - ] { - let path = home.join("auth.json"); - let mut json: serde_json::Value = - serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); - json["last_refresh"] = serde_json::json!(refreshed.to_rfc3339()); - std::fs::write(path, json.to_string()).unwrap(); - } - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let config = format!( - "chatgpt_base_url = \"http://{}\"\n", - listener.local_addr().unwrap() - ); - for home in [ambient.path(), managed.path()] { - std::fs::write(home.join("config.toml"), &config).unwrap(); - } - super::super::file_locations::with_ambient_codex_home(ambient.path().to_owned()); - let api = CodexAccountApi { - client: reqwest::Client::builder().no_proxy().build().unwrap(), - }; - let server = async { - let (mut stream, _) = timeout(Duration::from_secs(5), listener.accept()) - .await - .unwrap() - .unwrap(); - let mut headers = Vec::new(); - while !headers.ends_with(b"\r\n\r\n") { - headers.push( - timeout(Duration::from_secs(5), stream.read_u8()) - .await - .unwrap() - .unwrap(), - ); - } - stream.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}").await.unwrap(); - String::from_utf8(headers).unwrap().to_lowercase() - }; - let (result, headers) = - tokio::join!(api.fetch_snapshot(managed.path(), None, false), server); - super::super::file_locations::clear_ambient_codex_home_override(); - assert!(result.is_err()); - assert!(headers.contains(&format!("authorization: bearer {expected_token}"))); - let saved = load_credentials(managed.path()).unwrap(); - assert_eq!(saved.access_token, expected_token); - assert_eq!(saved.refresh_token, format!("refresh-{expected_token}")); - assert_eq!( - load_credentials(ambient.path()).unwrap().access_token, - if target_id == "active" { - expected_token - } else { - "ambient-token" - } - ); - } - } - - #[tokio::test] - async fn overlapping_fetches_reload_credentials_and_keep_other_homes_parallel() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - use tokio::net::{TcpListener, TcpStream}; - use tokio::time::{Duration, timeout}; - - async fn request(listener: &TcpListener) -> (TcpStream, String) { - let (mut stream, _) = timeout(Duration::from_secs(5), listener.accept()) - .await - .unwrap() - .unwrap(); - let mut headers = Vec::new(); - while !headers.ends_with(b"\r\n\r\n") { - headers.push( - timeout(Duration::from_secs(5), stream.read_u8()) - .await - .unwrap() - .unwrap(), - ); - } - (stream, String::from_utf8(headers).unwrap().to_lowercase()) - } - async fn respond(mut stream: TcpStream) { - stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}").await.unwrap(); - } - fn configure(home: &Path, address: std::net::SocketAddr, token: &str) { - std::fs::write( - home.join("config.toml"), - format!("chatgpt_base_url = \"http://{address}\"\n"), - ) - .unwrap(); - std::fs::write( - home.join("auth.json"), - serde_json::json!({"OPENAI_API_KEY":token}).to_string(), - ) - .unwrap(); - } - fn fetch( - home: std::path::PathBuf, - ) -> tokio::task::JoinHandle> { - tokio::spawn(async move { - let api = CodexAccountApi { - client: reqwest::Client::builder().no_proxy().build().unwrap(), - }; - // Exercise per-home concurrency independently of other tests - // that intentionally take the global account-switch write lock. - super::super::fetch_coordination::fetch_home_snapshot( - &api, &home, None, None, false, None, - ) - .await - }) - } - - let first_home = tempfile::tempdir().unwrap(); - let other_home = tempfile::tempdir().unwrap(); - let first_server = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let other_server = TcpListener::bind("127.0.0.1:0").await.unwrap(); - configure( - first_home.path(), - first_server.local_addr().unwrap(), - "old-token", - ); - configure( - other_home.path(), - other_server.local_addr().unwrap(), - "other-token", - ); - let first = fetch(first_home.path().to_owned()); - let (first_stream, headers) = request(&first_server).await; - assert!(headers.contains("authorization: bearer old-token")); - // A lexical alias of the same auth path must share the first lane. - let second = fetch(first_home.path().join(".")); - let other = fetch(other_home.path().to_owned()); - let (other_stream, headers) = request(&other_server).await; - assert!(headers.contains("authorization: bearer other-token")); - respond(other_stream).await; - timeout(Duration::from_secs(5), other) - .await - .unwrap() - .unwrap() - .unwrap(); - assert!( - timeout(Duration::from_millis(100), first_server.accept()) - .await - .is_err() - ); - - // Model a rotated token being persisted by the first in-flight fetch. - configure( - first_home.path(), - first_server.local_addr().unwrap(), - "rotated-token", - ); - respond(first_stream).await; - timeout(Duration::from_secs(5), first) - .await - .unwrap() - .unwrap() - .unwrap(); - let (second_stream, headers) = request(&first_server).await; - assert!(headers.contains("authorization: bearer rotated-token")); - respond(second_stream).await; - timeout(Duration::from_secs(5), second) - .await - .unwrap() - .unwrap() - .unwrap(); - } - - #[test] - fn parse_credentials_accepts_api_key() { - let creds = parse_credentials_json(r#"{"OPENAI_API_KEY":"sk-test"}"#).unwrap(); - assert_eq!(creds.access_token, "sk-test"); - assert_eq!(creds.account_id, None); - } - - #[test] - fn parse_credentials_accepts_tokens() { - let creds = parse_credentials_json( - r#"{"tokens":{"access_token":"at","refresh_token":"rt","account_id":"42"},"last_refresh":"2026-01-01T00:00:00Z"}"#, - ) - .unwrap(); - assert_eq!(creds.access_token, "at"); - assert_eq!(creds.refresh_token, "rt"); - assert_eq!(creds.account_id.as_deref(), Some("42")); - } - - #[test] - fn parse_credentials_missing_tokens_errors() { - assert!(parse_credentials_json(r#"{"foo":1}"#).is_err()); - } - - #[test] - fn jwt_payload_decodes() { - let payload = - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(br#"{"email":"a@b.c"}"#); - let token = format!("eyJhbGciOiJub25lIn0.{payload}."); - let parsed = jwt_payload(&token).unwrap(); - assert_eq!(parsed.get("email").and_then(|v| v.as_str()), Some("a@b.c")); - } - - #[test] - fn normalize_window_roles_orders_session_first() { - let weekly = UsageWindowSnapshot::new(10.0, None, 604_800); - let session = UsageWindowSnapshot::new(10.0, None, 18_000); - let (p, s) = normalize_window_roles(Some(weekly), Some(session)); - assert_eq!(p.unwrap().limit_window_seconds, 18_000); - assert_eq!(s.unwrap().limit_window_seconds, 604_800); - } - - #[test] - fn limit_reached_forces_100() { - let payload = make_rate_limit(); - let (p, s) = make_normalized_windows(Some(&payload)); - assert_eq!(p.as_ref().unwrap().used_percent, 100.0); - assert_eq!(s.as_ref().unwrap().used_percent, 100.0); - } - - fn make_rate_limit() -> serde_json::Map { - serde_json::from_str( - r#"{"allowed":true,"limit_reached":true,"primary_window":{"used_percent":40,"reset_at":0,"limit_window_seconds":18000},"secondary_window":{"used_percent":20,"reset_at":0,"limit_window_seconds":604800}}"#, - ) - .unwrap() - } - - #[test] - fn resolve_usage_url_default() { - let dir = tempfile::tempdir().unwrap(); - let url = resolve_usage_url(dir.path()); - assert_eq!(url, "https://chatgpt.com/backend-api/wham/usage"); - } - - #[test] - fn parse_chatgpt_base_url_parses_quoted() { - let url = parse_chatgpt_base_url( - "# comment\nchatgpt_base_url = \"https://example.com/backend-api\"\n", - ) - .unwrap(); - assert_eq!(url, "https://example.com/backend-api"); - } - - #[test] - fn equivalent_snapshots_match() { - let mk = || AccountUsageSnapshot { - email: Some("a@b.c".to_string()), - provider_account_id: Some("x".to_string()), - plan: Some("pro".to_string()), - allowed: Some(true), - limit_reached: None, - primary_window: Some(UsageWindowSnapshot::new(12.0, Some(Utc::now()), 18_000)), - secondary_window: None, - credits: None, - cost: None, - subscription: None, - updated_at: Utc::now(), - }; - assert!(is_equivalent(&mk(), &mk())); - let mut different = mk(); - different.plan = Some("plus".to_string()); - assert!(!is_equivalent(&mk(), &different)); - } - - #[test] - fn account_id_from_id_token_reads_auth() { - let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .encode(br#"{"https://api.openai.com/auth":{"chatgpt_account_id":"acct-99"}}"#); - let token = format!("h.{payload}.s"); - assert_eq!( - account_id_from_id_token(Some(&token)).as_deref(), - Some("acct-99") - ); - } -} +mod tests; diff --git a/rust/src/codex_accounts/api/tests.rs b/rust/src/codex_accounts/api/tests.rs new file mode 100644 index 0000000000..5b13bbbce9 --- /dev/null +++ b/rust/src/codex_accounts/api/tests.rs @@ -0,0 +1,322 @@ +use super::*; +use base64::Engine; + +#[test] +fn credential_error_states_exclude_permission_and_transport_failures() { + assert_eq!( + CodexApiError::AuthenticationRequired("missing credentials".to_string()).state_kind(), + ProviderStateKind::NeedsAuthentication + ); + assert_eq!( + CodexApiError::SessionExpired("unauthorized".to_string()).state_kind(), + ProviderStateKind::ExpiredSession + ); + assert_eq!( + CodexApiError::PermissionDenied("forbidden".to_string()).state_kind(), + ProviderStateKind::Unknown + ); + assert_eq!( + CodexApiError::Network("offline".to_string()).state_kind(), + ProviderStateKind::Unknown + ); +} + +#[tokio::test] +async fn active_fetches_use_and_sync_ambient_credentials_even_when_usage_fails() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio::time::{Duration, timeout}; + + for (target_id, managed_newer, expected_token) in [ + ("active", false, "ambient-token"), + ("active", true, "managed-token"), + ("other", true, "managed-token"), + ] { + let ambient = tempfile::tempdir().unwrap(); + let managed = tempfile::tempdir().unwrap(); + let credentials = |account: &str, token: &str| AuthCredentials { + access_token: token.into(), + refresh_token: format!("refresh-{token}"), + id_token: None, + account_id: Some(account.into()), + last_refresh: Some(Utc::now()), + }; + save_credentials(ambient.path(), &credentials("active", "ambient-token")).unwrap(); + save_credentials(managed.path(), &credentials(target_id, "managed-token")).unwrap(); + let now = Utc::now(); + for (home, refreshed) in [ + (ambient.path(), now - chrono::TimeDelta::hours(1)), + ( + managed.path(), + if managed_newer { + now + } else { + now - chrono::TimeDelta::hours(2) + }, + ), + ] { + let path = home.join("auth.json"); + let mut json: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); + json["last_refresh"] = serde_json::json!(refreshed.to_rfc3339()); + std::fs::write(path, json.to_string()).unwrap(); + } + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let config = format!( + "chatgpt_base_url = \"http://{}\"\n", + listener.local_addr().unwrap() + ); + for home in [ambient.path(), managed.path()] { + std::fs::write(home.join("config.toml"), &config).unwrap(); + } + super::super::file_locations::with_ambient_codex_home(ambient.path().to_owned()); + let api = CodexAccountApi { + client: reqwest::Client::builder().no_proxy().build().unwrap(), + }; + let server = async { + let (mut stream, _) = timeout(Duration::from_secs(5), listener.accept()) + .await + .unwrap() + .unwrap(); + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push( + timeout(Duration::from_secs(5), stream.read_u8()) + .await + .unwrap() + .unwrap(), + ); + } + stream.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}").await.unwrap(); + String::from_utf8(headers).unwrap().to_lowercase() + }; + let (result, headers) = + tokio::join!(api.fetch_snapshot(managed.path(), None, false), server); + super::super::file_locations::clear_ambient_codex_home_override(); + assert!(result.is_err()); + assert!(headers.contains(&format!("authorization: bearer {expected_token}"))); + let saved = load_credentials(managed.path()).unwrap(); + assert_eq!(saved.access_token, expected_token); + assert_eq!(saved.refresh_token, format!("refresh-{expected_token}")); + assert_eq!( + load_credentials(ambient.path()).unwrap().access_token, + if target_id == "active" { + expected_token + } else { + "ambient-token" + } + ); + } +} + +#[tokio::test] +async fn overlapping_fetches_reload_credentials_and_keep_other_homes_parallel() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::time::{Duration, timeout}; + + async fn request(listener: &TcpListener) -> (TcpStream, String) { + let (mut stream, _) = timeout(Duration::from_secs(5), listener.accept()) + .await + .unwrap() + .unwrap(); + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push( + timeout(Duration::from_secs(5), stream.read_u8()) + .await + .unwrap() + .unwrap(), + ); + } + (stream, String::from_utf8(headers).unwrap().to_lowercase()) + } + async fn respond(mut stream: TcpStream) { + stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}").await.unwrap(); + } + fn configure(home: &Path, address: std::net::SocketAddr, token: &str) { + std::fs::write( + home.join("config.toml"), + format!("chatgpt_base_url = \"http://{address}\"\n"), + ) + .unwrap(); + std::fs::write( + home.join("auth.json"), + serde_json::json!({"OPENAI_API_KEY":token}).to_string(), + ) + .unwrap(); + } + fn fetch( + home: std::path::PathBuf, + ) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + let api = CodexAccountApi { + client: reqwest::Client::builder().no_proxy().build().unwrap(), + }; + // Exercise per-home concurrency independently of other tests + // that intentionally take the global account-switch write lock. + super::super::fetch_coordination::fetch_home_snapshot( + &api, &home, None, None, false, None, + ) + .await + }) + } + + let first_home = tempfile::tempdir().unwrap(); + let other_home = tempfile::tempdir().unwrap(); + let first_server = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let other_server = TcpListener::bind("127.0.0.1:0").await.unwrap(); + configure( + first_home.path(), + first_server.local_addr().unwrap(), + "old-token", + ); + configure( + other_home.path(), + other_server.local_addr().unwrap(), + "other-token", + ); + let first = fetch(first_home.path().to_owned()); + let (first_stream, headers) = request(&first_server).await; + assert!(headers.contains("authorization: bearer old-token")); + // A lexical alias of the same auth path must share the first lane. + let second = fetch(first_home.path().join(".")); + let other = fetch(other_home.path().to_owned()); + let (other_stream, headers) = request(&other_server).await; + assert!(headers.contains("authorization: bearer other-token")); + respond(other_stream).await; + timeout(Duration::from_secs(5), other) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!( + timeout(Duration::from_millis(100), first_server.accept()) + .await + .is_err() + ); + + // Model a rotated token being persisted by the first in-flight fetch. + configure( + first_home.path(), + first_server.local_addr().unwrap(), + "rotated-token", + ); + respond(first_stream).await; + timeout(Duration::from_secs(5), first) + .await + .unwrap() + .unwrap() + .unwrap(); + let (second_stream, headers) = request(&first_server).await; + assert!(headers.contains("authorization: bearer rotated-token")); + respond(second_stream).await; + timeout(Duration::from_secs(5), second) + .await + .unwrap() + .unwrap() + .unwrap(); +} + +#[test] +fn parse_credentials_accepts_api_key() { + let creds = parse_credentials_json(r#"{"OPENAI_API_KEY":"sk-test"}"#).unwrap(); + assert_eq!(creds.access_token, "sk-test"); + assert_eq!(creds.account_id, None); +} + +#[test] +fn parse_credentials_accepts_tokens() { + let creds = parse_credentials_json( + r#"{"tokens":{"access_token":"at","refresh_token":"rt","account_id":"42"},"last_refresh":"2026-01-01T00:00:00Z"}"#, + ) + .unwrap(); + assert_eq!(creds.access_token, "at"); + assert_eq!(creds.refresh_token, "rt"); + assert_eq!(creds.account_id.as_deref(), Some("42")); +} + +#[test] +fn parse_credentials_missing_tokens_errors() { + assert!(parse_credentials_json(r#"{"foo":1}"#).is_err()); +} + +#[test] +fn jwt_payload_decodes() { + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(br#"{"email":"a@b.c"}"#); + let token = format!("eyJhbGciOiJub25lIn0.{payload}."); + let parsed = jwt_payload(&token).unwrap(); + assert_eq!(parsed.get("email").and_then(|v| v.as_str()), Some("a@b.c")); +} + +#[test] +fn normalize_window_roles_orders_session_first() { + let weekly = UsageWindowSnapshot::new(10.0, None, 604_800); + let session = UsageWindowSnapshot::new(10.0, None, 18_000); + let (p, s) = normalize_window_roles(Some(weekly), Some(session)); + assert_eq!(p.unwrap().limit_window_seconds, 18_000); + assert_eq!(s.unwrap().limit_window_seconds, 604_800); +} + +#[test] +fn limit_reached_forces_100() { + let payload = make_rate_limit(); + let (p, s) = make_normalized_windows(Some(&payload)); + assert_eq!(p.as_ref().unwrap().used_percent, 100.0); + assert_eq!(s.as_ref().unwrap().used_percent, 100.0); +} + +fn make_rate_limit() -> serde_json::Map { + serde_json::from_str( + r#"{"allowed":true,"limit_reached":true,"primary_window":{"used_percent":40,"reset_at":0,"limit_window_seconds":18000},"secondary_window":{"used_percent":20,"reset_at":0,"limit_window_seconds":604800}}"#, + ) + .unwrap() +} + +#[test] +fn resolve_usage_url_default() { + let dir = tempfile::tempdir().unwrap(); + let url = resolve_usage_url(dir.path()); + assert_eq!(url, "https://chatgpt.com/backend-api/wham/usage"); +} + +#[test] +fn parse_chatgpt_base_url_parses_quoted() { + let url = parse_chatgpt_base_url( + "# comment\nchatgpt_base_url = \"https://example.com/backend-api\"\n", + ) + .unwrap(); + assert_eq!(url, "https://example.com/backend-api"); +} + +#[test] +fn equivalent_snapshots_match() { + let mk = || AccountUsageSnapshot { + email: Some("a@b.c".to_string()), + provider_account_id: Some("x".to_string()), + plan: Some("pro".to_string()), + allowed: Some(true), + limit_reached: None, + primary_window: Some(UsageWindowSnapshot::new(12.0, Some(Utc::now()), 18_000)), + secondary_window: None, + credits: None, + cost: None, + subscription: None, + updated_at: Utc::now(), + }; + assert!(is_equivalent(&mk(), &mk())); + let mut different = mk(); + different.plan = Some("plus".to_string()); + assert!(!is_equivalent(&mk(), &different)); +} + +#[test] +fn account_id_from_id_token_reads_auth() { + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(br#"{"https://api.openai.com/auth":{"chatgpt_account_id":"acct-99"}}"#); + let token = format!("h.{payload}.s"); + assert_eq!( + account_id_from_id_token(Some(&token)).as_deref(), + Some("acct-99") + ); +} diff --git a/rust/src/codex_accounts/codex_desktop.rs b/rust/src/codex_accounts/codex_desktop.rs index ee31059013..e4a0527b67 100644 --- a/rust/src/codex_accounts/codex_desktop.rs +++ b/rust/src/codex_accounts/codex_desktop.rs @@ -14,8 +14,6 @@ use super::file_locations::{ ensure_directories, }; -pub const DEFAULT_RESTART_DELAY_SECONDS: f64 = 0.8; - /// Friendly Codex Desktop control error. #[derive(Debug, thiserror::Error)] pub enum CodexDesktopControlError { diff --git a/rust/src/codex_accounts/extra_usage.rs b/rust/src/codex_accounts/extra_usage.rs index 961b3a4213..7a21d1ac88 100644 --- a/rust/src/codex_accounts/extra_usage.rs +++ b/rust/src/codex_accounts/extra_usage.rs @@ -52,24 +52,7 @@ impl CodexExtraUsageCost { } } -impl AccountUsageSnapshot { - /// Build the account-scoped credit cost while preserving a separately - /// fetched monthly cap, if one is available. - pub fn extra_usage_cost( - &self, - attached: Option<&crate::core::CostSnapshot>, - ) -> Option { - let attached = attached.or(self.cost.as_ref()); - CodexExtraUsageCost::from_credits( - self.credits.as_ref(), - self.updated_at, - self.provider_account_id - .as_deref() - .or(self.email.as_deref()), - attached, - ) - } -} +impl AccountUsageSnapshot {} #[cfg(test)] mod extra_usage_tests { diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index 53757fb844..7fabc68105 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -46,11 +46,6 @@ fn contains_offset(text: &str) -> bool { digits && tail.contains(':') } -/// Format a UTC instant the same way CodexControl does (`...Z`). -pub fn format_datetime(value: Option>) -> Option { - value.map(|dt| dt.to_rfc3339_opts(chrono::SecondsFormat::AutoSi, true)) -} - pub fn utc_now() -> DateTime { Utc::now() } @@ -458,24 +453,6 @@ pub struct CreditsBalanceSnapshot { pub balance: Option, } -impl CreditsBalanceSnapshot { - pub fn display_value(&self) -> String { - if self.unlimited { - return "Unlimited".to_string(); - } - if let Some(balance) = self.balance { - return format!("{balance:.2}") - .trim_end_matches('0') - .trim_end_matches('.') - .to_string(); - } - if self.has_credits { - return "Available".to_string(); - } - "None".to_string() - } -} - /// A fetched snapshot for one Codex account. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -499,37 +476,6 @@ pub struct AccountUsageSnapshot { } impl AccountUsageSnapshot { - pub fn is_quota_blocked(&self) -> bool { - self.limit_reached == Some(true) || self.allowed == Some(false) - } - - pub fn has_quota_windows(&self) -> bool { - self.primary_window.is_some() || self.secondary_window.is_some() - } - - pub fn has_usable_quota_now(&self) -> bool { - if self.is_quota_blocked() { - return false; - } - let values = [self.primary_window.as_ref(), self.secondary_window.as_ref()] - .into_iter() - .flatten() - .map(|w| w.remaining_percent()); - let mut values = values.peekable(); - values.peek().is_some() && values.any(|v| v > 0.001) - } - - pub fn lowest_remaining_percent(&self) -> f64 { - if self.is_quota_blocked() { - return 0.0; - } - [self.secondary_window.as_ref(), self.primary_window.as_ref()] - .into_iter() - .flatten() - .map(|w| w.remaining_percent()) - .fold(f64::MAX, f64::min) - } - pub fn next_reset_at(&self) -> Option> { [self.primary_window.as_ref(), self.secondary_window.as_ref()] .into_iter() @@ -539,17 +485,6 @@ impl AccountUsageSnapshot { } } -/// Sort weight used to order accounts by practical usefulness. -pub fn account_sort_priority(snapshot: &AccountUsageSnapshot) -> u8 { - if snapshot.has_usable_quota_now() { - 0 - } else if snapshot.next_reset_at().is_some() { - 1 - } else { - 2 - } -} - #[cfg(test)] mod tests { use super::*; @@ -793,26 +728,6 @@ mod tests { ); } - #[test] - fn blocked_account_has_no_usable_quota() { - let snapshot = AccountUsageSnapshot { - email: None, - provider_account_id: None, - plan: None, - allowed: Some(false), - limit_reached: None, - primary_window: Some(UsageWindowSnapshot::new(10.0, None, 18_000)), - secondary_window: None, - credits: None, - cost: None, - subscription: None, - updated_at: utc_now(), - }; - assert!(snapshot.is_quota_blocked()); - assert!(!snapshot.has_usable_quota_now()); - assert_eq!(snapshot.lowest_remaining_percent(), 0.0); - } - #[test] fn parse_datetime_accepts_z_and_offset() { assert!(parse_datetime("2026-01-01T00:00:00Z").is_some()); @@ -820,46 +735,6 @@ mod tests { assert!(parse_datetime("").is_none()); } - #[test] - fn credits_display_value() { - assert_eq!( - CreditsBalanceSnapshot { - has_credits: true, - unlimited: true, - balance: None - } - .display_value(), - "Unlimited" - ); - assert_eq!( - CreditsBalanceSnapshot { - has_credits: true, - unlimited: false, - balance: Some(12.50) - } - .display_value(), - "12.5" - ); - assert_eq!( - CreditsBalanceSnapshot { - has_credits: true, - unlimited: false, - balance: None - } - .display_value(), - "Available" - ); - assert_eq!( - CreditsBalanceSnapshot { - has_credits: false, - unlimited: false, - balance: None - } - .display_value(), - "None" - ); - } - #[test] fn merge_prefers_managed_and_recency() { let mut managed = account( diff --git a/rust/src/codex_costs.rs b/rust/src/codex_costs.rs index e71bbe074e..2d42e51985 100644 --- a/rust/src/codex_costs.rs +++ b/rust/src/codex_costs.rs @@ -10,9 +10,8 @@ mod summary_contract; pub(crate) use host_costs::{CodexHostCostsArgs, HostOutputFormat, run_codex_host_costs}; pub use quota_windows::{CodexQuotaWindow, codex_quota_windows_from_cache}; pub(crate) use summary_contract::{ - CodexCostSummary, CodexHostCostReport, CodexHostCostWindow, CodexHostOutcome, - MAX_REMOTE_CODEX_COST_BYTES, REMOTE_CODEX_COST_INVALID, REMOTE_CODEX_COST_UNAVAILABLE, - decode_remote_codex_summary, + CodexCostSummary, CodexHostCostReport, CodexHostOutcome, MAX_REMOTE_CODEX_COST_BYTES, + REMOTE_CODEX_COST_INVALID, REMOTE_CODEX_COST_UNAVAILABLE, decode_remote_codex_summary, }; use chrono::{Duration, NaiveDate, Utc}; diff --git a/rust/src/codex_workspaces/indexer.rs b/rust/src/codex_workspaces/indexer.rs index 5a9b5f4f43..f62d649314 100644 --- a/rust/src/codex_workspaces/indexer.rs +++ b/rust/src/codex_workspaces/indexer.rs @@ -94,6 +94,7 @@ impl CodexWorkspacesIndex { self } + #[cfg(test)] pub fn with_sidecar_path(mut self, path: impl Into) -> Self { self.sidecar_path_override = Some(path.into()); self @@ -119,6 +120,7 @@ impl CodexWorkspacesIndex { Ok(WorkspaceUsageSidecar::new(path)) } + #[cfg(test)] pub fn load_cached_snapshot( &self, ) -> Result, IndexError> { @@ -127,12 +129,6 @@ impl CodexWorkspacesIndex { Ok(sidecar.load_latest_snapshot(scope.scope_signature(), self.history_days)?) } - pub fn clear_cached_snapshot(&self) -> Result<(), IndexError> { - let sidecar = self.sidecar()?; - sidecar.clear_snapshots()?; - Ok(()) - } - pub(super) fn load_snapshot_at( &self, force_refresh: bool, @@ -770,294 +766,4 @@ fn meta_day_out_of_range(path: &Path, range: &CostUsageDayRange) -> bool { } #[cfg(test)] -mod tests { - use super::*; - use crate::codex_workspaces::types::SourceStatus; - use chrono::Local; - use rusqlite::Connection; - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - - fn write_session( - dir: &Path, - day: &str, - name: &str, - cwd: &str, - model: &str, - input: i32, - output: i32, - ) { - // day = YYYY-MM-DD - let parts: Vec<_> = day.split('-').collect(); - let folder = dir.join(parts[0]).join(parts[1]).join(parts[2]); - fs::create_dir_all(&folder).unwrap(); - let path = folder.join(name); - let mut f = File::create(&path).unwrap(); - writeln!( - f, - r#"{{"timestamp":"{day}T10:00:00.000Z","type":"session_meta","payload":{{"session_id":"{id}","cwd":"{cwd}","originator":"codex_exec","source":"cli"}}}}"#, - id = name.trim_end_matches(".jsonl"), - cwd = cwd.replace('\\', "\\\\"), - ) - .unwrap(); - writeln!( - f, - r#"{{"timestamp":"{day}T10:00:01.000Z","type":"turn_context","payload":{{"model":"{model}"}}}}"# - ) - .unwrap(); - writeln!( - f, - r#"{{"timestamp":"{day}T10:00:02.000Z","type":"event_msg","payload":{{"type":"token_count","info":{{"last_token_usage":{{"input_tokens":{input},"cached_input_tokens":0,"output_tokens":{output}}}}}}}}}"# - ) - .unwrap(); - } - - #[test] - fn indexes_two_sessions_into_projects_and_daily() { - let tmp = TempDir::new().unwrap(); - let home = tmp.path().join("codex"); - let sessions = home.join("sessions"); - fs::create_dir_all(&sessions).unwrap(); - let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); - - write_session( - &sessions, - &day, - "sess-aaa.jsonl", - &tmp.path().join("proj-a").to_string_lossy(), - "gpt-5", - 1000, - 500, - ); - write_session( - &sessions, - &day, - "sess-bbb.jsonl", - &tmp.path().join("proj-b").to_string_lossy(), - "gpt-5", - 2000, - 100, - ); - - // Ensure project roots exist so resolve doesn't collapse oddly. - fs::create_dir_all(tmp.path().join("proj-a")).unwrap(); - fs::create_dir_all(tmp.path().join("proj-b")).unwrap(); - - let sidecar = tmp.path().join("sidecar.sqlite"); - let index = CodexWorkspacesIndex::new(30) - .with_codex_home(&home) - .with_sidecar_path(&sidecar); - - let snap = index.load_snapshot(true, |_| {}).expect("snapshot"); - assert_eq!(snap.indexed_file_count, 2); - assert_eq!(snap.projects.len(), 2); - assert!(!snap.daily.is_empty()); - assert_eq!(snap.source_status, SourceStatus::CatalogMissing); - let total_sessions: u32 = snap.projects.iter().map(|p| p.session_count).sum(); - assert_eq!(total_sessions, 2); - assert!(snap.total.input_tokens >= 3000); - assert!(snap.total.output_tokens >= 600); - - // Cached load works. - let cached = index.load_cached_snapshot().unwrap().expect("cached"); - assert_eq!(cached.indexed_file_count, 2); - } - - #[test] - fn archived_and_flat_rollouts_are_indexed_once() { - let tmp = TempDir::new().unwrap(); - let home = tmp.path().join("codex"); - let sessions = home.join("sessions"); - let archived = home.join("archived_sessions"); - fs::create_dir_all(&archived).unwrap(); - let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); - let parts: Vec<&str> = day.split('-').collect(); - let dated_dir = sessions.join(parts[0]).join(parts[1]).join(parts[2]); - let project = tmp.path().join("proj"); - fs::create_dir_all(&project).unwrap(); - for (name, input, output) in [ - ("sess-dated.jsonl", 1000, 10), - ("sess-archived.jsonl", 200, 20), - ("sess-legacy.jsonl", 30, 3), - ] { - write_session( - &sessions, - &day, - name, - &project.to_string_lossy(), - "gpt-5", - input, - output, - ); - } - // Codex archives a thread by moving its rollout into the flat - // archive; older builds kept rollouts flat in the sessions root. - fs::rename( - dated_dir.join("sess-archived.jsonl"), - archived.join("sess-archived.jsonl"), - ) - .unwrap(); - fs::rename( - dated_dir.join("sess-legacy.jsonl"), - sessions.join("sess-legacy.jsonl"), - ) - .unwrap(); - // An archived copy of a dated rollout is not counted twice. - fs::copy( - dated_dir.join("sess-dated.jsonl"), - archived.join("sess-dated.jsonl"), - ) - .unwrap(); - - let snap = CodexWorkspacesIndex::new(30) - .with_codex_home(&home) - .with_sidecar_path(tmp.path().join("sidecar.sqlite")) - .load_snapshot(true, |_| {}) - .expect("snapshot"); - - assert_eq!(snap.indexed_file_count, 3); - assert_eq!(snap.total.input_tokens, 1230); - assert_eq!(snap.total.output_tokens, 33); - let sessions_indexed: u32 = snap.projects.iter().map(|p| p.session_count).sum(); - assert_eq!(sessions_indexed, 3); - } - - #[test] - fn cached_snapshot_is_not_reused_for_another_codex_home() { - let tmp = TempDir::new().unwrap(); - let first_home = tmp.path().join("first-codex"); - let first_sessions = first_home.join("sessions"); - fs::create_dir_all(&first_sessions).unwrap(); - let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); - write_session( - &first_sessions, - &day, - "first-session.jsonl", - &tmp.path().join("first-project").to_string_lossy(), - "gpt-5", - 100, - 20, - ); - - let sidecar = tmp.path().join("sidecar.sqlite"); - let first = CodexWorkspacesIndex::new(30) - .with_codex_home(&first_home) - .with_sidecar_path(&sidecar); - let first_snapshot = first.load_snapshot(true, |_| {}).unwrap(); - assert_eq!(first_snapshot.indexed_file_count, 1); - - let second_home = tmp.path().join("second-codex"); - fs::create_dir_all(second_home.join("sessions")).unwrap(); - let second = CodexWorkspacesIndex::new(30) - .with_codex_home(&second_home) - .with_sidecar_path(&sidecar); - - assert!(second.load_cached_snapshot().unwrap().is_none()); - let second_snapshot = second.load_snapshot(false, |_| {}).unwrap(); - assert_eq!(second_snapshot.indexed_file_count, 0); - assert!(second_snapshot.projects.is_empty()); - assert_ne!( - first_snapshot.scope_signature, - second_snapshot.scope_signature - ); - } - - #[test] - fn daily_costs_use_each_day_rates() { - let sol = |day: &str| { - let models = HashMap::from([("gpt-5.6-sol".to_string(), (100, 10, 5))]); - let mut daily = HashMap::new(); - merge_daily(&mut daily, &HashMap::from([(day.to_string(), models)])); - daily[day].known_usd - }; - let historical = 90.0 * 5e-6 + 10.0 * 5e-7 + 5.0 * 3e-5; - let current = 90.0 * 4e-6 + 10.0 * 4e-7 + 5.0 * 2e-5; - assert!((sol("2026-08-20") - historical).abs() < 1e-12); - assert!((sol("2026-08-21") - current).abs() < 1e-12); - } - - #[test] - fn foreign_user_version_is_rejected() { - let tmp = TempDir::new().unwrap(); - let path = tmp.path().join("bad.sqlite"); - { - let conn = Connection::open(&path).unwrap(); - conn.pragma_update(None, "user_version", 99).unwrap(); - } - let home = tmp.path().join("codex"); - fs::create_dir_all(home.join("sessions")).unwrap(); - let index = CodexWorkspacesIndex::new(7) - .with_codex_home(home) - .with_sidecar_path(path); - let err = index.load_snapshot(true, |_| {}).unwrap_err(); - let msg = err.to_string(); - assert!( - msg.contains("incompatible") || msg.contains("user_version"), - "unexpected error: {msg}" - ); - } - - #[test] - fn missing_catalog_reports_catalog_missing() { - let tmp = TempDir::new().unwrap(); - let home = tmp.path().join("codex"); - fs::create_dir_all(home.join("sessions")).unwrap(); - let index = CodexWorkspacesIndex::new(7) - .with_codex_home(&home) - .with_sidecar_path(tmp.path().join("side.sqlite")); - let snap = index.load_snapshot(true, |_| {}).unwrap(); - assert_eq!(snap.source_status, SourceStatus::CatalogMissing); - } - - #[test] - fn privacy_redaction_strips_paths_and_titles() { - let mut snap = CodexLocalProjectUsageSnapshot { - updated_at: Utc::now(), - history_days: 30, - scope_signature: "x".into(), - indexed_file_count: 1, - skipped_file_count: 0, - total: UsageTotals::from_parts(10, 0, 5), - sessions: vec![SessionUsage { - id: "s1".into(), - project_id: "project-abc".into(), - display_title: "do not leak".into(), - cwd: Some("/Users/me/secret-repo".into()), - started_at: None, - latest_activity: None, - totals: UsageTotals::from_parts(10, 0, 5), - cost_estimate: CostEstimate::default(), - top_model: Some("gpt-5".into()), - }], - projects: vec![ProjectUsage { - id: "project-abc".into(), - display_name: "secret-repo".into(), - path: Some("/Users/me/secret-repo".into()), - totals: UsageTotals::from_parts(10, 0, 5), - cost_estimate: CostEstimate::default(), - session_count: 1, - latest_activity: None, - top_model: Some("gpt-5".into()), - top_sessions: vec![SessionUsage { - id: "s1".into(), - project_id: "project-abc".into(), - display_title: "do not leak".into(), - cwd: Some("/Users/me/secret-repo".into()), - started_at: None, - latest_activity: None, - totals: UsageTotals::from_parts(10, 0, 5), - cost_estimate: CostEstimate::default(), - top_model: Some("gpt-5".into()), - }], - }], - daily: vec![], - source_status: SourceStatus::Complete, - }; - snap.redact_for_privacy(); - assert_eq!(snap.projects[0].display_name, "Workspace"); - assert!(snap.projects[0].path.is_none()); - assert_eq!(snap.projects[0].top_sessions[0].display_title, "Session s1"); - assert!(snap.projects[0].top_sessions[0].cwd.is_none()); - } -} +mod tests; diff --git a/rust/src/codex_workspaces/indexer/tests.rs b/rust/src/codex_workspaces/indexer/tests.rs new file mode 100644 index 0000000000..be6a1d54ec --- /dev/null +++ b/rust/src/codex_workspaces/indexer/tests.rs @@ -0,0 +1,289 @@ +use super::*; +use crate::codex_workspaces::types::SourceStatus; +use chrono::Local; +use rusqlite::Connection; +use std::fs::File; +use std::io::Write; +use tempfile::TempDir; + +fn write_session( + dir: &Path, + day: &str, + name: &str, + cwd: &str, + model: &str, + input: i32, + output: i32, +) { + // day = YYYY-MM-DD + let parts: Vec<_> = day.split('-').collect(); + let folder = dir.join(parts[0]).join(parts[1]).join(parts[2]); + fs::create_dir_all(&folder).unwrap(); + let path = folder.join(name); + let mut f = File::create(&path).unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:00.000Z","type":"session_meta","payload":{{"session_id":"{id}","cwd":"{cwd}","originator":"codex_exec","source":"cli"}}}}"#, + id = name.trim_end_matches(".jsonl"), + cwd = cwd.replace('\\', "\\\\"), + ) + .unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:01.000Z","type":"turn_context","payload":{{"model":"{model}"}}}}"# + ) + .unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:02.000Z","type":"event_msg","payload":{{"type":"token_count","info":{{"last_token_usage":{{"input_tokens":{input},"cached_input_tokens":0,"output_tokens":{output}}}}}}}}}"# + ) + .unwrap(); +} + +#[test] +fn indexes_two_sessions_into_projects_and_daily() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("codex"); + let sessions = home.join("sessions"); + fs::create_dir_all(&sessions).unwrap(); + let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); + + write_session( + &sessions, + &day, + "sess-aaa.jsonl", + &tmp.path().join("proj-a").to_string_lossy(), + "gpt-5", + 1000, + 500, + ); + write_session( + &sessions, + &day, + "sess-bbb.jsonl", + &tmp.path().join("proj-b").to_string_lossy(), + "gpt-5", + 2000, + 100, + ); + + // Ensure project roots exist so resolve doesn't collapse oddly. + fs::create_dir_all(tmp.path().join("proj-a")).unwrap(); + fs::create_dir_all(tmp.path().join("proj-b")).unwrap(); + + let sidecar = tmp.path().join("sidecar.sqlite"); + let index = CodexWorkspacesIndex::new(30) + .with_codex_home(&home) + .with_sidecar_path(&sidecar); + + let snap = index.load_snapshot(true, |_| {}).expect("snapshot"); + assert_eq!(snap.indexed_file_count, 2); + assert_eq!(snap.projects.len(), 2); + assert!(!snap.daily.is_empty()); + assert_eq!(snap.source_status, SourceStatus::CatalogMissing); + let total_sessions: u32 = snap.projects.iter().map(|p| p.session_count).sum(); + assert_eq!(total_sessions, 2); + assert!(snap.total.input_tokens >= 3000); + assert!(snap.total.output_tokens >= 600); + + // Cached load works. + let cached = index.load_cached_snapshot().unwrap().expect("cached"); + assert_eq!(cached.indexed_file_count, 2); +} + +#[test] +fn archived_and_flat_rollouts_are_indexed_once() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("codex"); + let sessions = home.join("sessions"); + let archived = home.join("archived_sessions"); + fs::create_dir_all(&archived).unwrap(); + let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); + let parts: Vec<&str> = day.split('-').collect(); + let dated_dir = sessions.join(parts[0]).join(parts[1]).join(parts[2]); + let project = tmp.path().join("proj"); + fs::create_dir_all(&project).unwrap(); + for (name, input, output) in [ + ("sess-dated.jsonl", 1000, 10), + ("sess-archived.jsonl", 200, 20), + ("sess-legacy.jsonl", 30, 3), + ] { + write_session( + &sessions, + &day, + name, + &project.to_string_lossy(), + "gpt-5", + input, + output, + ); + } + // Codex archives a thread by moving its rollout into the flat + // archive; older builds kept rollouts flat in the sessions root. + fs::rename( + dated_dir.join("sess-archived.jsonl"), + archived.join("sess-archived.jsonl"), + ) + .unwrap(); + fs::rename( + dated_dir.join("sess-legacy.jsonl"), + sessions.join("sess-legacy.jsonl"), + ) + .unwrap(); + // An archived copy of a dated rollout is not counted twice. + fs::copy( + dated_dir.join("sess-dated.jsonl"), + archived.join("sess-dated.jsonl"), + ) + .unwrap(); + + let snap = CodexWorkspacesIndex::new(30) + .with_codex_home(&home) + .with_sidecar_path(tmp.path().join("sidecar.sqlite")) + .load_snapshot(true, |_| {}) + .expect("snapshot"); + + assert_eq!(snap.indexed_file_count, 3); + assert_eq!(snap.total.input_tokens, 1230); + assert_eq!(snap.total.output_tokens, 33); + let sessions_indexed: u32 = snap.projects.iter().map(|p| p.session_count).sum(); + assert_eq!(sessions_indexed, 3); +} + +#[test] +fn cached_snapshot_is_not_reused_for_another_codex_home() { + let tmp = TempDir::new().unwrap(); + let first_home = tmp.path().join("first-codex"); + let first_sessions = first_home.join("sessions"); + fs::create_dir_all(&first_sessions).unwrap(); + let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); + write_session( + &first_sessions, + &day, + "first-session.jsonl", + &tmp.path().join("first-project").to_string_lossy(), + "gpt-5", + 100, + 20, + ); + + let sidecar = tmp.path().join("sidecar.sqlite"); + let first = CodexWorkspacesIndex::new(30) + .with_codex_home(&first_home) + .with_sidecar_path(&sidecar); + let first_snapshot = first.load_snapshot(true, |_| {}).unwrap(); + assert_eq!(first_snapshot.indexed_file_count, 1); + + let second_home = tmp.path().join("second-codex"); + fs::create_dir_all(second_home.join("sessions")).unwrap(); + let second = CodexWorkspacesIndex::new(30) + .with_codex_home(&second_home) + .with_sidecar_path(&sidecar); + + assert!(second.load_cached_snapshot().unwrap().is_none()); + let second_snapshot = second.load_snapshot(false, |_| {}).unwrap(); + assert_eq!(second_snapshot.indexed_file_count, 0); + assert!(second_snapshot.projects.is_empty()); + assert_ne!( + first_snapshot.scope_signature, + second_snapshot.scope_signature + ); +} + +#[test] +fn daily_costs_use_each_day_rates() { + let sol = |day: &str| { + let models = HashMap::from([("gpt-5.6-sol".to_string(), (100, 10, 5))]); + let mut daily = HashMap::new(); + merge_daily(&mut daily, &HashMap::from([(day.to_string(), models)])); + daily[day].known_usd + }; + let historical = 90.0 * 5e-6 + 10.0 * 5e-7 + 5.0 * 3e-5; + let current = 90.0 * 4e-6 + 10.0 * 4e-7 + 5.0 * 2e-5; + assert!((sol("2026-08-20") - historical).abs() < 1e-12); + assert!((sol("2026-08-21") - current).abs() < 1e-12); +} + +#[test] +fn foreign_user_version_is_rejected() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("bad.sqlite"); + { + let conn = Connection::open(&path).unwrap(); + conn.pragma_update(None, "user_version", 99).unwrap(); + } + let home = tmp.path().join("codex"); + fs::create_dir_all(home.join("sessions")).unwrap(); + let index = CodexWorkspacesIndex::new(7) + .with_codex_home(home) + .with_sidecar_path(path); + let err = index.load_snapshot(true, |_| {}).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("incompatible") || msg.contains("user_version"), + "unexpected error: {msg}" + ); +} + +#[test] +fn missing_catalog_reports_catalog_missing() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("codex"); + fs::create_dir_all(home.join("sessions")).unwrap(); + let index = CodexWorkspacesIndex::new(7) + .with_codex_home(&home) + .with_sidecar_path(tmp.path().join("side.sqlite")); + let snap = index.load_snapshot(true, |_| {}).unwrap(); + assert_eq!(snap.source_status, SourceStatus::CatalogMissing); +} + +#[test] +fn privacy_redaction_strips_paths_and_titles() { + let mut snap = CodexLocalProjectUsageSnapshot { + updated_at: Utc::now(), + history_days: 30, + scope_signature: "x".into(), + indexed_file_count: 1, + skipped_file_count: 0, + total: UsageTotals::from_parts(10, 0, 5), + sessions: vec![SessionUsage { + id: "s1".into(), + project_id: "project-abc".into(), + display_title: "do not leak".into(), + cwd: Some("/Users/me/secret-repo".into()), + started_at: None, + latest_activity: None, + totals: UsageTotals::from_parts(10, 0, 5), + cost_estimate: CostEstimate::default(), + top_model: Some("gpt-5".into()), + }], + projects: vec![ProjectUsage { + id: "project-abc".into(), + display_name: "secret-repo".into(), + path: Some("/Users/me/secret-repo".into()), + totals: UsageTotals::from_parts(10, 0, 5), + cost_estimate: CostEstimate::default(), + session_count: 1, + latest_activity: None, + top_model: Some("gpt-5".into()), + top_sessions: vec![SessionUsage { + id: "s1".into(), + project_id: "project-abc".into(), + display_title: "do not leak".into(), + cwd: Some("/Users/me/secret-repo".into()), + started_at: None, + latest_activity: None, + totals: UsageTotals::from_parts(10, 0, 5), + cost_estimate: CostEstimate::default(), + top_model: Some("gpt-5".into()), + }], + }], + daily: vec![], + source_status: SourceStatus::Complete, + }; + snap.redact_for_privacy(); + assert_eq!(snap.projects[0].display_name, "Workspace"); + assert!(snap.projects[0].path.is_none()); + assert_eq!(snap.projects[0].top_sessions[0].display_title, "Session s1"); + assert!(snap.projects[0].top_sessions[0].cwd.is_none()); +} diff --git a/rust/src/codex_workspaces/sidecar.rs b/rust/src/codex_workspaces/sidecar.rs index a87aa84f10..940b138c68 100644 --- a/rust/src/codex_workspaces/sidecar.rs +++ b/rust/src/codex_workspaces/sidecar.rs @@ -138,15 +138,6 @@ impl WorkspaceUsageSidecar { Ok(()) } - pub fn clear_snapshots(&self) -> Result<(), SidecarError> { - if !self.path.exists() { - return Ok(()); - } - let conn = self.open(false)?; - conn.execute("DELETE FROM snapshot_payloads", [])?; - Ok(()) - } - fn open(&self, read_only: bool) -> Result { if read_only { let conn = Connection::open_with_flags( diff --git a/rust/src/core/adaptive_refresh.rs b/rust/src/core/adaptive_refresh.rs index 85da81679b..703335faee 100644 --- a/rust/src/core/adaptive_refresh.rs +++ b/rust/src/core/adaptive_refresh.rs @@ -52,9 +52,6 @@ const LONG_IDLE_DELAY: Duration = Duration::from_secs(30 * 60); const CONSTRAINED_DELAY: Duration = Duration::from_secs(30 * 60); const CODING_ACTIVITY_DELAY_CAP: Duration = Duration::from_secs(5 * 60); -/// Representative cadence for consumers that need one interval but cannot access live state. -pub const NOMINAL_INTERVAL_FOR_HEURISTICS: Duration = Duration::from_secs(5 * 60); - /// Compute the next adaptive refresh delay from activity / power inputs. pub fn next_delay(input: AdaptiveRefreshInput) -> AdaptiveRefreshDecision { if input.low_power_mode_enabled || input.thermal_pressure == ThermalPressure::Constrained { diff --git a/rust/src/core/credentials.rs b/rust/src/core/credentials.rs deleted file mode 100755 index f01a7070b3..0000000000 --- a/rust/src/core/credentials.rs +++ /dev/null @@ -1,110 +0,0 @@ -//! Credential storage abstraction - -use thiserror::Error; - -/// Errors that can occur with credential operations -#[derive(Debug, Error)] -pub enum CredentialError { - #[error("Credential not found")] - NotFound, - - #[error("Access denied")] - AccessDenied, - - #[error("Storage error: {0}")] - Storage(String), - - #[error("Invalid credential format")] - InvalidFormat, -} - -/// Trait for credential storage backends -pub trait CredentialStore: Send + Sync { - /// Get a credential by key - fn get(&self, service: &str, key: &str) -> Result; - - /// Set a credential - fn set(&self, service: &str, key: &str, value: &str) -> Result<(), CredentialError>; - - /// Delete a credential - fn delete(&self, service: &str, key: &str) -> Result<(), CredentialError>; -} - -/// Windows Credential Manager implementation -#[cfg(windows)] -pub struct WindowsCredentialStore; - -#[cfg(windows)] -impl WindowsCredentialStore { - pub fn new() -> Self { - Self - } -} - -#[cfg(windows)] -impl Default for WindowsCredentialStore { - fn default() -> Self { - Self::new() - } -} - -#[cfg(windows)] -impl CredentialStore for WindowsCredentialStore { - fn get(&self, service: &str, key: &str) -> Result { - let entry = keyring::Entry::new(service, key) - .map_err(|e| CredentialError::Storage(e.to_string()))?; - entry.get_password().map_err(|e| match e { - keyring::Error::NoEntry => CredentialError::NotFound, - keyring::Error::Ambiguous(_) => CredentialError::Storage("Ambiguous entry".to_string()), - _ => CredentialError::Storage(e.to_string()), - }) - } - - fn set(&self, service: &str, key: &str, value: &str) -> Result<(), CredentialError> { - let entry = keyring::Entry::new(service, key) - .map_err(|e| CredentialError::Storage(e.to_string()))?; - entry - .set_password(value) - .map_err(|e| CredentialError::Storage(e.to_string())) - } - - fn delete(&self, service: &str, key: &str) -> Result<(), CredentialError> { - let entry = keyring::Entry::new(service, key) - .map_err(|e| CredentialError::Storage(e.to_string()))?; - entry.delete_credential().map_err(|e| match e { - keyring::Error::NoEntry => CredentialError::NotFound, - _ => CredentialError::Storage(e.to_string()), - }) - } -} - -/// OAuth credentials structure -#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] -pub struct OAuthCredentials { - pub access_token: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub refresh_token: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub expires_at: Option>, - #[serde(default)] - pub scopes: Vec, - #[serde(skip_serializing_if = "Option::is_none")] - pub rate_limit_tier: Option, -} - -impl OAuthCredentials { - /// Check if the token is expired - pub fn is_expired(&self) -> bool { - if let Some(expires_at) = self.expires_at { - // Consider expired if within 5 minutes of expiry - expires_at <= chrono::Utc::now() + chrono::Duration::minutes(5) - } else { - false - } - } - - /// Check if the credentials have a specific scope - pub fn has_scope(&self, scope: &str) -> bool { - self.scopes.iter().any(|s| s == scope) - } -} diff --git a/rust/src/core/hook_transition.rs b/rust/src/core/hook_transition.rs index 48a70f9db5..af9fa36a26 100644 --- a/rust/src/core/hook_transition.rs +++ b/rust/src/core/hook_transition.rs @@ -225,14 +225,6 @@ impl HookTransitionDetector { } } - pub fn with_thresholds(reached_threshold: f64, reset_drop_threshold: f64) -> Self { - Self { - reached_threshold, - reset_drop_threshold, - ..Self::default() - } - } - /// Drops every baseline when the hook configuration changed. pub fn reset_if_configuration_changed(&mut self, revision: i64) { if self.config_revision == Some(revision) { @@ -571,662 +563,4 @@ pub fn dispatch_usage_updated_hook( } #[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - use std::path::PathBuf; - - const PROVIDER: &str = "codex"; - - fn lane_key(window: HookQuotaWindow, account: Option<&str>) -> HookQuotaLaneKey { - HookQuotaLaneKey::new(PROVIDER, window, account.map(str::to_string), None) - } - - fn rate_window(used_percent: f64, resets_at: Option>) -> RateWindow { - let mut w = RateWindow::new(used_percent); - w.window_minutes = Some(300); - w.resets_at = resets_at; - w - } - - fn informational_window(used_percent: f64) -> RateWindow { - let mut w = RateWindow::informational("placeholder"); - w.used_percent = used_percent; - w - } - - fn lane( - used_percent: Option, - key: HookQuotaLaneKey, - resets_at: Option>, - thresholds: &[f64], - informational: bool, - ) -> HookQuotaLaneObservation { - let label = key.window.display_name().to_string(); - let account = key.account_discriminator.clone(); - HookQuotaLaneObservation { - key, - label, - rate_window: used_percent.map(|p| { - if informational { - informational_window(p) - } else { - rate_window(p, resets_at) - } - }), - fallback_thresholds: thresholds.to_vec(), - account_display_name: account, - } - } - - fn observation( - lanes: Vec, - status: HookProviderStatus, - refresh_failure: Option<&str>, - ) -> HookProviderObservation { - HookProviderObservation { - provider: PROVIDER.into(), - lanes, - status, - refresh_failure_status: refresh_failure.map(str::to_string), - account_display_name: None, - successful_usage: None, - account_discriminator: None, - } - } - - fn rule(event: HookEventType, threshold: Option, provider: Option<&str>) -> HookRule { - HookRule { - enabled: true, - event: Some(event), - events: Vec::new(), - provider: provider.map(str::to_string), - threshold, - executable: PathBuf::from("/bin/true"), - arguments: Vec::new(), - timeout_secs: 10, - } - } - - fn config(enabled: bool, rules: Option>) -> HooksConfig { - HooksConfig { - enabled, - events: rules.unwrap_or_else(|| { - vec![ - rule(HookEventType::QuotaLow, None, None), - rule(HookEventType::QuotaReached, None, None), - rule(HookEventType::QuotaReset, None, None), - rule(HookEventType::ProviderUnavailable, None, None), - rule(HookEventType::ProviderRecovered, None, None), - rule(HookEventType::RefreshFailed, None, None), - ] - }), - } - } - - fn events_of(dispatches: &[HookDispatch]) -> Vec { - dispatches.iter().map(|d| d.event.event).collect() - } - - #[test] - fn first_sample_establishes_baseline_without_firing() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let dispatches = detector.evaluate( - &observation( - vec![lane( - Some(95.0), - lane_key(HookQuotaWindow::Session, None), - None, - &[0.8], - false, - )], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn quota_low_fires_once_on_upward_crossing() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - - let crossing = detector.evaluate( - &observation( - vec![lane(Some(85.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(events_of(&crossing), vec![HookEventType::QuotaLow]); - - let persisting = detector.evaluate( - &observation( - vec![lane(Some(90.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(persisting.is_empty()); - } - - #[test] - fn quota_low_dispatches_only_rule_whose_threshold_crossed() { - let mut detector = HookTransitionDetector::new(); - let cfg = config( - true, - Some(vec![ - rule(HookEventType::QuotaLow, Some(0.5), None), - rule(HookEventType::QuotaLow, Some(0.8), None), - ]), - ); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(60.0), key.clone(), None, &[], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(85.0), key, None, &[], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(dispatches.len(), 1); - let rules = dispatches[0].rules.as_ref().expect("narrowed rules"); - assert_eq!(rules.len(), 1); - assert_eq!(rules[0].threshold, Some(0.8)); - } - - #[test] - fn quota_reached_fires_on_upward_edge_only() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(90.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - - let reached = detector.evaluate( - &observation( - vec![lane(Some(100.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!( - reached - .iter() - .any(|d| d.event.event == HookEventType::QuotaReached) - ); - - let still_full = detector.evaluate( - &observation( - vec![lane(Some(100.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!( - !still_full - .iter() - .any(|d| d.event.event == HookEventType::QuotaReached) - ); - } - - #[test] - fn quota_reached_never_fires_for_weekly_lane() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Weekly, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(90.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(100.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!( - !dispatches - .iter() - .any(|d| d.event.event == HookEventType::QuotaReached) - ); - } - - #[test] - fn quota_reset_fires_when_reset_boundary_advances() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let first = Utc.timestamp_opt(1_000_000, 0).unwrap(); - let second = first + chrono::Duration::seconds(18_000); - - let _ = detector.evaluate( - &observation( - vec![lane(Some(100.0), key.clone(), Some(first), &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(0.0), key, Some(second), &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(events_of(&dispatches), vec![HookEventType::QuotaReset]); - } - - #[test] - fn quota_reset_fires_on_usage_drop_without_boundary() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(95.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(10.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(events_of(&dispatches), vec![HookEventType::QuotaReset]); - } - - #[test] - fn reset_suppresses_depletion_edge_in_same_poll() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(95.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(5.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!( - !dispatches - .iter() - .any(|d| d.event.event == HookEventType::QuotaReached) - ); - assert!( - !dispatches - .iter() - .any(|d| d.event.event == HookEventType::QuotaLow) - ); - } - - #[test] - fn provider_status_fires_outage_and_recovery_edges() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let _ = detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); - - let outage = detector.evaluate(&observation(vec![], HookProviderStatus::Major, None), &cfg); - assert_eq!(events_of(&outage), vec![HookEventType::ProviderUnavailable]); - - let persisting = detector.evaluate( - &observation(vec![], HookProviderStatus::Critical, None), - &cfg, - ); - assert!(persisting.is_empty()); - - let recovered = - detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); - assert_eq!( - events_of(&recovered), - vec![HookEventType::ProviderRecovered] - ); - } - - #[test] - fn unknown_and_maintenance_never_flip_status_state() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let _ = detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); - - assert!( - detector - .evaluate( - &observation(vec![], HookProviderStatus::Unknown, None), - &cfg - ) - .is_empty() - ); - assert!( - detector - .evaluate( - &observation(vec![], HookProviderStatus::Maintenance, None), - &cfg - ) - .is_empty() - ); - - let outage = detector.evaluate(&observation(vec![], HookProviderStatus::Major, None), &cfg); - assert_eq!(events_of(&outage), vec![HookEventType::ProviderUnavailable]); - } - - #[test] - fn first_definite_status_does_not_fire() { - let mut detector = HookTransitionDetector::new(); - let dispatches = detector.evaluate( - &observation(vec![], HookProviderStatus::Critical, None), - &config(true, None), - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn refresh_failure_emits_coarse_status_only() { - let mut detector = HookTransitionDetector::new(); - let dispatches = detector.evaluate( - &observation(vec![], HookProviderStatus::Unknown, Some("timeout")), - &config(true, None), - ); - assert_eq!(events_of(&dispatches), vec![HookEventType::RefreshFailed]); - assert_eq!(dispatches[0].event.status.as_deref(), Some("timeout")); - } - - #[test] - fn refresh_failure_does_not_disturb_quota_baselines() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let _ = detector.evaluate( - &observation(vec![], HookProviderStatus::Unknown, Some("offline")), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(85.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!( - dispatches - .iter() - .any(|d| d.event.event == HookEventType::QuotaLow) - ); - } - - #[test] - fn disabled_hooks_produce_no_events() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(false, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(95.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn configuration_change_clears_baselines() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - detector.reset_if_configuration_changed(1); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - detector.reset_if_configuration_changed(2); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(95.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn synthetic_placeholder_lane_never_fires() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(100.0), key, None, &[0.8], true)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn disappearing_lane_resets_baseline() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let _ = detector.evaluate( - &observation(vec![], HookProviderStatus::Unknown, None), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(95.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn accounts_on_same_provider_track_independently() { - let mut detector = HookTransitionDetector::new(); - let cfg = config(true, None); - let first = lane_key(HookQuotaWindow::Session, Some("a@example.com")); - let second = lane_key(HookQuotaWindow::Session, Some("b@example.com")); - let _ = detector.evaluate( - &observation( - vec![ - lane(Some(50.0), first.clone(), None, &[0.8], false), - lane(Some(50.0), second.clone(), None, &[0.8], false), - ], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![ - lane(Some(85.0), first, None, &[0.8], false), - lane(Some(55.0), second, None, &[0.8], false), - ], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(dispatches.len(), 1); - assert_eq!(dispatches[0].event.event, HookEventType::QuotaLow); - } - - #[test] - fn quota_low_respects_explicit_rule_threshold_over_fallback() { - let mut detector = HookTransitionDetector::new(); - let cfg = config( - true, - Some(vec![rule(HookEventType::QuotaLow, Some(0.9), None)]), - ); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let below = detector.evaluate( - &observation( - vec![lane(Some(85.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(below.is_empty()); - let above = detector.evaluate( - &observation( - vec![lane(Some(95.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert_eq!(events_of(&above), vec![HookEventType::QuotaLow]); - } - - #[test] - fn quota_low_ignores_rules_scoped_to_another_provider() { - let mut detector = HookTransitionDetector::new(); - let cfg = config( - true, - Some(vec![rule(HookEventType::QuotaLow, None, Some("claude"))]), - ); - let key = lane_key(HookQuotaWindow::Session, None); - let _ = detector.evaluate( - &observation( - vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - let dispatches = detector.evaluate( - &observation( - vec![lane(Some(95.0), key, None, &[0.8], false)], - HookProviderStatus::Unknown, - None, - ), - &cfg, - ); - assert!(dispatches.is_empty()); - } - - #[test] - fn rate_limiter_suppresses_duplicate_refresh_failed() { - use super::super::hooks::HookRateLimiter; - use std::time::Duration; - - let limiter = HookRateLimiter::new(Duration::from_secs(600)); - let event = HookEvent::new(HookEventType::RefreshFailed, "codex").with_status("timeout"); - assert!(limiter.allow(&event, None)); - assert!(!limiter.allow(&event, None)); - - // Quota events are not rate-limited by HookEventType::is_rate_limited. - assert!(!HookEventType::QuotaLow.is_rate_limited()); - assert!(HookEventType::UsageUpdated.is_rate_limited()); - assert!(HookEventType::RefreshFailed.is_rate_limited()); - assert!(HookEventType::ProviderUnavailable.is_rate_limited()); - } -} +mod tests; diff --git a/rust/src/core/hook_transition/tests.rs b/rust/src/core/hook_transition/tests.rs new file mode 100644 index 0000000000..2458869a4c --- /dev/null +++ b/rust/src/core/hook_transition/tests.rs @@ -0,0 +1,656 @@ +use super::*; +use chrono::TimeZone; +use std::path::PathBuf; + +const PROVIDER: &str = "codex"; + +fn lane_key(window: HookQuotaWindow, account: Option<&str>) -> HookQuotaLaneKey { + HookQuotaLaneKey::new(PROVIDER, window, account.map(str::to_string), None) +} + +fn rate_window(used_percent: f64, resets_at: Option>) -> RateWindow { + let mut w = RateWindow::new(used_percent); + w.window_minutes = Some(300); + w.resets_at = resets_at; + w +} + +fn informational_window(used_percent: f64) -> RateWindow { + let mut w = RateWindow::informational("placeholder"); + w.used_percent = used_percent; + w +} + +fn lane( + used_percent: Option, + key: HookQuotaLaneKey, + resets_at: Option>, + thresholds: &[f64], + informational: bool, +) -> HookQuotaLaneObservation { + let label = key.window.display_name().to_string(); + let account = key.account_discriminator.clone(); + HookQuotaLaneObservation { + key, + label, + rate_window: used_percent.map(|p| { + if informational { + informational_window(p) + } else { + rate_window(p, resets_at) + } + }), + fallback_thresholds: thresholds.to_vec(), + account_display_name: account, + } +} + +fn observation( + lanes: Vec, + status: HookProviderStatus, + refresh_failure: Option<&str>, +) -> HookProviderObservation { + HookProviderObservation { + provider: PROVIDER.into(), + lanes, + status, + refresh_failure_status: refresh_failure.map(str::to_string), + account_display_name: None, + successful_usage: None, + account_discriminator: None, + } +} + +fn rule(event: HookEventType, threshold: Option, provider: Option<&str>) -> HookRule { + HookRule { + enabled: true, + event: Some(event), + events: Vec::new(), + provider: provider.map(str::to_string), + threshold, + executable: PathBuf::from("/bin/true"), + arguments: Vec::new(), + timeout_secs: 10, + } +} + +fn config(enabled: bool, rules: Option>) -> HooksConfig { + HooksConfig { + enabled, + events: rules.unwrap_or_else(|| { + vec![ + rule(HookEventType::QuotaLow, None, None), + rule(HookEventType::QuotaReached, None, None), + rule(HookEventType::QuotaReset, None, None), + rule(HookEventType::ProviderUnavailable, None, None), + rule(HookEventType::ProviderRecovered, None, None), + rule(HookEventType::RefreshFailed, None, None), + ] + }), + } +} + +fn events_of(dispatches: &[HookDispatch]) -> Vec { + dispatches.iter().map(|d| d.event.event).collect() +} + +#[test] +fn first_sample_establishes_baseline_without_firing() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let dispatches = detector.evaluate( + &observation( + vec![lane( + Some(95.0), + lane_key(HookQuotaWindow::Session, None), + None, + &[0.8], + false, + )], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn quota_low_fires_once_on_upward_crossing() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + + let crossing = detector.evaluate( + &observation( + vec![lane(Some(85.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(events_of(&crossing), vec![HookEventType::QuotaLow]); + + let persisting = detector.evaluate( + &observation( + vec![lane(Some(90.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(persisting.is_empty()); +} + +#[test] +fn quota_low_dispatches_only_rule_whose_threshold_crossed() { + let mut detector = HookTransitionDetector::new(); + let cfg = config( + true, + Some(vec![ + rule(HookEventType::QuotaLow, Some(0.5), None), + rule(HookEventType::QuotaLow, Some(0.8), None), + ]), + ); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(60.0), key.clone(), None, &[], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(85.0), key, None, &[], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(dispatches.len(), 1); + let rules = dispatches[0].rules.as_ref().expect("narrowed rules"); + assert_eq!(rules.len(), 1); + assert_eq!(rules[0].threshold, Some(0.8)); +} + +#[test] +fn quota_reached_fires_on_upward_edge_only() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(90.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + + let reached = detector.evaluate( + &observation( + vec![lane(Some(100.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!( + reached + .iter() + .any(|d| d.event.event == HookEventType::QuotaReached) + ); + + let still_full = detector.evaluate( + &observation( + vec![lane(Some(100.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!( + !still_full + .iter() + .any(|d| d.event.event == HookEventType::QuotaReached) + ); +} + +#[test] +fn quota_reached_never_fires_for_weekly_lane() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Weekly, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(90.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(100.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!( + !dispatches + .iter() + .any(|d| d.event.event == HookEventType::QuotaReached) + ); +} + +#[test] +fn quota_reset_fires_when_reset_boundary_advances() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let first = Utc.timestamp_opt(1_000_000, 0).unwrap(); + let second = first + chrono::Duration::seconds(18_000); + + let _ = detector.evaluate( + &observation( + vec![lane(Some(100.0), key.clone(), Some(first), &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(0.0), key, Some(second), &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(events_of(&dispatches), vec![HookEventType::QuotaReset]); +} + +#[test] +fn quota_reset_fires_on_usage_drop_without_boundary() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(95.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(10.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(events_of(&dispatches), vec![HookEventType::QuotaReset]); +} + +#[test] +fn reset_suppresses_depletion_edge_in_same_poll() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(95.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(5.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!( + !dispatches + .iter() + .any(|d| d.event.event == HookEventType::QuotaReached) + ); + assert!( + !dispatches + .iter() + .any(|d| d.event.event == HookEventType::QuotaLow) + ); +} + +#[test] +fn provider_status_fires_outage_and_recovery_edges() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let _ = detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); + + let outage = detector.evaluate(&observation(vec![], HookProviderStatus::Major, None), &cfg); + assert_eq!(events_of(&outage), vec![HookEventType::ProviderUnavailable]); + + let persisting = detector.evaluate( + &observation(vec![], HookProviderStatus::Critical, None), + &cfg, + ); + assert!(persisting.is_empty()); + + let recovered = detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); + assert_eq!( + events_of(&recovered), + vec![HookEventType::ProviderRecovered] + ); +} + +#[test] +fn unknown_and_maintenance_never_flip_status_state() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let _ = detector.evaluate(&observation(vec![], HookProviderStatus::None, None), &cfg); + + assert!( + detector + .evaluate( + &observation(vec![], HookProviderStatus::Unknown, None), + &cfg + ) + .is_empty() + ); + assert!( + detector + .evaluate( + &observation(vec![], HookProviderStatus::Maintenance, None), + &cfg + ) + .is_empty() + ); + + let outage = detector.evaluate(&observation(vec![], HookProviderStatus::Major, None), &cfg); + assert_eq!(events_of(&outage), vec![HookEventType::ProviderUnavailable]); +} + +#[test] +fn first_definite_status_does_not_fire() { + let mut detector = HookTransitionDetector::new(); + let dispatches = detector.evaluate( + &observation(vec![], HookProviderStatus::Critical, None), + &config(true, None), + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn refresh_failure_emits_coarse_status_only() { + let mut detector = HookTransitionDetector::new(); + let dispatches = detector.evaluate( + &observation(vec![], HookProviderStatus::Unknown, Some("timeout")), + &config(true, None), + ); + assert_eq!(events_of(&dispatches), vec![HookEventType::RefreshFailed]); + assert_eq!(dispatches[0].event.status.as_deref(), Some("timeout")); +} + +#[test] +fn refresh_failure_does_not_disturb_quota_baselines() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let _ = detector.evaluate( + &observation(vec![], HookProviderStatus::Unknown, Some("offline")), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(85.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!( + dispatches + .iter() + .any(|d| d.event.event == HookEventType::QuotaLow) + ); +} + +#[test] +fn disabled_hooks_produce_no_events() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(false, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(95.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn configuration_change_clears_baselines() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + detector.reset_if_configuration_changed(1); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + detector.reset_if_configuration_changed(2); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(95.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn synthetic_placeholder_lane_never_fires() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(100.0), key, None, &[0.8], true)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn disappearing_lane_resets_baseline() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let _ = detector.evaluate( + &observation(vec![], HookProviderStatus::Unknown, None), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(95.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn accounts_on_same_provider_track_independently() { + let mut detector = HookTransitionDetector::new(); + let cfg = config(true, None); + let first = lane_key(HookQuotaWindow::Session, Some("a@example.com")); + let second = lane_key(HookQuotaWindow::Session, Some("b@example.com")); + let _ = detector.evaluate( + &observation( + vec![ + lane(Some(50.0), first.clone(), None, &[0.8], false), + lane(Some(50.0), second.clone(), None, &[0.8], false), + ], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![ + lane(Some(85.0), first, None, &[0.8], false), + lane(Some(55.0), second, None, &[0.8], false), + ], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(dispatches.len(), 1); + assert_eq!(dispatches[0].event.event, HookEventType::QuotaLow); +} + +#[test] +fn quota_low_respects_explicit_rule_threshold_over_fallback() { + let mut detector = HookTransitionDetector::new(); + let cfg = config( + true, + Some(vec![rule(HookEventType::QuotaLow, Some(0.9), None)]), + ); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let below = detector.evaluate( + &observation( + vec![lane(Some(85.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(below.is_empty()); + let above = detector.evaluate( + &observation( + vec![lane(Some(95.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert_eq!(events_of(&above), vec![HookEventType::QuotaLow]); +} + +#[test] +fn quota_low_ignores_rules_scoped_to_another_provider() { + let mut detector = HookTransitionDetector::new(); + let cfg = config( + true, + Some(vec![rule(HookEventType::QuotaLow, None, Some("claude"))]), + ); + let key = lane_key(HookQuotaWindow::Session, None); + let _ = detector.evaluate( + &observation( + vec![lane(Some(50.0), key.clone(), None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + let dispatches = detector.evaluate( + &observation( + vec![lane(Some(95.0), key, None, &[0.8], false)], + HookProviderStatus::Unknown, + None, + ), + &cfg, + ); + assert!(dispatches.is_empty()); +} + +#[test] +fn rate_limiter_suppresses_duplicate_refresh_failed() { + use super::super::hooks::HookRateLimiter; + use std::time::Duration; + + let limiter = HookRateLimiter::new(Duration::from_secs(600)); + let event = HookEvent::new(HookEventType::RefreshFailed, "codex").with_status("timeout"); + assert!(limiter.allow(&event, None)); + assert!(!limiter.allow(&event, None)); + + // Quota events are not rate-limited by HookEventType::is_rate_limited. + assert!(!HookEventType::QuotaLow.is_rate_limited()); + assert!(HookEventType::UsageUpdated.is_rate_limited()); + assert!(HookEventType::RefreshFailed.is_rate_limited()); + assert!(HookEventType::ProviderUnavailable.is_rate_limited()); +} diff --git a/rust/src/core/hooks.rs b/rust/src/core/hooks.rs index 6febcea286..116063a5d7 100644 --- a/rust/src/core/hooks.rs +++ b/rust/src/core/hooks.rs @@ -9,7 +9,7 @@ use std::io::Write; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; use std::sync::Mutex; -use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant}; /// Stable event names used in config, env vars, and the JSON payload. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -97,17 +97,10 @@ impl HookEvent { secondary_window_minutes: None, secondary_reset_at: None, status: None, - timestamp: utc_now_iso(), + timestamp: format_utc(chrono::Utc::now()), } } - pub fn with_remaining_percent(mut self, remaining: f64) -> Self { - let remaining = remaining.clamp(0.0, 100.0); - self.remaining_percent = Some(remaining); - self.usage_percent = Some(((100.0 - remaining) / 100.0).clamp(0.0, 1.0)); - self - } - pub fn with_used_percent(mut self, used: f64) -> Self { let used = used.clamp(0.0, 100.0); self.usage_percent = Some(used / 100.0); @@ -163,7 +156,7 @@ impl HookEvent { } pub fn with_timestamp(mut self, ts: chrono::DateTime) -> Self { - self.timestamp = format_unix_utc(ts.timestamp().max(0) as u64); + self.timestamp = format_utc(ts); self } @@ -350,17 +343,6 @@ impl HooksConfig { .filter(|rule| rule.matches(event)) .collect() } - - /// Rules that match the event ignoring the top-level `enabled` flag (for `hooks test`). - pub fn matching_rules_ignoring_master_switch(&self, event: &HookEvent) -> Vec<&HookRule> { - if self.events.len() > Self::MAX_RULES { - return Vec::new(); - } - self.events - .iter() - .filter(|rule| rule.matches(event)) - .collect() - } } /// Env keys forwarded from the parent process (never secrets). @@ -579,54 +561,11 @@ thread_local! { static HOOK_RATE_LIMITER: HookRateLimiter = HookRateLimiter::default(); } -fn utc_now_iso() -> String { - let secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0); - // RFC3339-ish UTC without chrono dependency in this module path. - // Good enough for hooks; refresh path also has chrono elsewhere. - format_unix_utc(secs) -} - -fn format_unix_utc(secs: u64) -> String { - // Manual UTC formatting: YYYY-MM-DDTHH:MM:SSZ - const SECS_PER_DAY: u64 = 86_400; - let days = secs / SECS_PER_DAY; - let tod = secs % SECS_PER_DAY; - let hour = tod / 3600; - let min = (tod % 3600) / 60; - let sec = tod % 60; - #[allow( - clippy::cast_possible_wrap, - reason = "secs comes from duration since the UNIX epoch, so days fits i64 by a wide margin" - )] - let (year, month, day) = civil_from_days(days as i64); - format!("{year:04}-{month:02}-{day:02}T{hour:02}:{min:02}:{sec:02}Z") -} - -/// Howard Hinnant civil_from_days (proleptic Gregorian). -fn civil_from_days(z: i64) -> (i32, u32, u32) { - let z = z + 719_468; - let era = if z >= 0 { z } else { z - 146_096 } / 146_097; - let doe = (z - era * 146_097) as u64; - let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146_096) / 365; - #[allow( - clippy::cast_possible_wrap, - reason = "yoe < 400 by the era math and era * 400 stays far inside i64 for any representable day count" - )] - let y = yoe as i64 + era * 400; - let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); - let mp = (5 * doy + 2) / 153; - let d = doy - (153 * mp + 2) / 5 + 1; - let m = if mp < 10 { mp + 3 } else { mp - 9 }; - let y = if m <= 2 { y + 1 } else { y }; - #[allow( - clippy::cast_possible_truncation, - reason = "Hinnant's algorithm bounds m to 1..=12 and d to 1..=31; y only exceeds i32 for day offsets billions of years from real clock time" - )] - let date = (y as i32, m as u32, d as u32); - date +/// `YYYY-MM-DDTHH:MM:SSZ`, clamped to the epoch for pre-1970 clocks. +fn format_utc(ts: chrono::DateTime) -> String { + ts.max(chrono::DateTime::UNIX_EPOCH) + .format("%Y-%m-%dT%H:%M:%SZ") + .to_string() } fn format_number(value: f64) -> String { @@ -740,6 +679,25 @@ mod tests { use super::*; use std::sync::Arc; + #[test] + fn timestamps_are_utc_seconds_with_z_and_clamp_before_epoch() { + let at = |secs: i64, nanos: u32| { + HookEvent::new(HookEventType::QuotaLow, "codex") + .with_timestamp(chrono::DateTime::from_timestamp(secs, nanos).unwrap()) + .timestamp + }; + assert_eq!(at(0, 0), "1970-01-01T00:00:00Z"); + assert_eq!(at(-5, 0), "1970-01-01T00:00:00Z"); + assert_eq!(at(-1, 500_000_000), "1970-01-01T00:00:00Z"); + assert_eq!(at(951_782_400, 0), "2000-02-29T00:00:00Z"); + assert_eq!(at(1_760_000_000, 999_000_000), "2025-10-09T08:53:20Z"); + assert_eq!(at(4_102_444_799, 0), "2099-12-31T23:59:59Z"); + let now = HookEvent::new(HookEventType::QuotaLow, "codex").timestamp; + assert_eq!(now.len(), 20); + assert!(now.ends_with('Z')); + assert!(chrono::DateTime::parse_from_rfc3339(&now).is_ok()); + } + #[test] fn payload_env_includes_event_and_remaining() { let event = HookEvent::new(HookEventType::QuotaLow, "claude") diff --git a/rust/src/core/http_proxy.rs b/rust/src/core/http_proxy.rs index 445ba503de..6fc8d54cbf 100644 --- a/rust/src/core/http_proxy.rs +++ b/rust/src/core/http_proxy.rs @@ -105,14 +105,6 @@ pub fn apply_app_proxy(builder: ClientBuilder) -> ClientBuilder { apply_proxy_to_builder(builder, &HttpProxySettings::from_app_settings()) } -/// Validate proxy settings for UI (returns `None` when OK / disabled). -pub fn validation_error(settings: &HttpProxySettings) -> Option { - if !settings.enabled { - return None; - } - resolve_proxy(settings).err() -} - #[cfg(test)] mod tests { use super::*; @@ -161,14 +153,8 @@ mod tests { } #[test] - fn validation_error_none_when_disabled() { - let s = HttpProxySettings::from_parts(false, "bad", "", ""); - assert!(validation_error(&s).is_none()); - } - - #[test] - fn validation_error_when_enabled_and_bad() { + fn resolve_proxy_errors_when_enabled_and_bad() { let s = HttpProxySettings::from_parts(true, "ftp://x", "", ""); - assert!(validation_error(&s).is_some()); + assert!(resolve_proxy(&s).is_err()); } } diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index 8f0ae2bdca..1aec9c7db3 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -26,13 +26,11 @@ mod quota_burndown; mod rate_window; mod redactor; mod session_equivalent_forecast; -mod session_quota; mod sqlite; mod timezone; mod token_accounts; mod usage_pace; mod usage_snapshot; -mod widget_snapshot; pub use adaptive_refresh::*; pub use aws_signing::*; @@ -61,10 +59,8 @@ pub use quota_burndown::{ pub use rate_window::*; pub use redactor::*; pub use session_equivalent_forecast::*; -pub use session_quota::*; pub use sqlite::*; pub use timezone::{local_timezone_name, try_local_timezone_name}; pub use token_accounts::*; pub use usage_pace::*; pub use usage_snapshot::*; -pub use widget_snapshot::*; diff --git a/rust/src/core/openai_dashboard.rs b/rust/src/core/openai_dashboard.rs index b2d1d88372..33e3491368 100755 --- a/rust/src/core/openai_dashboard.rs +++ b/rust/src/core/openai_dashboard.rs @@ -2,14 +2,6 @@ //! //! Data structures for OpenAI/Codex dashboard usage breakdown and credits tracking. -// Serde data model mirroring the OpenAI dashboard payload plus its cache -// store; several fields and helpers have no consumer yet but must survive -// (de)serialization round-trips. -#![allow( - dead_code, - reason = "dashboard snapshot fields mirror the upstream OpenAI payload; not all are consumed yet" -)] - use chrono::{DateTime, NaiveDate, Utc}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -124,36 +116,6 @@ impl OpenAIDashboardSnapshot { .collect() } - /// Set usage breakdown data - pub fn with_usage_breakdown(mut self, breakdown: Vec) -> Self { - self.usage_breakdown = breakdown; - self - } - - /// Set primary limit - pub fn with_primary_limit(mut self, limit: RateWindow) -> Self { - self.primary_limit = Some(limit); - self - } - - /// Set secondary limit - pub fn with_secondary_limit(mut self, limit: RateWindow) -> Self { - self.secondary_limit = Some(limit); - self - } - - /// Set credits remaining - pub fn with_credits_remaining(mut self, credits: f64) -> Self { - self.credits_remaining = Some(credits); - self - } - - /// Set account plan - pub fn with_account_plan(mut self, plan: impl Into) -> Self { - self.account_plan = Some(plan.into()); - self - } - /// Set explicitly observed subscription lifecycle dates. pub fn with_subscription(mut self, subscription: Option) -> Self { self.subscription = subscription; @@ -183,11 +145,6 @@ impl CreditEvent { description: None, } } - - pub fn with_description(mut self, description: impl Into) -> Self { - self.description = Some(description.into()); - self - } } /// Daily usage breakdown diff --git a/rust/src/core/provider_state.rs b/rust/src/core/provider_state.rs index d9e9c5f99c..4bcc593ec0 100644 --- a/rust/src/core/provider_state.rs +++ b/rust/src/core/provider_state.rs @@ -39,12 +39,6 @@ pub enum ProviderStateKind { } impl ProviderStateKind { - /// Whether this state should mark the provider as having a problem in - /// presentation surfaces (tone changes, hides usage details). - pub const fn is_problem(self) -> bool { - !matches!(self, ProviderStateKind::Ready) - } - /// Whether this state means the account must sign in again. The only /// states that qualify for credential-expiry alerts; quota, permission, /// rate-limit and transport failures map elsewhere and never do. @@ -143,19 +137,6 @@ mod tests { assert_eq!(E::Other("API error 500".into()).state_kind(), K::Unknown); } - #[test] - fn ready_is_the_only_non_problem_state() { - assert!(!ProviderStateKind::Ready.is_problem()); - for kind in [ - ProviderStateKind::NeedsAuthentication, - ProviderStateKind::ExpiredSession, - ProviderStateKind::LocalRuntimeOffline, - ProviderStateKind::Unknown, - ] { - assert!(kind.is_problem()); - } - } - #[test] fn only_authentication_states_need_sign_in() { assert!(ProviderStateKind::NeedsAuthentication.needs_sign_in()); diff --git a/rust/src/core/quota_burndown.rs b/rust/src/core/quota_burndown.rs index ee92b5a659..5f3fca90e8 100644 --- a/rust/src/core/quota_burndown.rs +++ b/rust/src/core/quota_burndown.rs @@ -40,8 +40,6 @@ pub const MAX_SERIES_SAMPLES: usize = 17_520; pub enum QuotaBurndownStoreError { #[error("quota burndown account key is empty")] EmptyAccountKey, - #[error("quota burndown account key does not match the requested partition")] - AccountKeyMismatch, #[error("failed to read quota burndown history: {0}")] Read(#[source] std::io::Error), #[error("failed to decode quota burndown history: {0}")] diff --git a/rust/src/core/rate_window.rs b/rust/src/core/rate_window.rs index cf250ef73d..464a992651 100755 --- a/rust/src/core/rate_window.rs +++ b/rust/src/core/rate_window.rs @@ -217,11 +217,6 @@ impl RateWindow { self.used_percent >= 100.0 } - /// Check if the window is nearly exhausted (>= 90% used) - pub fn is_nearly_exhausted(&self) -> bool { - self.used_percent >= 90.0 - } - /// Format the reset time as a countdown string pub fn format_countdown(&self) -> Option { Some(format_countdown_until(self.resets_at?, Utc::now())) diff --git a/rust/src/core/redactor.rs b/rust/src/core/redactor.rs index 03858bc066..29d617420b 100755 --- a/rust/src/core/redactor.rs +++ b/rust/src/core/redactor.rs @@ -91,27 +91,6 @@ impl SecretRedactor { pub struct PersonalInfoRedactor; impl PersonalInfoRedactor { - /// Redact a single email address if privacy mode is enabled - /// - /// # Arguments - /// * `email` - The email address to potentially redact - /// * `is_enabled` - Whether privacy/redaction mode is enabled - /// - /// # Returns - /// The original email if disabled, or "Hidden" if enabled - pub fn redact_email(email: Option<&str>, is_enabled: bool) -> String { - match email { - Some(e) if !e.trim().is_empty() => { - if is_enabled { - EMAIL_PLACEHOLDER.to_string() - } else { - e.to_string() - } - } - _ => String::new(), - } - } - /// Redact all email addresses in a text string /// /// # Arguments @@ -161,31 +140,6 @@ impl PersonalInfoRedactor { mod tests { use super::*; - #[test] - fn test_redact_email_disabled() { - let email = "test@example.com"; - assert_eq!( - PersonalInfoRedactor::redact_email(Some(email), false), - email - ); - } - - #[test] - fn test_redact_email_enabled() { - let email = "test@example.com"; - assert_eq!( - PersonalInfoRedactor::redact_email(Some(email), true), - EMAIL_PLACEHOLDER - ); - } - - #[test] - fn test_redact_email_none() { - assert_eq!(PersonalInfoRedactor::redact_email(None, true), ""); - assert_eq!(PersonalInfoRedactor::redact_email(Some(""), true), ""); - assert_eq!(PersonalInfoRedactor::redact_email(Some(" "), true), ""); - } - #[test] fn test_redact_emails_in_text() { let text = "Contact me at user@example.com or admin@test.org for help"; diff --git a/rust/src/core/session_equivalent_forecast.rs b/rust/src/core/session_equivalent_forecast.rs index 6df3233a69..8f3fa5b79d 100644 --- a/rust/src/core/session_equivalent_forecast.rs +++ b/rust/src/core/session_equivalent_forecast.rs @@ -55,14 +55,7 @@ pub enum PlanUtilizationSeriesName { Weekly, } -impl PlanUtilizationSeriesName { - pub fn canonical_window_minutes(self) -> u32 { - match self { - Self::Session => SESSION_WINDOW_MINUTES, - Self::Weekly => WEEKLY_WINDOW_MINUTES, - } - } -} +impl PlanUtilizationSeriesName {} /// Chronological history for one series. #[derive(Debug, Clone, PartialEq)] @@ -516,866 +509,8 @@ fn is_workday(date: DateTime, work_days: u8) -> bool { iso <= work_days } -// ── In-process history store ───────────────────────────────────────── - -/// Identity that forecast history is scoped to. -/// -/// History must never be shared across accounts. Plan sizes differ, so blending -/// observations from two accounts yields a silently wrong median — see the account -/// isolation regression test below. -/// -/// ponytail: when this is persisted to disk, hash `account_key` rather than writing the -/// raw address; in-process it stays plain because the email is already resident in -/// `UsageSnapshot`. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct ForecastScope { - pub provider_id: String, - pub account_key: Option, -} - -impl ForecastScope { - pub fn new(provider_id: &str, account_key: Option<&str>) -> Self { - Self { - provider_id: provider_id.to_string(), - account_key: account_key.map(str::to_string), - } - } -} - -/// Append-only ring of session/weekly observations for burn estimation. -/// -/// ponytail: history is process-local and lost on restart; upgrade path = disk cache -/// keyed by provider+account once forecast quality justifies persistence. -#[derive(Debug, Default)] -pub struct SessionEquivalentHistoryStore { - by_scope: HashMap, -} - -#[derive(Debug, Default)] -struct ProviderHistory { - session: Vec, - weekly: Vec, -} - -static HISTORY_STORE: LazyLock> = - LazyLock::new(|| Mutex::new(SessionEquivalentHistoryStore::default())); - -impl SessionEquivalentHistoryStore { - /// Record one observation pair for a provider+account (session + optional weekly). - pub fn record( - &mut self, - scope: &ForecastScope, - session: Option, - weekly: Option, - sample_limit: usize, - ) { - let limit = sample_limit.max(1); - // Keep more raw points than completed groups so grouping still has density. - let ring = limit.saturating_mul(8).max(24); - let hist = self.by_scope.entry(scope.clone()).or_default(); - if let Some(entry) = session { - push_ring(&mut hist.session, entry, ring); - } - if let Some(entry) = weekly { - push_ring(&mut hist.weekly, entry, ring); - } - } - - pub fn histories(&self, scope: &ForecastScope) -> Vec { - let Some(hist) = self.by_scope.get(scope) else { - return Vec::new(); - }; - let mut out = Vec::new(); - if !hist.session.is_empty() { - out.push(PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Session, - window_minutes: SESSION_WINDOW_MINUTES, - entries: hist.session.clone(), - }); - } - if !hist.weekly.is_empty() { - out.push(PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Weekly, - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: hist.weekly.clone(), - }); - } - out - } - - /// Merge persisted series into the in-process store for `scope`. Older - /// in-process observations win ties; the result stays sorted so the - /// estimator's chronological check keeps passing. - pub fn adopt_persisted( - &mut self, - scope: &ForecastScope, - persisted: &[quota_burndown::PersistedPlanSeries], - ) { - for series in persisted { - let entries: Vec = series - .entries - .iter() - .filter_map(PersistedEntry::to_history) - .collect(); - if entries.is_empty() { - continue; - } - let hist = self.by_scope.entry(scope.clone()).or_default(); - let ring = quota_burndown::MAX_SERIES_SAMPLES; - match series.name.as_str() { - "session" => { - merge_ring(&mut hist.session, entries, ring); - } - "weekly" => { - merge_ring(&mut hist.weekly, entries, ring); - } - _ => {} - } - } - } -} - -fn push_ring( - buf: &mut Vec, - entry: PlanUtilizationHistoryEntry, - ring: usize, -) { - buf.push(entry); - if buf.len() > ring { - let drop = buf.len() - ring; - buf.drain(0..drop); - } -} - -/// Merge loaded entries into a ring buffer: append, sort, dedup, cap. -fn merge_ring( - buf: &mut Vec, - entries: Vec, - ring: usize, -) { - buf.extend(entries); - buf.sort_by(|left, right| { - left.captured_at - .cmp(&right.captured_at) - .then_with(|| left.used_percent.total_cmp(&right.used_percent)) - }); - buf.dedup(); - if buf.len() > ring { - let drop = buf.len() - ring; - buf.drain(0..drop); - } -} - -/// Load the persisted burndown history for one provider+account into the -/// process-local store, so history survives a restart. Runs once per scope; -/// later calls are cheap no-ops (the loaded marker is kept in the store). -pub fn load_persisted_history( - provider_id: &str, - account_key: Option<&str>, -) -> Vec { - static LOADED: LazyLock>> = - LazyLock::new(|| Mutex::new(HashSet::new())); - let scope = ForecastScope::new(provider_id, account_key); - if let Ok(guard) = LOADED.lock() - && guard.contains(&scope) - { - return global_history_store() - .lock() - .ok() - .map(|store| store.histories(&scope)) - .unwrap_or_default(); - } - let series = quota_burndown::load_persisted_series(provider_id, account_key); - if let Ok(mut loaded) = LOADED.lock() { - loaded.insert(scope.clone()); - } - if let Ok(mut store) = global_history_store().lock() { - store.adopt_persisted(&scope, &series); - } - global_history_store() - .lock() - .ok() - .map(|store| store.histories(&scope)) - .unwrap_or_default() -} - -/// Global in-process history used by Claude/Codex snapshot refresh. -pub fn global_history_store() -> &'static Mutex { - &HISTORY_STORE -} - -/// Record session/weekly windows from a live provider usage snapshot. -/// -/// `account_key` scopes the history so switching accounts on one provider does not -/// blend burn observations across plans. -pub fn record_provider_windows( - provider_id: &str, - account_key: Option<&str>, - session: &crate::core::RateWindow, - weekly: Option<&crate::core::RateWindow>, - now: DateTime, -) { - if session.is_informational - || session.window_minutes != Some(SESSION_WINDOW_MINUTES) - || !session.used_percent.is_finite() - { - return; - } - let session_entry = PlanUtilizationHistoryEntry { - captured_at: now, - used_percent: session.used_percent.clamp(0.0, 100.0), - resets_at: session.resets_at, - }; - let weekly_entry = weekly.and_then(|w| { - if w.is_informational - || w.window_minutes != Some(WEEKLY_WINDOW_MINUTES) - || !w.used_percent.is_finite() - { - return None; - } - Some(PlanUtilizationHistoryEntry { - captured_at: now, - used_percent: w.used_percent.clamp(0.0, 100.0), - resets_at: w.resets_at, - }) - }); - - if let Ok(mut guard) = global_history_store().lock() { - guard.record( - &ForecastScope::new(provider_id, account_key), - Some(session_entry.clone()), - weekly_entry.clone(), - SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, - ); - } - - // Best-effort persistence for the burndown chart (upstream #4085 keeps a - // durable store; a failed write only costs the chart, never the forecast). - if let Err(error) = quota_burndown::persist_recorded_windows( - provider_id, - account_key, - Some(&session_entry), - weekly_entry.as_ref(), - now, - ) { - tracing::debug!(%error, "quota burndown persistence failed"); - } -} - -/// Best-effort persistence of one recorded observation pair for the burndown -/// chart. Failures are logged at debug level and never surface. -pub fn persist_recorded_windows( - provider_id: &str, - account_key: Option<&str>, - session: Option<&PlanUtilizationHistoryEntry>, - weekly: Option<&PlanUtilizationHistoryEntry>, - _now: DateTime, -) -> Result<(), quota_burndown::QuotaBurndownStoreError> { - let config_root = match crate::logging::config_root() { - Some(root) => root, - None => return Ok(()), - }; - let mut session_series = quota_burndown::PersistedPlanSeries { - name: "session".to_string(), - window_minutes: SESSION_WINDOW_MINUTES, - entries: Vec::new(), - }; - if let Some(entry) = session { - session_series - .entries - .push(quota_burndown::PersistedPlanEntry::from_history(entry)); - } - let mut weekly_series = quota_burndown::PersistedPlanSeries { - name: "weekly".to_string(), - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: Vec::new(), - }; - if let Some(entry) = weekly { - weekly_series - .entries - .push(quota_burndown::PersistedPlanEntry::from_history(entry)); - } - let key = quota_burndown::persisted_account_key(account_key); - if key.is_empty() { - return Ok(()); - } - quota_burndown::merge_and_persist_series( - &config_root, - provider_id, - &key, - vec![session_series, weekly_series], - ) - .map(|_| ()) -} - -/// Last learned full-session burn estimate retained across idle refreshes. -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct RetainedFullSessionEstimate { - pub estimate: f64, - pub updated_at: DateTime, -} - -/// Keep the last positive full-session estimate when a refresh yields none. -/// -/// `fresh` wins when present and valid; otherwise `previous` is kept so the UI -/// does not blank while the session window is idle (upstream #2336). -pub fn retain_last_full_session_estimate( - previous: Option, - fresh: Option, - now: DateTime, -) -> Option { - match fresh { - Some(estimate) if estimate.is_finite() && estimate > 0.0 => { - Some(RetainedFullSessionEstimate { - estimate, - updated_at: now, - }) - } - _ => previous, - } -} - -// ponytail: in-memory only, lost on restart; upgrade = persist to cache file. -#[derive(Debug, Default)] -struct LastFullSessionEstimateStore { - by_scope: HashMap, -} - -static LAST_FULL_SESSION_ESTIMATE_STORE: LazyLock> = - LazyLock::new(|| Mutex::new(LastFullSessionEstimateStore::default())); - -fn last_full_session_estimate_store() -> &'static Mutex { - &LAST_FULL_SESSION_ESTIMATE_STORE -} - -/// Remember/recall the last learned full-session burn estimate for a provider+account. -pub fn remember_full_session_estimate( - scope: &ForecastScope, - fresh: Option, - now: DateTime, -) -> Option { - let Ok(mut guard) = last_full_session_estimate_store().lock() else { - return fresh.filter(|v| v.is_finite() && *v > 0.0); - }; - let previous = guard.by_scope.get(scope).copied(); - let retained = retain_last_full_session_estimate(previous, fresh, now); - if let Some(entry) = retained { - guard.by_scope.insert(scope.clone(), entry); - Some(entry.estimate) - } else { - guard.by_scope.remove(scope); - None - } -} - -/// Compute forecast for a provider+account using the in-process history ring. -pub fn forecast_for_provider( - provider_id: &str, - account_key: Option<&str>, - session: &crate::core::RateWindow, - weekly: &crate::core::RateWindow, - now: DateTime, - work_days: Option, -) -> Option { - let scope = ForecastScope::new(provider_id, account_key); - let histories = global_history_store() - .lock() - .ok() - .map(|g| g.histories(&scope)) - .unwrap_or_default(); - let fresh_burn = SessionEquivalentBurnEstimator::estimate( - &histories, - session.resets_at, - now, - SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, - ); - let fresh_median = fresh_burn - .as_ref() - .map(|b| b.median_weekly_percent_per_window); - let sample_count = fresh_burn - .as_ref() - .map(|b| b.sample_count) - .unwrap_or(SessionEquivalentBurnEstimator::MINIMUM_SAMPLE_COUNT); - let median = remember_full_session_estimate(&scope, fresh_median, now)?; - let burn = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: median, - sample_count, - }; - SessionEquivalentForecast::make(session, weekly, &burn, None, now, work_days) -} +mod history; +pub use history::*; #[cfg(test)] -mod tests { - use super::*; - use crate::core::RateWindow; - use chrono::TimeZone; - - fn ts(secs: i64) -> DateTime { - Utc.timestamp_opt(secs, 0).single().unwrap() - } - - fn entry(captured: i64, used: f64, resets: i64) -> PlanUtilizationHistoryEntry { - PlanUtilizationHistoryEntry { - captured_at: ts(captured), - used_percent: used, - resets_at: Some(ts(resets)), - } - } - - fn session_secs() -> i64 { - i64::from(SESSION_WINDOW_MINUTES) * 60 - } - - fn weekly_secs() -> i64 { - i64::from(WEEKLY_WINDOW_MINUTES) * 60 - } - - /// Build three completed full-burn sessions with aligned weekly observations. - fn three_sample_histories( - base: i64, - ) -> ( - Vec, - DateTime, - DateTime, - ) { - let s = session_secs(); - let w = weekly_secs(); - let weekly_reset = base + w; - - let mut session_entries = Vec::new(); - let mut weekly_entries = Vec::new(); - - // Three completed sessions ending at base+s, base+2s, base+3s. - for i in 0..3 { - let start = base + i * s; - let end = start + s; - // Window-start weekly anchor (session used 0). - weekly_entries.push(entry(start - 30, 10.0 + i as f64 * 5.0, weekly_reset)); - // Mid / end session captures with exact weekly alignment. - session_entries.push(entry(start + 60, 40.0, end)); - weekly_entries.push(entry(start + 60, 12.0 + i as f64 * 5.0, weekly_reset)); - session_entries.push(entry(end - 60, 100.0, end)); - weekly_entries.push(entry(end - 60, 15.0 + i as f64 * 5.0, weekly_reset)); - // End-of-window weekly for max>=100 path. - weekly_entries.push(entry(end, 15.5 + i as f64 * 5.0, weekly_reset)); - } - - // Current open session. - let current_start = base + 3 * s; - let current_end = current_start + s; - let now_secs = current_start + 30 * 60; - session_entries.push(entry(now_secs, 20.0, current_end)); - weekly_entries.push(entry(now_secs, 30.0, weekly_reset)); - - session_entries.sort_by_key(|e| e.captured_at); - weekly_entries.sort_by_key(|e| e.captured_at); - - let histories = vec![ - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Session, - window_minutes: SESSION_WINDOW_MINUTES, - entries: session_entries, - }, - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Weekly, - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: weekly_entries, - }, - ]; - (histories, ts(now_secs), ts(current_end)) - } - - #[test] - fn median_math_from_three_full_sessions() { - let base = 1_700_000_000_i64; - let (histories, now, current_end) = three_sample_histories(base); - let estimate = SessionEquivalentBurnEstimator::estimate( - &histories, - Some(current_end), - now, - SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, - ) - .expect("estimate"); - assert_eq!(estimate.sample_count, 3); - // Mid-session captures dominate when start/end anchors are optional: - // weekly 12→15 over session 40→100 ⇒ 100 * 3 / 60 = 5.0 - // (window-start + end anchors yield the same full-allowance burn here). - assert!( - (estimate.median_weekly_percent_per_window - 5.0).abs() < 0.01, - "median={}", - estimate.median_weekly_percent_per_window - ); - } - - #[test] - fn grouping_tolerance_merges_nearby_resets() { - let s = session_secs(); - let base = 1_700_100_000_i64; - let weekly_reset = base + weekly_secs(); - let end = base + s; - - // Nearby resets within ±120s must merge into one group (not out-of-order). - let single_group = vec![ - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Session, - window_minutes: SESSION_WINDOW_MINUTES, - entries: vec![ - entry(base + 60, 30.0, end), - entry(base + 120, 60.0, end + 100), - entry(base + 180, 90.0, end + 110), - ], - }, - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Weekly, - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: vec![ - entry(base - 30, 5.0, weekly_reset), - entry(base + 60, 8.0, weekly_reset), - entry(base + 120, 10.0, weekly_reset), - entry(base + 180, 12.0, weekly_reset), - ], - }, - ]; - // One completed group → below minimum sample count, but not an ordering reject. - assert!( - SessionEquivalentBurnEstimator::estimate( - &single_group, - Some(ts(end + s)), - ts(end + 60), - 7 - ) - .is_none() - ); - - // Strictly decreasing resets_at across groups is rejected. - let out_of_order = vec![ - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Session, - window_minutes: SESSION_WINDOW_MINUTES, - entries: vec![ - entry(base + 60, 30.0, end + s), - entry(base + 120, 60.0, end), // jumps backward beyond tolerance - ], - }, - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Weekly, - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: vec![ - entry(base + 60, 8.0, weekly_reset), - entry(base + 120, 10.0, weekly_reset), - ], - }, - ]; - assert!( - SessionEquivalentBurnEstimator::estimate( - &out_of_order, - Some(ts(end + 2 * s)), - ts(end + s + 60), - 7 - ) - .is_none() - ); - } - - #[test] - fn insufficient_samples_returns_none() { - let base = 1_700_200_000_i64; - let s = session_secs(); - let weekly_reset = base + weekly_secs(); - let end = base + s; - let histories = vec![ - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Session, - window_minutes: SESSION_WINDOW_MINUTES, - entries: vec![entry(base + 60, 50.0, end), entry(base + 120, 80.0, end)], - }, - PlanUtilizationSeriesHistory { - name: PlanUtilizationSeriesName::Weekly, - window_minutes: WEEKLY_WINDOW_MINUTES, - entries: vec![ - entry(base - 30, 10.0, weekly_reset), - entry(base + 60, 12.0, weekly_reset), - entry(base + 120, 14.0, weekly_reset), - ], - }, - ]; - assert!( - SessionEquivalentBurnEstimator::estimate( - &histories, - Some(ts(end + s)), - ts(end + 60), - 7 - ) - .is_none() - ); - } - - #[test] - fn forecast_make_fractional_windows() { - let now = ts(1_700_300_000); - let session = RateWindow::with_details( - 25.0, - Some(SESSION_WINDOW_MINUTES), - Some(now + Duration::hours(3)), - None, - ); - let weekly = RateWindow::with_details( - 40.0, - Some(WEEKLY_WINDOW_MINUTES), - Some(now + Duration::days(3)), - None, - ); - let burn = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: 12.0, - sample_count: 3, - }; - let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) - .expect("forecast"); - // remaining weekly 60 / 12 = 5.0 - assert!((forecast.estimated_windows_to_exhaust_weekly - 5.0).abs() < 1e-9); - assert!(forecast.available_windows_until_reset > 0.0); - // Mirrors the production whole-window truncation. - #[allow( - clippy::cast_possible_truncation, - reason = "assertion mirrors the production whole-window truncation" - )] - let expected_windows = forecast.available_windows_until_reset.floor() as i64; - assert_eq!(forecast.windows_until_reset, expected_windows); - assert_eq!(forecast.sample_count, 3); - assert!((forecast.weekly_used_percent - 40.0).abs() < 1e-9); - } - - #[test] - fn work_days_reduces_available_windows_vs_wall_clock() { - // Friday 18:00 UTC → Monday 18:00 UTC (weekend in between). - let friday = Utc - .with_ymd_and_hms(2024, 6, 14, 18, 0, 0) - .single() - .unwrap(); - assert_eq!(friday.weekday(), Weekday::Fri); - let monday = Utc - .with_ymd_and_hms(2024, 6, 17, 18, 0, 0) - .single() - .unwrap(); - - let wall = effective_remaining_seconds(friday, monday, None); - let work5 = effective_remaining_seconds(friday, monday, Some(5)); - assert!((wall - 3.0 * 24.0 * 3600.0).abs() < 1.0); - // Fri 18→24 = 6h, Mon 0→18 = 18h → 24h work seconds (no Sat/Sun). - assert!((work5 - 24.0 * 3600.0).abs() < 1.0); - assert!(work5 < wall); - - let session = RateWindow::with_details( - 10.0, - Some(SESSION_WINDOW_MINUTES), - Some(friday + Duration::hours(4)), - None, - ); - let weekly = - RateWindow::with_details(20.0, Some(WEEKLY_WINDOW_MINUTES), Some(monday), None); - let burn = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: 10.0, - sample_count: 4, - }; - let wall_f = - SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, None).unwrap(); - let work_f = - SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, Some(5)) - .unwrap(); - assert!(work_f.available_windows_until_reset < wall_f.available_windows_until_reset); - // Mirrors the production whole-window truncation. - #[allow( - clippy::cast_possible_truncation, - reason = "assertion mirrors the production whole-window truncation" - )] - let expected_work_windows = work_f.available_windows_until_reset.floor() as i64; - assert_eq!(work_f.windows_until_reset, expected_work_windows); - } - - #[test] - fn work_days_out_of_range_uses_wall_clock() { - let now = ts(1_700_400_000); - let reset = now + Duration::days(2); - let wall = effective_remaining_seconds(now, reset, None); - assert!((effective_remaining_seconds(now, reset, Some(1)) - wall).abs() < 1e-9); - assert!((effective_remaining_seconds(now, reset, Some(7)) - wall).abs() < 1e-9); - } - - #[test] - fn history_ring_retains_latest_samples() { - let mut store = SessionEquivalentHistoryStore::default(); - let scope = ForecastScope::new("claude", None); - let base = ts(1_700_500_000); - for i in 0..30 { - store.record( - &scope, - Some(PlanUtilizationHistoryEntry { - captured_at: base + Duration::minutes(i), - used_percent: i as f64, - resets_at: Some(base + Duration::hours(5)), - }), - None, - 7, - ); - } - let h = store.histories(&scope); - assert_eq!(h.len(), 1); - assert!(h[0].entries.len() <= 7 * 8); - assert!((h[0].entries.last().unwrap().used_percent - 29.0).abs() < 1e-9); - } - - /// Two accounts on one provider must not share burn samples. - /// - /// Regression: history was keyed by `provider_id` alone, so switching the active - /// account blended both accounts' observations into one ring and produced a median - /// drawn from a mixture of plans. - #[test] - fn history_is_isolated_per_account() { - let mut store = SessionEquivalentHistoryStore::default(); - let base = ts(1_700_500_000); - let alice = ForecastScope::new("codex", Some("alice@example.com")); - let bob = ForecastScope::new("codex", Some("bob@example.com")); - - for i in 0..5 { - store.record( - &alice, - Some(PlanUtilizationHistoryEntry { - captured_at: base + Duration::minutes(i), - used_percent: 10.0, - resets_at: Some(base + Duration::hours(5)), - }), - None, - 7, - ); - } - store.record( - &bob, - Some(PlanUtilizationHistoryEntry { - captured_at: base + Duration::minutes(99), - used_percent: 90.0, - resets_at: Some(base + Duration::hours(5)), - }), - None, - 7, - ); - - let alice_hist = store.histories(&alice); - let bob_hist = store.histories(&bob); - assert_eq!(alice_hist[0].entries.len(), 5); - assert_eq!(bob_hist[0].entries.len(), 1); - assert!( - alice_hist[0] - .entries - .iter() - .all(|e| (e.used_percent - 10.0).abs() < 1e-9), - "bob's 90% sample leaked into alice's history" - ); - assert!((bob_hist[0].entries[0].used_percent - 90.0).abs() < 1e-9); - - // A provider with no account discriminator is its own bucket, not a catch-all. - assert!( - store - .histories(&ForecastScope::new("codex", None)) - .is_empty() - ); - } - - #[test] - fn make_rejects_bad_windows_and_zero_remaining() { - let now = ts(1_700_600_000); - let burn = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: 10.0, - sample_count: 3, - }; - let session = RateWindow::with_details( - 10.0, - Some(SESSION_WINDOW_MINUTES), - Some(now + Duration::hours(2)), - None, - ); - let exhausted = RateWindow::with_details( - 100.0, - Some(WEEKLY_WINDOW_MINUTES), - Some(now + Duration::days(2)), - None, - ); - assert!( - SessionEquivalentForecast::make(&session, &exhausted, &burn, None, now, None).is_none() - ); - - let low_samples = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: 10.0, - sample_count: 2, - }; - let weekly = RateWindow::with_details( - 50.0, - Some(WEEKLY_WINDOW_MINUTES), - Some(now + Duration::days(2)), - None, - ); - assert!( - SessionEquivalentForecast::make(&session, &weekly, &low_samples, None, now, None) - .is_none() - ); - } - - #[test] - fn retain_last_full_session_estimate_survives_idle_refresh() { - let t0 = ts(1_700_000_000); - let t1 = t0 + Duration::hours(1); - let learned = retain_last_full_session_estimate(None, Some(12.5), t0) - .expect("fresh estimate should stick"); - assert_eq!(learned.estimate, 12.5); - assert_eq!(learned.updated_at, t0); - - // Idle refresh yields no fresh sample — keep the last learned value. - let idle = retain_last_full_session_estimate(Some(learned), None, t1) - .expect("idle refresh must keep last estimate"); - assert_eq!(idle.estimate, 12.5); - assert_eq!(idle.updated_at, t0); - - // Non-positive / non-finite fresh values also keep previous. - let still = retain_last_full_session_estimate(Some(idle), Some(0.0), t1) - .expect("zero fresh must not clear"); - assert_eq!(still.estimate, 12.5); - let still = retain_last_full_session_estimate(Some(still), Some(f64::NAN), t1) - .expect("nan fresh must not clear"); - assert_eq!(still.estimate, 12.5); - } - - #[test] - fn retain_last_full_session_estimate_replaced_by_fresh_sample() { - let t0 = ts(1_700_000_000); - let t1 = t0 + Duration::hours(2); - let previous = retain_last_full_session_estimate(None, Some(8.0), t0).unwrap(); - let next = retain_last_full_session_estimate(Some(previous), Some(15.0), t1) - .expect("fresh sample replaces previous"); - assert_eq!(next.estimate, 15.0); - assert_eq!(next.updated_at, t1); - } - - #[test] - fn make_keeps_forecast_when_session_reset_is_expired_idle() { - let now = ts(1_700_600_000); - let burn = SessionEquivalentBurnEstimate { - median_weekly_percent_per_window: 10.0, - sample_count: 3, - }; - // Session reset already passed (idle) — forecast still derives from weekly. - let session = RateWindow::with_details( - 0.0, - Some(SESSION_WINDOW_MINUTES), - Some(now - Duration::hours(1)), - None, - ); - let weekly = RateWindow::with_details( - 40.0, - Some(WEEKLY_WINDOW_MINUTES), - Some(now + Duration::days(3)), - None, - ); - let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) - .expect("idle expired session should still forecast"); - assert!((forecast.estimated_windows_to_exhaust_weekly - 6.0).abs() < 1e-9); - } -} +mod tests; diff --git a/rust/src/core/session_equivalent_forecast/history.rs b/rust/src/core/session_equivalent_forecast/history.rs new file mode 100644 index 0000000000..09dafab46d --- /dev/null +++ b/rust/src/core/session_equivalent_forecast/history.rs @@ -0,0 +1,347 @@ +//! In-process history store plus the last-full-session estimate cache. + +use super::*; + +/// Identity that forecast history is scoped to. +/// +/// History must never be shared across accounts. Plan sizes differ, so blending +/// observations from two accounts yields a silently wrong median — see the account +/// isolation regression test below. +/// +/// ponytail: when this is persisted to disk, hash `account_key` rather than writing the +/// raw address; in-process it stays plain because the email is already resident in +/// `UsageSnapshot`. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ForecastScope { + pub provider_id: String, + pub account_key: Option, +} + +impl ForecastScope { + pub fn new(provider_id: &str, account_key: Option<&str>) -> Self { + Self { + provider_id: provider_id.to_string(), + account_key: account_key.map(str::to_string), + } + } +} + +/// Append-only ring of session/weekly observations for burn estimation. +/// +/// ponytail: history is process-local and lost on restart; upgrade path = disk cache +/// keyed by provider+account once forecast quality justifies persistence. +#[derive(Debug, Default)] +pub struct SessionEquivalentHistoryStore { + by_scope: HashMap, +} + +/// The tracked series, in output order: (name, persisted key, window minutes). +const SERIES: [(PlanUtilizationSeriesName, &str, u32); 2] = [ + ( + PlanUtilizationSeriesName::Session, + "session", + SESSION_WINDOW_MINUTES, + ), + ( + PlanUtilizationSeriesName::Weekly, + "weekly", + WEEKLY_WINDOW_MINUTES, + ), +]; + +/// One ring per `SERIES` entry, same index. +type ProviderHistory = [Vec; SERIES.len()]; + +static HISTORY_STORE: LazyLock> = + LazyLock::new(|| Mutex::new(SessionEquivalentHistoryStore::default())); + +impl SessionEquivalentHistoryStore { + /// Record one observation pair for a provider+account (session + optional weekly). + pub fn record( + &mut self, + scope: &ForecastScope, + session: Option, + weekly: Option, + sample_limit: usize, + ) { + let limit = sample_limit.max(1); + // Keep more raw points than completed groups so grouping still has density. + let ring = limit.saturating_mul(8).max(24); + let hist = self.by_scope.entry(scope.clone()).or_default(); + for (buf, entry) in hist.iter_mut().zip([session, weekly]) { + if let Some(entry) = entry { + push_ring(buf, entry, ring); + } + } + } + + pub fn histories(&self, scope: &ForecastScope) -> Vec { + let Some(hist) = self.by_scope.get(scope) else { + return Vec::new(); + }; + SERIES + .iter() + .zip(hist) + .filter(|(_, entries)| !entries.is_empty()) + .map( + |(&(name, _, window_minutes), entries)| PlanUtilizationSeriesHistory { + name, + window_minutes, + entries: entries.clone(), + }, + ) + .collect() + } + + /// Merge persisted series into the in-process store for `scope`. Older + /// in-process observations win ties; the result stays sorted so the + /// estimator's chronological check keeps passing. + pub fn adopt_persisted( + &mut self, + scope: &ForecastScope, + persisted: &[quota_burndown::PersistedPlanSeries], + ) { + for series in persisted { + let entries: Vec = series + .entries + .iter() + .filter_map(PersistedEntry::to_history) + .collect(); + if entries.is_empty() { + continue; + } + let Some(index) = SERIES.iter().position(|(_, key, _)| *key == series.name) else { + continue; + }; + let hist = self.by_scope.entry(scope.clone()).or_default(); + merge_ring( + &mut hist[index], + entries, + quota_burndown::MAX_SERIES_SAMPLES, + ); + } + } +} + +fn push_ring( + buf: &mut Vec, + entry: PlanUtilizationHistoryEntry, + ring: usize, +) { + buf.push(entry); + if buf.len() > ring { + let drop = buf.len() - ring; + buf.drain(0..drop); + } +} + +/// Merge loaded entries into a ring buffer: append, sort, dedup, cap. +fn merge_ring( + buf: &mut Vec, + entries: Vec, + ring: usize, +) { + buf.extend(entries); + buf.sort_by(|left, right| { + left.captured_at + .cmp(&right.captured_at) + .then_with(|| left.used_percent.total_cmp(&right.used_percent)) + }); + buf.dedup(); + if buf.len() > ring { + let drop = buf.len() - ring; + buf.drain(0..drop); + } +} + +/// Load the persisted burndown history for one provider+account into the +/// process-local store, so history survives a restart. Runs once per scope; +/// later calls are cheap no-ops (the loaded marker is kept in the store). +pub fn load_persisted_history( + provider_id: &str, + account_key: Option<&str>, +) -> Vec { + static LOADED: LazyLock>> = + LazyLock::new(|| Mutex::new(HashSet::new())); + let scope = ForecastScope::new(provider_id, account_key); + if let Ok(guard) = LOADED.lock() + && guard.contains(&scope) + { + return global_history_store() + .lock() + .ok() + .map(|store| store.histories(&scope)) + .unwrap_or_default(); + } + let series = quota_burndown::load_persisted_series(provider_id, account_key); + if let Ok(mut loaded) = LOADED.lock() { + loaded.insert(scope.clone()); + } + if let Ok(mut store) = global_history_store().lock() { + store.adopt_persisted(&scope, &series); + } + global_history_store() + .lock() + .ok() + .map(|store| store.histories(&scope)) + .unwrap_or_default() +} + +/// Global in-process history used by Claude/Codex snapshot refresh. +pub fn global_history_store() -> &'static Mutex { + &HISTORY_STORE +} + +/// Record session/weekly windows from a live provider usage snapshot. +/// +/// `account_key` scopes the history so switching accounts on one provider does not +/// blend burn observations across plans. +pub fn record_provider_windows( + provider_id: &str, + account_key: Option<&str>, + session: &crate::core::RateWindow, + weekly: Option<&crate::core::RateWindow>, + now: DateTime, +) { + let Some(session_entry) = history_entry(session, SESSION_WINDOW_MINUTES, now) else { + return; + }; + let weekly_entry = weekly.and_then(|w| history_entry(w, WEEKLY_WINDOW_MINUTES, now)); + + if let Ok(mut guard) = global_history_store().lock() { + guard.record( + &ForecastScope::new(provider_id, account_key), + Some(session_entry.clone()), + weekly_entry.clone(), + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ); + } + + // Best-effort persistence for the burndown chart (upstream #4085 keeps a + // durable store; a failed write only costs the chart, never the forecast). + if let Err(error) = quota_burndown::persist_recorded_windows( + provider_id, + account_key, + Some(&session_entry), + weekly_entry.as_ref(), + now, + ) { + tracing::debug!(%error, "quota burndown persistence failed"); + } +} + +/// A clamped observation of `window`, or `None` when it is informational, has +/// another length than `minutes`, or carries a non-finite percent. +fn history_entry( + window: &crate::core::RateWindow, + minutes: u32, + now: DateTime, +) -> Option { + if window.is_informational + || window.window_minutes != Some(minutes) + || !window.used_percent.is_finite() + { + return None; + } + Some(PlanUtilizationHistoryEntry { + captured_at: now, + used_percent: window.used_percent.clamp(0.0, 100.0), + resets_at: window.resets_at, + }) +} + +/// Last learned full-session burn estimate retained across idle refreshes. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RetainedFullSessionEstimate { + pub estimate: f64, + pub updated_at: DateTime, +} + +/// Keep the last positive full-session estimate when a refresh yields none. +/// +/// `fresh` wins when present and valid; otherwise `previous` is kept so the UI +/// does not blank while the session window is idle (upstream #2336). +pub fn retain_last_full_session_estimate( + previous: Option, + fresh: Option, + now: DateTime, +) -> Option { + match fresh { + Some(estimate) if estimate.is_finite() && estimate > 0.0 => { + Some(RetainedFullSessionEstimate { + estimate, + updated_at: now, + }) + } + _ => previous, + } +} + +// ponytail: in-memory only, lost on restart; upgrade = persist to cache file. +#[derive(Debug, Default)] +struct LastFullSessionEstimateStore { + by_scope: HashMap, +} + +static LAST_FULL_SESSION_ESTIMATE_STORE: LazyLock> = + LazyLock::new(|| Mutex::new(LastFullSessionEstimateStore::default())); + +fn last_full_session_estimate_store() -> &'static Mutex { + &LAST_FULL_SESSION_ESTIMATE_STORE +} + +/// Remember/recall the last learned full-session burn estimate for a provider+account. +pub fn remember_full_session_estimate( + scope: &ForecastScope, + fresh: Option, + now: DateTime, +) -> Option { + let Ok(mut guard) = last_full_session_estimate_store().lock() else { + return fresh.filter(|v| v.is_finite() && *v > 0.0); + }; + let previous = guard.by_scope.get(scope).copied(); + let retained = retain_last_full_session_estimate(previous, fresh, now); + if let Some(entry) = retained { + guard.by_scope.insert(scope.clone(), entry); + Some(entry.estimate) + } else { + guard.by_scope.remove(scope); + None + } +} + +/// Compute forecast for a provider+account using the in-process history ring. +pub fn forecast_for_provider( + provider_id: &str, + account_key: Option<&str>, + session: &crate::core::RateWindow, + weekly: &crate::core::RateWindow, + now: DateTime, + work_days: Option, +) -> Option { + let scope = ForecastScope::new(provider_id, account_key); + let histories = global_history_store() + .lock() + .ok() + .map(|g| g.histories(&scope)) + .unwrap_or_default(); + let fresh_burn = SessionEquivalentBurnEstimator::estimate( + &histories, + session.resets_at, + now, + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ); + let fresh_median = fresh_burn + .as_ref() + .map(|b| b.median_weekly_percent_per_window); + let sample_count = fresh_burn + .as_ref() + .map(|b| b.sample_count) + .unwrap_or(SessionEquivalentBurnEstimator::MINIMUM_SAMPLE_COUNT); + let median = remember_full_session_estimate(&scope, fresh_median, now)?; + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: median, + sample_count, + }; + SessionEquivalentForecast::make(session, weekly, &burn, None, now, work_days) +} diff --git a/rust/src/core/session_equivalent_forecast/tests.rs b/rust/src/core/session_equivalent_forecast/tests.rs new file mode 100644 index 0000000000..ec1af10887 --- /dev/null +++ b/rust/src/core/session_equivalent_forecast/tests.rs @@ -0,0 +1,599 @@ +use super::*; +use crate::core::RateWindow; +use chrono::TimeZone; + +fn ts(secs: i64) -> DateTime { + Utc.timestamp_opt(secs, 0).single().unwrap() +} + +fn entry(captured: i64, used: f64, resets: i64) -> PlanUtilizationHistoryEntry { + PlanUtilizationHistoryEntry { + captured_at: ts(captured), + used_percent: used, + resets_at: Some(ts(resets)), + } +} + +fn session_secs() -> i64 { + i64::from(SESSION_WINDOW_MINUTES) * 60 +} + +fn weekly_secs() -> i64 { + i64::from(WEEKLY_WINDOW_MINUTES) * 60 +} + +#[test] +fn persisted_series_map_by_name_in_session_then_weekly_order() { + let mut store = SessionEquivalentHistoryStore::default(); + let scope = ForecastScope::new("pin-series", None); + let persisted = |name: &str, used: f64| quota_burndown::PersistedPlanSeries { + name: name.to_string(), + window_minutes: 1, + entries: vec![PersistedEntry::from_history(&entry( + 1_700_000_000, + used, + 1_700_018_000, + ))], + }; + store.adopt_persisted( + &scope, + &[ + persisted("weekly", 40.0), + persisted("monthly", 70.0), + persisted("session", 10.0), + ], + ); + let shape: Vec<_> = store + .histories(&scope) + .iter() + .map(|h| (h.name, h.window_minutes, h.entries[0].used_percent)) + .collect(); + assert_eq!( + shape, + vec![ + (PlanUtilizationSeriesName::Session, 300, 10.0), + (PlanUtilizationSeriesName::Weekly, 10_080, 40.0), + ] + ); + + let mut weekly_only = SessionEquivalentHistoryStore::default(); + weekly_only.record( + &scope, + None, + Some(entry(1_700_000_000, 5.0, 1_700_600_000)), + 7, + ); + let names: Vec<_> = weekly_only + .histories(&scope) + .iter() + .map(|h| h.name) + .collect(); + assert_eq!(names, vec![PlanUtilizationSeriesName::Weekly]); +} + +#[test] +fn record_provider_windows_filters_and_clamps_both_windows() { + let now = ts(1_700_000_000); + let window = |minutes: u32, used: f64| RateWindow { + used_percent: used, + window_minutes: Some(minutes), + resets_at: Some(ts(1_700_018_000)), + ..RateWindow::new(0.0) + }; + let recorded = |provider: &str| { + global_history_store() + .lock() + .unwrap() + .histories(&ForecastScope::new(provider, None)) + .iter() + .map(|h| { + ( + h.name, + h.entries.iter().map(|e| e.used_percent).collect::>(), + ) + }) + .collect::>() + }; + + record_provider_windows( + "pin-rec-a", + None, + &window(300, 120.0), + Some(&window(10_080, -3.0)), + now, + ); + assert_eq!( + recorded("pin-rec-a"), + vec![ + (PlanUtilizationSeriesName::Session, vec![100.0]), + (PlanUtilizationSeriesName::Weekly, vec![0.0]), + ] + ); + + let informational = RateWindow { + is_informational: true, + ..window(10_080, 50.0) + }; + record_provider_windows( + "pin-rec-b", + None, + &window(300, 20.0), + Some(&informational), + now, + ); + record_provider_windows( + "pin-rec-b", + None, + &window(300, 30.0), + Some(&window(43_200, 50.0)), + now, + ); + record_provider_windows( + "pin-rec-b", + None, + &window(300, 40.0), + Some(&window(10_080, f64::NAN)), + now, + ); + assert_eq!( + recorded("pin-rec-b"), + vec![(PlanUtilizationSeriesName::Session, vec![20.0, 30.0, 40.0])] + ); + + let informational_session = RateWindow { + is_informational: true, + ..window(300, 20.0) + }; + record_provider_windows( + "pin-rec-c", + None, + &informational_session, + Some(&window(10_080, 50.0)), + now, + ); + record_provider_windows( + "pin-rec-c", + None, + &window(60, 20.0), + Some(&window(10_080, 50.0)), + now, + ); + record_provider_windows("pin-rec-c", None, &window(300, f64::INFINITY), None, now); + assert!(recorded("pin-rec-c").is_empty()); +} + +/// Build three completed full-burn sessions with aligned weekly observations. +fn three_sample_histories( + base: i64, +) -> ( + Vec, + DateTime, + DateTime, +) { + let s = session_secs(); + let w = weekly_secs(); + let weekly_reset = base + w; + + let mut session_entries = Vec::new(); + let mut weekly_entries = Vec::new(); + + // Three completed sessions ending at base+s, base+2s, base+3s. + for i in 0..3 { + let start = base + i * s; + let end = start + s; + // Window-start weekly anchor (session used 0). + weekly_entries.push(entry(start - 30, 10.0 + i as f64 * 5.0, weekly_reset)); + // Mid / end session captures with exact weekly alignment. + session_entries.push(entry(start + 60, 40.0, end)); + weekly_entries.push(entry(start + 60, 12.0 + i as f64 * 5.0, weekly_reset)); + session_entries.push(entry(end - 60, 100.0, end)); + weekly_entries.push(entry(end - 60, 15.0 + i as f64 * 5.0, weekly_reset)); + // End-of-window weekly for max>=100 path. + weekly_entries.push(entry(end, 15.5 + i as f64 * 5.0, weekly_reset)); + } + + // Current open session. + let current_start = base + 3 * s; + let current_end = current_start + s; + let now_secs = current_start + 30 * 60; + session_entries.push(entry(now_secs, 20.0, current_end)); + weekly_entries.push(entry(now_secs, 30.0, weekly_reset)); + + session_entries.sort_by_key(|e| e.captured_at); + weekly_entries.sort_by_key(|e| e.captured_at); + + let histories = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: session_entries, + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: weekly_entries, + }, + ]; + (histories, ts(now_secs), ts(current_end)) +} + +#[test] +fn median_math_from_three_full_sessions() { + let base = 1_700_000_000_i64; + let (histories, now, current_end) = three_sample_histories(base); + let estimate = SessionEquivalentBurnEstimator::estimate( + &histories, + Some(current_end), + now, + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ) + .expect("estimate"); + assert_eq!(estimate.sample_count, 3); + // Mid-session captures dominate when start/end anchors are optional: + // weekly 12→15 over session 40→100 ⇒ 100 * 3 / 60 = 5.0 + // (window-start + end anchors yield the same full-allowance burn here). + assert!( + (estimate.median_weekly_percent_per_window - 5.0).abs() < 0.01, + "median={}", + estimate.median_weekly_percent_per_window + ); +} + +#[test] +fn grouping_tolerance_merges_nearby_resets() { + let s = session_secs(); + let base = 1_700_100_000_i64; + let weekly_reset = base + weekly_secs(); + let end = base + s; + + // Nearby resets within ±120s must merge into one group (not out-of-order). + let single_group = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 30.0, end), + entry(base + 120, 60.0, end + 100), + entry(base + 180, 90.0, end + 110), + ], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base - 30, 5.0, weekly_reset), + entry(base + 60, 8.0, weekly_reset), + entry(base + 120, 10.0, weekly_reset), + entry(base + 180, 12.0, weekly_reset), + ], + }, + ]; + // One completed group → below minimum sample count, but not an ordering reject. + assert!( + SessionEquivalentBurnEstimator::estimate(&single_group, Some(ts(end + s)), ts(end + 60), 7) + .is_none() + ); + + // Strictly decreasing resets_at across groups is rejected. + let out_of_order = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 30.0, end + s), + entry(base + 120, 60.0, end), // jumps backward beyond tolerance + ], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 8.0, weekly_reset), + entry(base + 120, 10.0, weekly_reset), + ], + }, + ]; + assert!( + SessionEquivalentBurnEstimator::estimate( + &out_of_order, + Some(ts(end + 2 * s)), + ts(end + s + 60), + 7 + ) + .is_none() + ); +} + +#[test] +fn insufficient_samples_returns_none() { + let base = 1_700_200_000_i64; + let s = session_secs(); + let weekly_reset = base + weekly_secs(); + let end = base + s; + let histories = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![entry(base + 60, 50.0, end), entry(base + 120, 80.0, end)], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base - 30, 10.0, weekly_reset), + entry(base + 60, 12.0, weekly_reset), + entry(base + 120, 14.0, weekly_reset), + ], + }, + ]; + assert!( + SessionEquivalentBurnEstimator::estimate(&histories, Some(ts(end + s)), ts(end + 60), 7) + .is_none() + ); +} + +#[test] +fn forecast_make_fractional_windows() { + let now = ts(1_700_300_000); + let session = RateWindow::with_details( + 25.0, + Some(SESSION_WINDOW_MINUTES), + Some(now + Duration::hours(3)), + None, + ); + let weekly = RateWindow::with_details( + 40.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(3)), + None, + ); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 12.0, + sample_count: 3, + }; + let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) + .expect("forecast"); + // remaining weekly 60 / 12 = 5.0 + assert!((forecast.estimated_windows_to_exhaust_weekly - 5.0).abs() < 1e-9); + assert!(forecast.available_windows_until_reset > 0.0); + // Mirrors the production whole-window truncation. + #[allow( + clippy::cast_possible_truncation, + reason = "assertion mirrors the production whole-window truncation" + )] + let expected_windows = forecast.available_windows_until_reset.floor() as i64; + assert_eq!(forecast.windows_until_reset, expected_windows); + assert_eq!(forecast.sample_count, 3); + assert!((forecast.weekly_used_percent - 40.0).abs() < 1e-9); +} + +#[test] +fn work_days_reduces_available_windows_vs_wall_clock() { + // Friday 18:00 UTC → Monday 18:00 UTC (weekend in between). + let friday = Utc + .with_ymd_and_hms(2024, 6, 14, 18, 0, 0) + .single() + .unwrap(); + assert_eq!(friday.weekday(), Weekday::Fri); + let monday = Utc + .with_ymd_and_hms(2024, 6, 17, 18, 0, 0) + .single() + .unwrap(); + + let wall = effective_remaining_seconds(friday, monday, None); + let work5 = effective_remaining_seconds(friday, monday, Some(5)); + assert!((wall - 3.0 * 24.0 * 3600.0).abs() < 1.0); + // Fri 18→24 = 6h, Mon 0→18 = 18h → 24h work seconds (no Sat/Sun). + assert!((work5 - 24.0 * 3600.0).abs() < 1.0); + assert!(work5 < wall); + + let session = RateWindow::with_details( + 10.0, + Some(SESSION_WINDOW_MINUTES), + Some(friday + Duration::hours(4)), + None, + ); + let weekly = RateWindow::with_details(20.0, Some(WEEKLY_WINDOW_MINUTES), Some(monday), None); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 4, + }; + let wall_f = + SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, None).unwrap(); + let work_f = + SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, Some(5)).unwrap(); + assert!(work_f.available_windows_until_reset < wall_f.available_windows_until_reset); + // Mirrors the production whole-window truncation. + #[allow( + clippy::cast_possible_truncation, + reason = "assertion mirrors the production whole-window truncation" + )] + let expected_work_windows = work_f.available_windows_until_reset.floor() as i64; + assert_eq!(work_f.windows_until_reset, expected_work_windows); +} + +#[test] +fn work_days_out_of_range_uses_wall_clock() { + let now = ts(1_700_400_000); + let reset = now + Duration::days(2); + let wall = effective_remaining_seconds(now, reset, None); + assert!((effective_remaining_seconds(now, reset, Some(1)) - wall).abs() < 1e-9); + assert!((effective_remaining_seconds(now, reset, Some(7)) - wall).abs() < 1e-9); +} + +#[test] +fn history_ring_retains_latest_samples() { + let mut store = SessionEquivalentHistoryStore::default(); + let scope = ForecastScope::new("claude", None); + let base = ts(1_700_500_000); + for i in 0..30 { + store.record( + &scope, + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(i), + used_percent: i as f64, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + } + let h = store.histories(&scope); + assert_eq!(h.len(), 1); + assert!(h[0].entries.len() <= 7 * 8); + assert!((h[0].entries.last().unwrap().used_percent - 29.0).abs() < 1e-9); +} + +/// Two accounts on one provider must not share burn samples. +/// +/// Regression: history was keyed by `provider_id` alone, so switching the active +/// account blended both accounts' observations into one ring and produced a median +/// drawn from a mixture of plans. +#[test] +fn history_is_isolated_per_account() { + let mut store = SessionEquivalentHistoryStore::default(); + let base = ts(1_700_500_000); + let alice = ForecastScope::new("codex", Some("alice@example.com")); + let bob = ForecastScope::new("codex", Some("bob@example.com")); + + for i in 0..5 { + store.record( + &alice, + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(i), + used_percent: 10.0, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + } + store.record( + &bob, + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(99), + used_percent: 90.0, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + + let alice_hist = store.histories(&alice); + let bob_hist = store.histories(&bob); + assert_eq!(alice_hist[0].entries.len(), 5); + assert_eq!(bob_hist[0].entries.len(), 1); + assert!( + alice_hist[0] + .entries + .iter() + .all(|e| (e.used_percent - 10.0).abs() < 1e-9), + "bob's 90% sample leaked into alice's history" + ); + assert!((bob_hist[0].entries[0].used_percent - 90.0).abs() < 1e-9); + + // A provider with no account discriminator is its own bucket, not a catch-all. + assert!( + store + .histories(&ForecastScope::new("codex", None)) + .is_empty() + ); +} + +#[test] +fn make_rejects_bad_windows_and_zero_remaining() { + let now = ts(1_700_600_000); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 3, + }; + let session = RateWindow::with_details( + 10.0, + Some(SESSION_WINDOW_MINUTES), + Some(now + Duration::hours(2)), + None, + ); + let exhausted = RateWindow::with_details( + 100.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(2)), + None, + ); + assert!( + SessionEquivalentForecast::make(&session, &exhausted, &burn, None, now, None).is_none() + ); + + let low_samples = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 2, + }; + let weekly = RateWindow::with_details( + 50.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(2)), + None, + ); + assert!( + SessionEquivalentForecast::make(&session, &weekly, &low_samples, None, now, None).is_none() + ); +} + +#[test] +fn retain_last_full_session_estimate_survives_idle_refresh() { + let t0 = ts(1_700_000_000); + let t1 = t0 + Duration::hours(1); + let learned = retain_last_full_session_estimate(None, Some(12.5), t0) + .expect("fresh estimate should stick"); + assert_eq!(learned.estimate, 12.5); + assert_eq!(learned.updated_at, t0); + + // Idle refresh yields no fresh sample — keep the last learned value. + let idle = retain_last_full_session_estimate(Some(learned), None, t1) + .expect("idle refresh must keep last estimate"); + assert_eq!(idle.estimate, 12.5); + assert_eq!(idle.updated_at, t0); + + // Non-positive / non-finite fresh values also keep previous. + let still = retain_last_full_session_estimate(Some(idle), Some(0.0), t1) + .expect("zero fresh must not clear"); + assert_eq!(still.estimate, 12.5); + let still = retain_last_full_session_estimate(Some(still), Some(f64::NAN), t1) + .expect("nan fresh must not clear"); + assert_eq!(still.estimate, 12.5); +} + +#[test] +fn retain_last_full_session_estimate_replaced_by_fresh_sample() { + let t0 = ts(1_700_000_000); + let t1 = t0 + Duration::hours(2); + let previous = retain_last_full_session_estimate(None, Some(8.0), t0).unwrap(); + let next = retain_last_full_session_estimate(Some(previous), Some(15.0), t1) + .expect("fresh sample replaces previous"); + assert_eq!(next.estimate, 15.0); + assert_eq!(next.updated_at, t1); +} + +#[test] +fn make_keeps_forecast_when_session_reset_is_expired_idle() { + let now = ts(1_700_600_000); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 3, + }; + // Session reset already passed (idle) — forecast still derives from weekly. + let session = RateWindow::with_details( + 0.0, + Some(SESSION_WINDOW_MINUTES), + Some(now - Duration::hours(1)), + None, + ); + let weekly = RateWindow::with_details( + 40.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(3)), + None, + ); + let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) + .expect("idle expired session should still forecast"); + assert!((forecast.estimated_windows_to_exhaust_weekly - 6.0).abs() < 1e-9); +} diff --git a/rust/src/core/session_quota.rs b/rust/src/core/session_quota.rs deleted file mode 100755 index f37f406b66..0000000000 --- a/rust/src/core/session_quota.rs +++ /dev/null @@ -1,202 +0,0 @@ -//! Session Quota Notifications -//! -//! Monitors session quota changes and sends notifications when: -//! - Quota becomes depleted (0% remaining) -//! - Quota is restored (becomes available again) - -#![allow( - dead_code, - reason = "quota types reserved for future session-limit enforcement" -)] -#![allow( - unused_imports, - reason = "imports needed for future session-limit enforcement wiring" -)] - -use crate::core::ProviderId; - -/// Threshold for considering quota as depleted (0.01%) -const DEPLETED_THRESHOLD: f64 = 0.0001; - -/// Session quota transition states -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SessionQuotaTransition { - /// No change in quota state - None, - /// Quota just became depleted - Depleted, - /// Quota was restored after being depleted - Restored, -} - -/// Logic for detecting session quota transitions -pub struct SessionQuotaLogic; - -impl SessionQuotaLogic { - /// Check if remaining percentage is considered depleted - pub fn is_depleted(remaining: Option) -> bool { - match remaining { - Some(r) => r <= DEPLETED_THRESHOLD, - None => false, - } - } - - /// Detect transition between previous and current remaining percentages - pub fn transition( - previous_remaining: Option, - current_remaining: Option, - ) -> SessionQuotaTransition { - let Some(current) = current_remaining else { - return SessionQuotaTransition::None; - }; - let Some(previous) = previous_remaining else { - return SessionQuotaTransition::None; - }; - - let was_depleted = previous <= DEPLETED_THRESHOLD; - let is_depleted = current <= DEPLETED_THRESHOLD; - - if !was_depleted && is_depleted { - SessionQuotaTransition::Depleted - } else if was_depleted && !is_depleted { - SessionQuotaTransition::Restored - } else { - SessionQuotaTransition::None - } - } -} - -/// Session quota notifier that tracks state and sends notifications -pub struct SessionQuotaNotifier { - /// Previous remaining percentages by provider - previous_remaining: std::collections::HashMap, -} - -impl SessionQuotaNotifier { - /// Create a new notifier - pub fn new() -> Self { - Self { - previous_remaining: std::collections::HashMap::new(), - } - } - - /// Update quota for a provider and send notification if state changed - pub fn update(&mut self, provider: ProviderId, current_remaining: Option) { - let previous = self.previous_remaining.get(&provider).copied(); - let transition = SessionQuotaLogic::transition(previous, current_remaining); - - if transition != SessionQuotaTransition::None { - self.post_notification(transition, provider); - } - - // Store current as previous for next comparison - if let Some(remaining) = current_remaining { - self.previous_remaining.insert(provider, remaining); - } - } - - /// Post a notification for the given transition - fn post_notification(&self, transition: SessionQuotaTransition, provider: ProviderId) { - let provider_name = provider.display_name(); - - let (title, body) = match transition { - SessionQuotaTransition::None => return, - SessionQuotaTransition::Depleted => ( - format!("{} session depleted", provider_name), - "0% left. Will notify when it's available again.".to_string(), - ), - SessionQuotaTransition::Restored => ( - format!("{} session restored", provider_name), - "Session quota is available again.".to_string(), - ), - }; - - tracing::info!( - provider = %provider.cli_name(), - transition = ?transition, - "Session quota notification: {} - {}", - title, - body - ); - - // Note: Actual toast notification is handled by NotificationManager.check_session_transition() - // This function is for logging and integration with the notification system - } - - /// Reset tracking for a provider - pub fn reset(&mut self, provider: ProviderId) { - self.previous_remaining.remove(&provider); - } - - /// Reset all tracking - pub fn reset_all(&mut self) { - self.previous_remaining.clear(); - } -} - -impl Default for SessionQuotaNotifier { - fn default() -> Self { - Self::new() - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_is_depleted() { - assert!(SessionQuotaLogic::is_depleted(Some(0.0))); - assert!(SessionQuotaLogic::is_depleted(Some(0.00001))); - assert!(!SessionQuotaLogic::is_depleted(Some(0.01))); - assert!(!SessionQuotaLogic::is_depleted(Some(50.0))); - assert!(!SessionQuotaLogic::is_depleted(None)); - } - - #[test] - fn test_transition_none() { - // No change when both are above threshold - assert_eq!( - SessionQuotaLogic::transition(Some(50.0), Some(40.0)), - SessionQuotaTransition::None - ); - // No change when both are below threshold - assert_eq!( - SessionQuotaLogic::transition(Some(0.0), Some(0.0)), - SessionQuotaTransition::None - ); - // No change with None values - assert_eq!( - SessionQuotaLogic::transition(None, Some(50.0)), - SessionQuotaTransition::None - ); - assert_eq!( - SessionQuotaLogic::transition(Some(50.0), None), - SessionQuotaTransition::None - ); - } - - #[test] - fn test_transition_depleted() { - assert_eq!( - SessionQuotaLogic::transition(Some(10.0), Some(0.0)), - SessionQuotaTransition::Depleted - ); - assert_eq!( - SessionQuotaLogic::transition(Some(0.001), Some(0.00001)), - SessionQuotaTransition::Depleted - ); - } - - #[test] - fn test_transition_restored() { - assert_eq!( - SessionQuotaLogic::transition(Some(0.0), Some(10.0)), - SessionQuotaTransition::Restored - ); - assert_eq!( - SessionQuotaLogic::transition(Some(0.00001), Some(1.0)), - SessionQuotaTransition::Restored - ); - } -} diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index bae251693b..204ec5c441 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -50,335 +50,252 @@ pub struct TokenAccountSupport { } impl TokenAccountSupport { + fn cookie(subtitle: &'static str, placeholder: &'static str) -> Self { + Self { + title: "Session tokens", + subtitle, + placeholder, + injection: TokenInjection::CookieHeader, + requires_manual_cookie_source: true, + cookie_name: None, + } + } + + fn api_key( + title: &'static str, + env_key: &str, + subtitle: &'static str, + placeholder: &'static str, + ) -> Self { + Self { + title, + subtitle, + placeholder, + injection: TokenInjection::Environment { + key: env_key.to_string(), + }, + requires_manual_cookie_source: false, + cookie_name: None, + } + } + + fn with_cookie_name(mut self, cookie_name: &'static str) -> Self { + self.cookie_name = Some(cookie_name); + self + } + /// Get token account support for a provider pub fn for_provider(provider: ProviderId) -> Option { match provider { - ProviderId::Claude => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store Claude sessionKey cookies for settings-page usage. OAuth tokens are kept as a legacy fallback.", - placeholder: "Paste sessionKey value or Cookie: sessionKey=...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("sessionKey"), - }), - ProviderId::Zai => Some(TokenAccountSupport { - title: "API tokens", - subtitle: "Stored locally in token-accounts.json. Team usage can use workspace_id as organization|project.", - placeholder: "Paste token...", - injection: TokenInjection::Environment { - key: "Z_AI_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Cursor => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Cursor Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::OpenCode => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple OpenCode Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Factory => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Factory Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Alibaba => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Alibaba Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::AlibabaTokenPlan => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Alibaba Token Plan Cookie headers.", - placeholder: "Cookie: cna=...; login_aliyunid_csrf=...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::MiniMax => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple MiniMax Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Augment => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Augment Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Amp => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Amp Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Ollama => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Ollama Cookie headers or __Secure-session values.", - placeholder: "__Secure-session value or Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("__Secure-session"), - }), - ProviderId::T3Chat => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple T3 Chat Cookie headers or full browser cURL captures.", - placeholder: "Cookie: ... or curl ... -H 'Cookie: ...'", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::ZoomMate => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple ZoomMate Cookie headers or credits/status cURL captures.", - placeholder: "Cookie: ... or curl 'https://ai.zoom.us/.../credits/status' -H 'Authorization: Bearer ...'", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Mistral => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Mistral Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Manus => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Manus session_id values.", - placeholder: "session_id value or Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("session_id"), - }), - ProviderId::MiMo => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Xiaomi MiMo Cookie headers.", - placeholder: "Cookie: api-platform_serviceToken=...; userId=...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::CommandCode => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Command Code Cookie headers or Better Auth values.", - placeholder: "Cookie: __Secure-commandcode_prod_.session_token=... or better-auth value", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("__Secure-better-auth.session_token"), - }), - ProviderId::Qoder => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Qoder Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::CodeBuddy => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store CodeBuddy CN Cookie headers (from plans-usage DevTools cURL).", - placeholder: "Cookie: session=...; ... (or paste full Cookie header)", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Sakana => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Sakana Console Cookie headers.", - placeholder: "Cookie: ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: None, - }), - ProviderId::Notion => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Notion Cookie headers or token_v2 values.", - placeholder: "Cookie: token_v2=... or paste the token_v2 value", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("token_v2"), - }), - ProviderId::Replicate => Some(TokenAccountSupport { - title: "Session tokens", - subtitle: "Store multiple Replicate Cookie headers from the billing page.", - placeholder: "Cookie: sessionid=...; ...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("sessionid"), - }), - ProviderId::Sub2Api => Some(TokenAccountSupport { - title: "Group API keys", - subtitle: "Store multiple sub2api group API keys with labels such as Claude, Codex, or Gemini.", - placeholder: "sk-...", - injection: TokenInjection::Environment { - key: "SUB2API_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::DeepInfra => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple DeepInfra API keys.", - placeholder: "API key from deepinfra.com/dash", - injection: TokenInjection::Environment { - key: "DEEPINFRA_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::HuggingFace => Some(TokenAccountSupport { - title: "API tokens", - subtitle: "Store multiple Hugging Face access tokens.", - placeholder: "Paste a Hugging Face access token", - injection: TokenInjection::Environment { - key: "CODEXBAR_HUGGINGFACE_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::AiAnd => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple ai& API keys.", - placeholder: "API key from console.aiand.com", - injection: TokenInjection::Environment { - key: "AIAND_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::ZenMux => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple ZenMux Management API keys.", - placeholder: "Management API key", - injection: TokenInjection::Environment { - key: "ZENMUX_MANAGEMENT_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::ClinePass => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple ClinePass API keys. Without one, CodexBar reads your existing Cline session (run cline auth) without copying it.", - placeholder: "API key", - injection: TokenInjection::Environment { - key: "CLINE_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Neuralwatt => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple Neuralwatt API keys.", - placeholder: "API key", - injection: TokenInjection::Environment { - key: "NEURALWATT_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Grok => Some(TokenAccountSupport { + ProviderId::Claude => Some( + Self::cookie( + "Store Claude sessionKey cookies for settings-page usage. OAuth tokens are kept as a legacy fallback.", + "Paste sessionKey value or Cookie: sessionKey=...", + ) + .with_cookie_name("sessionKey"), + ), + ProviderId::Zai => Some(Self::api_key( + "API tokens", + "Z_AI_API_KEY", + "Stored locally in token-accounts.json. Team usage can use workspace_id as organization|project.", + "Paste token...", + )), + ProviderId::Cursor => Some(Self::cookie( + "Store multiple Cursor Cookie headers.", + "Cookie: ...", + )), + ProviderId::OpenCode => Some(Self::cookie( + "Store multiple OpenCode Cookie headers.", + "Cookie: ...", + )), + ProviderId::Factory => Some(Self::cookie( + "Store multiple Factory Cookie headers.", + "Cookie: ...", + )), + ProviderId::Alibaba => Some(Self::cookie( + "Store multiple Alibaba Cookie headers.", + "Cookie: ...", + )), + ProviderId::AlibabaTokenPlan => Some(Self::cookie( + "Store multiple Alibaba Token Plan Cookie headers.", + "Cookie: cna=...; login_aliyunid_csrf=...", + )), + ProviderId::MiniMax => Some(Self::cookie( + "Store multiple MiniMax Cookie headers.", + "Cookie: ...", + )), + ProviderId::Augment => Some(Self::cookie( + "Store multiple Augment Cookie headers.", + "Cookie: ...", + )), + ProviderId::Amp => Some(Self::cookie( + "Store multiple Amp Cookie headers.", + "Cookie: ...", + )), + ProviderId::Ollama => Some( + Self::cookie( + "Store multiple Ollama Cookie headers or __Secure-session values.", + "__Secure-session value or Cookie: ...", + ) + .with_cookie_name("__Secure-session"), + ), + ProviderId::T3Chat => Some(Self::cookie( + "Store multiple T3 Chat Cookie headers or full browser cURL captures.", + "Cookie: ... or curl ... -H 'Cookie: ...'", + )), + ProviderId::ZoomMate => Some(Self::cookie( + "Store multiple ZoomMate Cookie headers or credits/status cURL captures.", + "Cookie: ... or curl 'https://ai.zoom.us/.../credits/status' -H 'Authorization: Bearer ...'", + )), + ProviderId::Mistral => Some(Self::cookie( + "Store multiple Mistral Cookie headers.", + "Cookie: ...", + )), + ProviderId::Manus => Some( + Self::cookie( + "Store multiple Manus session_id values.", + "session_id value or Cookie: ...", + ) + .with_cookie_name("session_id"), + ), + ProviderId::MiMo => Some(Self::cookie( + "Store multiple Xiaomi MiMo Cookie headers.", + "Cookie: api-platform_serviceToken=...; userId=...", + )), + ProviderId::CommandCode => Some( + Self::cookie( + "Store multiple Command Code Cookie headers or Better Auth values.", + "Cookie: __Secure-commandcode_prod_.session_token=... or better-auth value", + ) + .with_cookie_name("__Secure-better-auth.session_token"), + ), + ProviderId::Qoder => Some(Self::cookie( + "Store multiple Qoder Cookie headers.", + "Cookie: ...", + )), + ProviderId::CodeBuddy => Some(Self::cookie( + "Store CodeBuddy CN Cookie headers (from plans-usage DevTools cURL).", + "Cookie: session=...; ... (or paste full Cookie header)", + )), + ProviderId::Sakana => Some(Self::cookie( + "Store multiple Sakana Console Cookie headers.", + "Cookie: ...", + )), + ProviderId::Notion => Some( + Self::cookie( + "Store multiple Notion Cookie headers or token_v2 values.", + "Cookie: token_v2=... or paste the token_v2 value", + ) + .with_cookie_name("token_v2"), + ), + ProviderId::Replicate => Some( + Self::cookie( + "Store multiple Replicate Cookie headers from the billing page.", + "Cookie: sessionid=...; ...", + ) + .with_cookie_name("sessionid"), + ), + ProviderId::Sub2Api => Some(Self::api_key( + "Group API keys", + "SUB2API_API_KEY", + "Store multiple sub2api group API keys with labels such as Claude, Codex, or Gemini.", + "sk-...", + )), + ProviderId::DeepInfra => Some(Self::api_key( + "API keys", + "DEEPINFRA_API_KEY", + "Store multiple DeepInfra API keys.", + "API key from deepinfra.com/dash", + )), + ProviderId::HuggingFace => Some(Self::api_key( + "API tokens", + "CODEXBAR_HUGGINGFACE_API_KEY", + "Store multiple Hugging Face access tokens.", + "Paste a Hugging Face access token", + )), + ProviderId::AiAnd => Some(Self::api_key( + "API keys", + "AIAND_API_KEY", + "Store multiple ai& API keys.", + "API key from console.aiand.com", + )), + ProviderId::ZenMux => Some(Self::api_key( + "API keys", + "ZENMUX_MANAGEMENT_API_KEY", + "Store multiple ZenMux Management API keys.", + "Management API key", + )), + ProviderId::ClinePass => Some(Self::api_key( + "API keys", + "CLINE_API_KEY", + "Store multiple ClinePass API keys. Without one, CodexBar reads your existing Cline session (run cline auth) without copying it.", + "API key", + )), + ProviderId::Neuralwatt => Some(Self::api_key( + "API keys", + "NEURALWATT_API_KEY", + "Store multiple Neuralwatt API keys.", + "API key", + )), + ProviderId::Grok => Some(Self { title: "Grok credentials", - subtitle: "Store SuperGrok bearer tokens or grok.com Cookie headers.", - placeholder: "Bearer token or Cookie: ...", - injection: TokenInjection::CookieHeader, requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Xai => Some(TokenAccountSupport { - title: "Management API keys", - subtitle: "Store multiple xAI Management API keys. Team ID is set separately under provider settings.", - placeholder: "xai-... Management API key from console.x.ai", - injection: TokenInjection::Environment { - key: "XAI_MANAGEMENT_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), + ..Self::cookie( + "Store SuperGrok bearer tokens or grok.com Cookie headers.", + "Bearer token or Cookie: ...", + ) + }), + ProviderId::Xai => Some(Self::api_key( + "Management API keys", + "XAI_MANAGEMENT_API_KEY", + "Store multiple xAI Management API keys. Team ID is set separately under provider settings.", + "xai-... Management API key from console.x.ai", + )), // Upstream 0.45 #2271: labeled OpenRouter API keys via token accounts. - ProviderId::OpenRouter => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple OpenRouter API keys.", - placeholder: "sk-or-v1-...", - injection: TokenInjection::Environment { - key: "OPENROUTER_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Copilot => Some(TokenAccountSupport { - title: "GitHub accounts", - subtitle: "Store GitHub OAuth tokens for Copilot plan usage.", - placeholder: "Sign in with GitHub or paste a GitHub OAuth token...", - injection: TokenInjection::Environment { - key: "GITHUB_TOKEN".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::Kimi => Some(TokenAccountSupport { + ProviderId::OpenRouter => Some(Self::api_key( + "API keys", + "OPENROUTER_API_KEY", + "Store multiple OpenRouter API keys.", + "sk-or-v1-...", + )), + ProviderId::Copilot => Some(Self::api_key( + "GitHub accounts", + "GITHUB_TOKEN", + "Store GitHub OAuth tokens for Copilot plan usage.", + "Sign in with GitHub or paste a GitHub OAuth token...", + )), + ProviderId::Kimi => Some(Self { title: "Web sessions", - subtitle: "Store labeled Kimi kimi-auth web sessions.", - placeholder: "kimi-auth value or Cookie: kimi-auth=...", - injection: TokenInjection::CookieHeader, - requires_manual_cookie_source: true, - cookie_name: Some("kimi-auth"), - }), - ProviderId::Doubao => Some(TokenAccountSupport { - title: "Ark API keys", - subtitle: "Store labeled Volcengine Ark API keys.", - placeholder: "Ark API key", - injection: TokenInjection::Environment { - key: "ARK_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), - ProviderId::OpenCodeGo => Some(TokenAccountSupport { - title: "API keys or sessions", - subtitle: "Store labeled OpenCode Go API keys or Cookie headers.", - placeholder: "API key or Cookie: ...", + ..Self::cookie( + "Store labeled Kimi kimi-auth web sessions.", + "kimi-auth value or Cookie: kimi-auth=...", + ) + .with_cookie_name("kimi-auth") + }), + ProviderId::Doubao => Some(Self::api_key( + "Ark API keys", + "ARK_API_KEY", + "Store labeled Volcengine Ark API keys.", + "Ark API key", + )), + ProviderId::OpenCodeGo => Some(Self { injection: TokenInjection::EnvironmentOrCookie { key: "OPENCODE_API_KEY".to_string(), }, - requires_manual_cookie_source: false, - cookie_name: None, + ..Self::api_key( + "API keys or sessions", + "OPENCODE_API_KEY", + "Store labeled OpenCode Go API keys or Cookie headers.", + "API key or Cookie: ...", + ) }), // Upstream 0.67: labeled Aixy API keys via token accounts. - ProviderId::Aixy => Some(TokenAccountSupport { - title: "API keys", - subtitle: "Store multiple Aixy API keys.", - placeholder: "Paste Aixy API key…", - injection: TokenInjection::Environment { - key: "AIXY_API_KEY".to_string(), - }, - requires_manual_cookie_source: false, - cookie_name: None, - }), + ProviderId::Aixy => Some(Self::api_key( + "API keys", + "AIXY_API_KEY", + "Store multiple Aixy API keys.", + "Paste Aixy API key…", + )), // These providers don't support token accounts ProviderId::Codex | ProviderId::Pi @@ -669,12 +586,6 @@ impl ProviderAccountData { self.accounts.get(self.clamped_active_index()) } - /// Get the active account mutably - pub fn active_account_mut(&mut self) -> Option<&mut TokenAccount> { - let idx = self.clamped_active_index(); - self.accounts.get_mut(idx) - } - /// Add a new account pub fn add_account(&mut self, account: TokenAccount) { self.accounts.push(account); @@ -706,11 +617,6 @@ impl ProviderAccountData { } } - /// Check if this provider has multiple accounts - pub fn has_multiple(&self) -> bool { - self.accounts.len() > 1 - } - /// Get account count pub fn count(&self) -> usize { self.accounts.len() @@ -805,15 +711,6 @@ impl TokenAccountStore { Ok(()) } - /// Ensure the accounts file exists - pub fn ensure_exists(&self) -> Result { - if self.file_path.exists() { - return Ok(self.file_path.clone()); - } - self.save(&HashMap::new())?; - Ok(self.file_path.clone()) - } - /// Load accounts for a specific provider pub fn load_provider( &self, @@ -893,282 +790,5 @@ impl TokenAccountOverride { } } -/// Maximum number of accounts to fetch per provider -pub const MAX_ACCOUNTS_PER_FETCH: usize = 6; - #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_token_account_support() { - assert!(TokenAccountSupport::is_supported(ProviderId::Claude)); - assert!(TokenAccountSupport::is_supported(ProviderId::Cursor)); - assert!(TokenAccountSupport::is_supported(ProviderId::Copilot)); - assert!(TokenAccountSupport::is_supported(ProviderId::OpenRouter)); - assert!(TokenAccountSupport::is_supported(ProviderId::Grok)); - assert!(TokenAccountSupport::is_supported(ProviderId::Kimi)); - assert!(TokenAccountSupport::is_supported(ProviderId::Doubao)); - assert!(TokenAccountSupport::is_supported(ProviderId::OpenCodeGo)); - assert!(TokenAccountSupport::is_supported(ProviderId::Aixy)); - assert!(!TokenAccountSupport::is_supported(ProviderId::Codex)); - assert!(!TokenAccountSupport::is_supported(ProviderId::Gemini)); - assert!(!TokenAccountSupport::is_supported(ProviderId::Hyper)); - assert!(!TokenAccountSupport::is_supported(ProviderId::GitKraken)); - assert!(!TokenAccountSupport::is_supported(ProviderId::Bifrost)); - } - - #[test] - fn upstream_account_sources_normalize_and_classify_selected_credentials() { - assert_eq!( - TokenAccountSupport::normalized_cookie_header(ProviderId::Kimi, "selected-session"), - "kimi-auth=selected-session" - ); - assert_eq!( - TokenAccountSupport::account_kind(ProviderId::Kimi, "selected-session"), - TokenAccountKind::Cookie - ); - assert_eq!( - TokenAccountSupport::account_kind(ProviderId::Doubao, "ark-key"), - TokenAccountKind::ApiKey - ); - assert_eq!( - TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, "opencode-key"), - TokenAccountKind::ApiKey - ); - assert_eq!( - TokenAccountSupport::account_kind( - ProviderId::OpenCodeGo, - "Cookie: session=opencode-session" - ), - TokenAccountKind::Cookie - ); - assert_eq!( - TokenAccountSupport::env_override(ProviderId::OpenCodeGo, "opencode-key") - .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) - .as_deref(), - Some("opencode-key") - ); - assert!( - TokenAccountSupport::env_override( - ProviderId::OpenCodeGo, - "Cookie: session=opencode-session" - ) - .is_none() - ); - for malformed in ["", " ", "Cookie: broken", "auth=fixture", "two words"] { - assert_eq!( - TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, malformed), - TokenAccountKind::Cookie - ); - assert!(TokenAccountSupport::env_override(ProviderId::OpenCodeGo, malformed).is_none()); - } - assert_eq!( - TokenAccountSupport::env_override(ProviderId::OpenCodeGo, " 'go_key' ") - .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) - .as_deref(), - Some("go_key") - ); - } - - #[test] - fn selected_account_effective_source_normalization() { - let cases = [ - ( - ProviderId::Kimi, - "kimi-session", - SourceMode::Auto, - Some(SourceMode::Web), - ), - ( - ProviderId::Kimi, - "kimi-session", - SourceMode::OAuth, - Some(SourceMode::Web), - ), - ( - ProviderId::Kimi, - "kimi-session", - SourceMode::Cli, - Some(SourceMode::Web), - ), - ( - ProviderId::Doubao, - "ark-key", - SourceMode::Cli, - Some(SourceMode::OAuth), - ), - ( - ProviderId::Doubao, - "ark-key", - SourceMode::Web, - Some(SourceMode::OAuth), - ), - ( - ProviderId::OpenCodeGo, - "Cookie: session=web", - SourceMode::Auto, - Some(SourceMode::Web), - ), - ( - ProviderId::OpenCodeGo, - "Cookie: session=web", - SourceMode::Cli, - Some(SourceMode::Cli), - ), - ( - ProviderId::OpenCodeGo, - "api-key", - SourceMode::Auto, - Some(SourceMode::Auto), - ), - ( - ProviderId::OpenCodeGo, - "api-key", - SourceMode::Web, - Some(SourceMode::Web), - ), - ( - ProviderId::OpenCodeGo, - "api-key", - SourceMode::Cli, - Some(SourceMode::Cli), - ), - (ProviderId::OpenRouter, "api-key", SourceMode::Auto, None), - ]; - - for (provider, token, requested, expected) in cases { - let account = - TokenAccountOverride::from_account(provider, TokenAccount::new("selected", token)); - assert_eq!(account.effective_source_mode(requested), expected); - } - } - - #[test] - fn aixy_token_accounts_inject_api_key_env() { - let support = TokenAccountSupport::for_provider(ProviderId::Aixy).unwrap(); - assert_eq!(support.title, "API keys"); - assert_eq!(support.placeholder, "Paste Aixy API key…"); - assert!(!support.requires_manual_cookie_source); - let env = TokenAccountSupport::env_override(ProviderId::Aixy, "gak_fixture").unwrap(); - assert_eq!( - env.get("AIXY_API_KEY").map(String::as_str), - Some("gak_fixture") - ); - } - - #[test] - fn grok_token_accounts_route_bearer_and_cookie_credentials() { - let bearer = - TokenAccountSupport::env_override(ProviderId::Grok, "Bearer oauth-token").unwrap(); - assert_eq!( - bearer.get("CODEXBAR_GROK_OAUTH_TOKEN").map(String::as_str), - Some("oauth-token") - ); - assert!(TokenAccountSupport::env_override(ProviderId::Grok, "Cookie: sso=abc").is_none()); - assert_eq!( - TokenAccountSupport::normalized_cookie_header(ProviderId::Grok, "Cookie: sso=abc"), - "Cookie: sso=abc" - ); - } - - #[test] - fn openrouter_token_accounts_inject_api_key_env() { - let support = TokenAccountSupport::for_provider(ProviderId::OpenRouter).unwrap(); - assert_eq!(support.title, "API keys"); - assert_eq!(support.placeholder, "sk-or-v1-..."); - assert!(!support.requires_manual_cookie_source); - match &support.injection { - TokenInjection::Environment { key } => assert_eq!(key, "OPENROUTER_API_KEY"), - other => panic!("expected environment injection, got {other:?}"), - } - - let mut data = ProviderAccountData::new(); - data.add_account(TokenAccount::new("Personal", "sk-or-v1-personal")); - data.add_account(TokenAccount::new("Work", "sk-or-v1-work")); - data.set_active(1); - - let active = data.active_account().unwrap(); - assert_eq!(active.label, "Work"); - let env = TokenAccountSupport::env_override(ProviderId::OpenRouter, &active.token).unwrap(); - assert_eq!( - env.get("OPENROUTER_API_KEY").map(String::as_str), - Some("sk-or-v1-work") - ); - - let override_data = - TokenAccountOverride::from_account(ProviderId::OpenRouter, active.clone()); - assert_eq!( - override_data - .env_override - .as_ref() - .and_then(|m| m.get("OPENROUTER_API_KEY")) - .map(String::as_str), - Some("sk-or-v1-work") - ); - assert!(override_data.cookie_header.is_none()); - } - - #[test] - fn test_claude_oauth_detection() { - assert!(TokenAccountSupport::is_claude_oauth_token( - "sk-ant-oat01-abc123" - )); - assert!(TokenAccountSupport::is_claude_oauth_token( - "Bearer sk-ant-oat01-abc123" - )); - assert!(!TokenAccountSupport::is_claude_oauth_token( - "sessionKey=abc123" - )); - assert!(!TokenAccountSupport::is_claude_oauth_token( - "Cookie: foo=bar" - )); - } - - #[test] - fn test_normalize_cookie_header() { - let header = - TokenAccountSupport::normalized_cookie_header(ProviderId::Claude, "abc123token"); - assert_eq!(header, "sessionKey=abc123token"); - - let header = TokenAccountSupport::normalized_cookie_header( - ProviderId::Claude, - "sessionKey=already_formatted", - ); - assert_eq!(header, "sessionKey=already_formatted"); - - let header = TokenAccountSupport::normalized_cookie_header(ProviderId::Ollama, "abc123"); - assert_eq!(header, "__Secure-session=abc123"); - } - - #[test] - fn test_provider_account_data() { - let mut data = ProviderAccountData::new(); - assert_eq!(data.clamped_active_index(), 0); - assert!(data.active_account().is_none()); - - let account = TokenAccount::new("Test", "token123"); - let id = account.id; - data.add_account(account); - - assert_eq!(data.count(), 1); - assert!(data.active_account().is_some()); - assert_eq!(data.active_account().unwrap().label, "Test"); - - data.remove_account(id); - assert_eq!(data.count(), 0); - } - - #[test] - fn test_multiple_accounts() { - let mut data = ProviderAccountData::new(); - data.add_account(TokenAccount::new("Account 1", "token1")); - data.add_account(TokenAccount::new("Account 2", "token2")); - - assert!(data.has_multiple()); - assert_eq!(data.active_account().unwrap().label, "Account 1"); - - data.set_active(1); - assert_eq!(data.active_account().unwrap().label, "Account 2"); - } -} +mod tests; diff --git a/rust/src/core/token_accounts/tests.rs b/rust/src/core/token_accounts/tests.rs new file mode 100644 index 0000000000..8f4e2d900d --- /dev/null +++ b/rust/src/core/token_accounts/tests.rs @@ -0,0 +1,339 @@ +use super::*; + +/// Pins every provider's token-account metadata so table refactors stay byte-identical. +#[test] +fn for_provider_table_is_pinned() { + let rows: Vec = ProviderId::all() + .iter() + .filter_map(|&provider| { + let support = TokenAccountSupport::for_provider(provider)?; + let injection = match &support.injection { + TokenInjection::CookieHeader => "cookie".to_string(), + TokenInjection::Environment { key } => format!("env:{key}"), + TokenInjection::EnvironmentOrCookie { key } => format!("env_or_cookie:{key}"), + }; + Some(format!( + "{provider:?}|{}|{}|{}|{injection}|{}|{:?}", + support.title, + support.subtitle, + support.placeholder, + support.requires_manual_cookie_source, + support.cookie_name + )) + }) + .collect(); + let mut expected: Vec<&str> = EXPECTED_SUPPORT_ROWS.to_vec(); + let mut actual: Vec<&str> = rows.iter().map(String::as_str).collect(); + expected.sort_unstable(); + actual.sort_unstable(); + assert_eq!(actual, expected); +} + +const EXPECTED_SUPPORT_ROWS: &[&str] = &[ + "Claude|Session tokens|Store Claude sessionKey cookies for settings-page usage. OAuth tokens are kept as a legacy fallback.|Paste sessionKey value or Cookie: sessionKey=...|cookie|true|Some(\"sessionKey\")", + "Zai|API tokens|Stored locally in token-accounts.json. Team usage can use workspace_id as organization|project.|Paste token...|env:Z_AI_API_KEY|false|None", + "Cursor|Session tokens|Store multiple Cursor Cookie headers.|Cookie: ...|cookie|true|None", + "OpenCode|Session tokens|Store multiple OpenCode Cookie headers.|Cookie: ...|cookie|true|None", + "Factory|Session tokens|Store multiple Factory Cookie headers.|Cookie: ...|cookie|true|None", + "Alibaba|Session tokens|Store multiple Alibaba Cookie headers.|Cookie: ...|cookie|true|None", + "AlibabaTokenPlan|Session tokens|Store multiple Alibaba Token Plan Cookie headers.|Cookie: cna=...; login_aliyunid_csrf=...|cookie|true|None", + "MiniMax|Session tokens|Store multiple MiniMax Cookie headers.|Cookie: ...|cookie|true|None", + "Augment|Session tokens|Store multiple Augment Cookie headers.|Cookie: ...|cookie|true|None", + "Amp|Session tokens|Store multiple Amp Cookie headers.|Cookie: ...|cookie|true|None", + "Ollama|Session tokens|Store multiple Ollama Cookie headers or __Secure-session values.|__Secure-session value or Cookie: ...|cookie|true|Some(\"__Secure-session\")", + "T3Chat|Session tokens|Store multiple T3 Chat Cookie headers or full browser cURL captures.|Cookie: ... or curl ... -H 'Cookie: ...'|cookie|true|None", + "ZoomMate|Session tokens|Store multiple ZoomMate Cookie headers or credits/status cURL captures.|Cookie: ... or curl 'https://ai.zoom.us/.../credits/status' -H 'Authorization: Bearer ...'|cookie|true|None", + "Mistral|Session tokens|Store multiple Mistral Cookie headers.|Cookie: ...|cookie|true|None", + "Manus|Session tokens|Store multiple Manus session_id values.|session_id value or Cookie: ...|cookie|true|Some(\"session_id\")", + "MiMo|Session tokens|Store multiple Xiaomi MiMo Cookie headers.|Cookie: api-platform_serviceToken=...; userId=...|cookie|true|None", + "CommandCode|Session tokens|Store multiple Command Code Cookie headers or Better Auth values.|Cookie: __Secure-commandcode_prod_.session_token=... or better-auth value|cookie|true|Some(\"__Secure-better-auth.session_token\")", + "Qoder|Session tokens|Store multiple Qoder Cookie headers.|Cookie: ...|cookie|true|None", + "CodeBuddy|Session tokens|Store CodeBuddy CN Cookie headers (from plans-usage DevTools cURL).|Cookie: session=...; ... (or paste full Cookie header)|cookie|true|None", + "Sakana|Session tokens|Store multiple Sakana Console Cookie headers.|Cookie: ...|cookie|true|None", + "Notion|Session tokens|Store multiple Notion Cookie headers or token_v2 values.|Cookie: token_v2=... or paste the token_v2 value|cookie|true|Some(\"token_v2\")", + "Replicate|Session tokens|Store multiple Replicate Cookie headers from the billing page.|Cookie: sessionid=...; ...|cookie|true|Some(\"sessionid\")", + "Sub2Api|Group API keys|Store multiple sub2api group API keys with labels such as Claude, Codex, or Gemini.|sk-...|env:SUB2API_API_KEY|false|None", + "DeepInfra|API keys|Store multiple DeepInfra API keys.|API key from deepinfra.com/dash|env:DEEPINFRA_API_KEY|false|None", + "HuggingFace|API tokens|Store multiple Hugging Face access tokens.|Paste a Hugging Face access token|env:CODEXBAR_HUGGINGFACE_API_KEY|false|None", + "AiAnd|API keys|Store multiple ai& API keys.|API key from console.aiand.com|env:AIAND_API_KEY|false|None", + "ZenMux|API keys|Store multiple ZenMux Management API keys.|Management API key|env:ZENMUX_MANAGEMENT_API_KEY|false|None", + "ClinePass|API keys|Store multiple ClinePass API keys. Without one, CodexBar reads your existing Cline session (run cline auth) without copying it.|API key|env:CLINE_API_KEY|false|None", + "Neuralwatt|API keys|Store multiple Neuralwatt API keys.|API key|env:NEURALWATT_API_KEY|false|None", + "Grok|Grok credentials|Store SuperGrok bearer tokens or grok.com Cookie headers.|Bearer token or Cookie: ...|cookie|false|None", + "Xai|Management API keys|Store multiple xAI Management API keys. Team ID is set separately under provider settings.|xai-... Management API key from console.x.ai|env:XAI_MANAGEMENT_API_KEY|false|None", + "OpenRouter|API keys|Store multiple OpenRouter API keys.|sk-or-v1-...|env:OPENROUTER_API_KEY|false|None", + "Copilot|GitHub accounts|Store GitHub OAuth tokens for Copilot plan usage.|Sign in with GitHub or paste a GitHub OAuth token...|env:GITHUB_TOKEN|false|None", + "Kimi|Web sessions|Store labeled Kimi kimi-auth web sessions.|kimi-auth value or Cookie: kimi-auth=...|cookie|true|Some(\"kimi-auth\")", + "Doubao|Ark API keys|Store labeled Volcengine Ark API keys.|Ark API key|env:ARK_API_KEY|false|None", + "OpenCodeGo|API keys or sessions|Store labeled OpenCode Go API keys or Cookie headers.|API key or Cookie: ...|env_or_cookie:OPENCODE_API_KEY|false|None", + "Aixy|API keys|Store multiple Aixy API keys.|Paste Aixy API key…|env:AIXY_API_KEY|false|None", +]; + +#[test] +fn test_token_account_support() { + assert!(TokenAccountSupport::is_supported(ProviderId::Claude)); + assert!(TokenAccountSupport::is_supported(ProviderId::Cursor)); + assert!(TokenAccountSupport::is_supported(ProviderId::Copilot)); + assert!(TokenAccountSupport::is_supported(ProviderId::OpenRouter)); + assert!(TokenAccountSupport::is_supported(ProviderId::Grok)); + assert!(TokenAccountSupport::is_supported(ProviderId::Kimi)); + assert!(TokenAccountSupport::is_supported(ProviderId::Doubao)); + assert!(TokenAccountSupport::is_supported(ProviderId::OpenCodeGo)); + assert!(TokenAccountSupport::is_supported(ProviderId::Aixy)); + assert!(!TokenAccountSupport::is_supported(ProviderId::Codex)); + assert!(!TokenAccountSupport::is_supported(ProviderId::Gemini)); + assert!(!TokenAccountSupport::is_supported(ProviderId::Hyper)); + assert!(!TokenAccountSupport::is_supported(ProviderId::GitKraken)); + assert!(!TokenAccountSupport::is_supported(ProviderId::Bifrost)); +} + +#[test] +fn upstream_account_sources_normalize_and_classify_selected_credentials() { + assert_eq!( + TokenAccountSupport::normalized_cookie_header(ProviderId::Kimi, "selected-session"), + "kimi-auth=selected-session" + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::Kimi, "selected-session"), + TokenAccountKind::Cookie + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::Doubao, "ark-key"), + TokenAccountKind::ApiKey + ); + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, "opencode-key"), + TokenAccountKind::ApiKey + ); + assert_eq!( + TokenAccountSupport::account_kind( + ProviderId::OpenCodeGo, + "Cookie: session=opencode-session" + ), + TokenAccountKind::Cookie + ); + assert_eq!( + TokenAccountSupport::env_override(ProviderId::OpenCodeGo, "opencode-key") + .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) + .as_deref(), + Some("opencode-key") + ); + assert!( + TokenAccountSupport::env_override( + ProviderId::OpenCodeGo, + "Cookie: session=opencode-session" + ) + .is_none() + ); + for malformed in ["", " ", "Cookie: broken", "auth=fixture", "two words"] { + assert_eq!( + TokenAccountSupport::account_kind(ProviderId::OpenCodeGo, malformed), + TokenAccountKind::Cookie + ); + assert!(TokenAccountSupport::env_override(ProviderId::OpenCodeGo, malformed).is_none()); + } + assert_eq!( + TokenAccountSupport::env_override(ProviderId::OpenCodeGo, " 'go_key' ") + .and_then(|env| env.get("OPENCODE_API_KEY").cloned()) + .as_deref(), + Some("go_key") + ); +} + +#[test] +fn selected_account_effective_source_normalization() { + let cases = [ + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::Auto, + Some(SourceMode::Web), + ), + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::OAuth, + Some(SourceMode::Web), + ), + ( + ProviderId::Kimi, + "kimi-session", + SourceMode::Cli, + Some(SourceMode::Web), + ), + ( + ProviderId::Doubao, + "ark-key", + SourceMode::Cli, + Some(SourceMode::OAuth), + ), + ( + ProviderId::Doubao, + "ark-key", + SourceMode::Web, + Some(SourceMode::OAuth), + ), + ( + ProviderId::OpenCodeGo, + "Cookie: session=web", + SourceMode::Auto, + Some(SourceMode::Web), + ), + ( + ProviderId::OpenCodeGo, + "Cookie: session=web", + SourceMode::Cli, + Some(SourceMode::Cli), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Auto, + Some(SourceMode::Auto), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Web, + Some(SourceMode::Web), + ), + ( + ProviderId::OpenCodeGo, + "api-key", + SourceMode::Cli, + Some(SourceMode::Cli), + ), + (ProviderId::OpenRouter, "api-key", SourceMode::Auto, None), + ]; + + for (provider, token, requested, expected) in cases { + let account = + TokenAccountOverride::from_account(provider, TokenAccount::new("selected", token)); + assert_eq!(account.effective_source_mode(requested), expected); + } +} + +#[test] +fn aixy_token_accounts_inject_api_key_env() { + let support = TokenAccountSupport::for_provider(ProviderId::Aixy).unwrap(); + assert_eq!(support.title, "API keys"); + assert_eq!(support.placeholder, "Paste Aixy API key…"); + assert!(!support.requires_manual_cookie_source); + let env = TokenAccountSupport::env_override(ProviderId::Aixy, "gak_fixture").unwrap(); + assert_eq!( + env.get("AIXY_API_KEY").map(String::as_str), + Some("gak_fixture") + ); +} + +#[test] +fn grok_token_accounts_route_bearer_and_cookie_credentials() { + let bearer = TokenAccountSupport::env_override(ProviderId::Grok, "Bearer oauth-token").unwrap(); + assert_eq!( + bearer.get("CODEXBAR_GROK_OAUTH_TOKEN").map(String::as_str), + Some("oauth-token") + ); + assert!(TokenAccountSupport::env_override(ProviderId::Grok, "Cookie: sso=abc").is_none()); + assert_eq!( + TokenAccountSupport::normalized_cookie_header(ProviderId::Grok, "Cookie: sso=abc"), + "Cookie: sso=abc" + ); +} + +#[test] +fn openrouter_token_accounts_inject_api_key_env() { + let support = TokenAccountSupport::for_provider(ProviderId::OpenRouter).unwrap(); + assert_eq!(support.title, "API keys"); + assert_eq!(support.placeholder, "sk-or-v1-..."); + assert!(!support.requires_manual_cookie_source); + match &support.injection { + TokenInjection::Environment { key } => assert_eq!(key, "OPENROUTER_API_KEY"), + other => panic!("expected environment injection, got {other:?}"), + } + + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Personal", "sk-or-v1-personal")); + data.add_account(TokenAccount::new("Work", "sk-or-v1-work")); + data.set_active(1); + + let active = data.active_account().unwrap(); + assert_eq!(active.label, "Work"); + let env = TokenAccountSupport::env_override(ProviderId::OpenRouter, &active.token).unwrap(); + assert_eq!( + env.get("OPENROUTER_API_KEY").map(String::as_str), + Some("sk-or-v1-work") + ); + + let override_data = TokenAccountOverride::from_account(ProviderId::OpenRouter, active.clone()); + assert_eq!( + override_data + .env_override + .as_ref() + .and_then(|m| m.get("OPENROUTER_API_KEY")) + .map(String::as_str), + Some("sk-or-v1-work") + ); + assert!(override_data.cookie_header.is_none()); +} + +#[test] +fn test_claude_oauth_detection() { + assert!(TokenAccountSupport::is_claude_oauth_token( + "sk-ant-oat01-abc123" + )); + assert!(TokenAccountSupport::is_claude_oauth_token( + "Bearer sk-ant-oat01-abc123" + )); + assert!(!TokenAccountSupport::is_claude_oauth_token( + "sessionKey=abc123" + )); + assert!(!TokenAccountSupport::is_claude_oauth_token( + "Cookie: foo=bar" + )); +} + +#[test] +fn test_normalize_cookie_header() { + let header = TokenAccountSupport::normalized_cookie_header(ProviderId::Claude, "abc123token"); + assert_eq!(header, "sessionKey=abc123token"); + + let header = TokenAccountSupport::normalized_cookie_header( + ProviderId::Claude, + "sessionKey=already_formatted", + ); + assert_eq!(header, "sessionKey=already_formatted"); + + let header = TokenAccountSupport::normalized_cookie_header(ProviderId::Ollama, "abc123"); + assert_eq!(header, "__Secure-session=abc123"); +} + +#[test] +fn test_provider_account_data() { + let mut data = ProviderAccountData::new(); + assert_eq!(data.clamped_active_index(), 0); + assert!(data.active_account().is_none()); + + let account = TokenAccount::new("Test", "token123"); + let id = account.id; + data.add_account(account); + + assert_eq!(data.count(), 1); + assert!(data.active_account().is_some()); + assert_eq!(data.active_account().unwrap().label, "Test"); + + data.remove_account(id); + assert_eq!(data.count(), 0); +} + +#[test] +fn test_multiple_accounts() { + let mut data = ProviderAccountData::new(); + data.add_account(TokenAccount::new("Account 1", "token1")); + data.add_account(TokenAccount::new("Account 2", "token2")); + + assert_eq!(data.count(), 2); + assert_eq!(data.active_account().unwrap().label, "Account 1"); + + data.set_active(1); + assert_eq!(data.active_account().unwrap().label, "Account 2"); +} diff --git a/rust/src/core/usage_pace.rs b/rust/src/core/usage_pace.rs index c48057f01c..b79b893b1d 100755 --- a/rust/src/core/usage_pace.rs +++ b/rust/src/core/usage_pace.rs @@ -48,22 +48,6 @@ impl PaceStage { PaceStage::SlightlyBehind | PaceStage::Behind | PaceStage::FarBehind => "🐢", } } - - /// Whether the user is consuming faster than expected - pub fn is_ahead(&self) -> bool { - matches!( - self, - PaceStage::SlightlyAhead | PaceStage::Ahead | PaceStage::FarAhead - ) - } - - /// Whether the user is consuming slower than expected - pub fn is_behind(&self) -> bool { - matches!( - self, - PaceStage::SlightlyBehind | PaceStage::Behind | PaceStage::FarBehind - ) - } } /// Usage pace prediction result @@ -265,7 +249,10 @@ mod tests { let window = RateWindow::with_details(80.0, Some(10080), Some(resets_at), None); let pace = UsagePace::weekly(&window, Some(now), 10080).unwrap(); - assert!(pace.stage.is_ahead()); + assert!(matches!( + pace.stage, + PaceStage::SlightlyAhead | PaceStage::Ahead | PaceStage::FarAhead + )); assert!(pace.delta_percent > 0.0); } diff --git a/rust/src/core/widget_snapshot.rs b/rust/src/core/widget_snapshot.rs deleted file mode 100755 index c7e8eaf41a..0000000000 --- a/rust/src/core/widget_snapshot.rs +++ /dev/null @@ -1,372 +0,0 @@ -//! Widget Snapshot -//! -//! Data export structures for widgets and external integrations. -//! Provides a serializable snapshot of all provider usage data. - -#![allow( - dead_code, - reason = "widget snapshot types reserved for future UI integration" -)] - -use crate::core::{ProviderId, RateWindow}; -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use std::fs; -use std::path::PathBuf; - -/// Token usage summary for a provider -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct TokenUsageSummary { - /// Session cost in USD - #[serde(skip_serializing_if = "Option::is_none")] - pub session_cost_usd: Option, - /// Session token count - #[serde(skip_serializing_if = "Option::is_none")] - pub session_tokens: Option, - /// Last 30 days cost in USD - #[serde(skip_serializing_if = "Option::is_none")] - pub last_30_days_cost_usd: Option, - /// Last 30 days token count - #[serde(skip_serializing_if = "Option::is_none")] - pub last_30_days_tokens: Option, -} - -impl TokenUsageSummary { - pub fn new() -> Self { - Self { - session_cost_usd: None, - session_tokens: None, - last_30_days_cost_usd: None, - last_30_days_tokens: None, - } - } - - pub fn with_session(mut self, cost_usd: f64, tokens: i64) -> Self { - self.session_cost_usd = Some(cost_usd); - self.session_tokens = Some(tokens); - self - } - - pub fn with_last_30_days(mut self, cost_usd: f64, tokens: i64) -> Self { - self.last_30_days_cost_usd = Some(cost_usd); - self.last_30_days_tokens = Some(tokens); - self - } -} - -impl Default for TokenUsageSummary { - fn default() -> Self { - Self::new() - } -} - -/// Daily usage data point -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct DailyUsagePoint { - /// Day key in yyyy-MM-dd format - pub day_key: String, - /// Total tokens used - #[serde(skip_serializing_if = "Option::is_none")] - pub total_tokens: Option, - /// Cost in USD - #[serde(skip_serializing_if = "Option::is_none")] - pub cost_usd: Option, -} - -impl DailyUsagePoint { - pub fn new(day_key: impl Into) -> Self { - Self { - day_key: day_key.into(), - total_tokens: None, - cost_usd: None, - } - } - - pub fn with_tokens(mut self, tokens: i64) -> Self { - self.total_tokens = Some(tokens); - self - } - - pub fn with_cost(mut self, cost_usd: f64) -> Self { - self.cost_usd = Some(cost_usd); - self - } -} - -/// Widget entry for a single provider -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WidgetProviderEntry { - /// Provider identifier - pub provider: ProviderId, - /// When this data was last updated - pub updated_at: DateTime, - /// Primary rate limit - #[serde(skip_serializing_if = "Option::is_none")] - pub primary: Option, - /// Secondary rate limit - #[serde(skip_serializing_if = "Option::is_none")] - pub secondary: Option, - /// Tertiary rate limit - #[serde(skip_serializing_if = "Option::is_none")] - pub tertiary: Option, - /// Credits remaining - #[serde(skip_serializing_if = "Option::is_none")] - pub credits_remaining: Option, - /// Code review remaining percentage (0-100) - #[serde(skip_serializing_if = "Option::is_none")] - pub code_review_remaining_percent: Option, - /// Token usage summary - #[serde(skip_serializing_if = "Option::is_none")] - pub token_usage: Option, - /// Daily usage data points - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub daily_usage: Vec, - /// Account email if available - #[serde(skip_serializing_if = "Option::is_none")] - pub account_email: Option, - /// Login method/plan info (e.g., "Claude Pro", "Claude Max") - #[serde(skip_serializing_if = "Option::is_none")] - pub login_method: Option, -} - -impl WidgetProviderEntry { - pub fn new(provider: ProviderId, updated_at: DateTime) -> Self { - Self { - provider, - updated_at, - primary: None, - secondary: None, - tertiary: None, - credits_remaining: None, - code_review_remaining_percent: None, - token_usage: None, - daily_usage: Vec::new(), - account_email: None, - login_method: None, - } - } - - pub fn with_primary(mut self, rate: RateWindow) -> Self { - self.primary = Some(rate); - self - } - - pub fn with_secondary(mut self, rate: RateWindow) -> Self { - self.secondary = Some(rate); - self - } - - pub fn with_tertiary(mut self, rate: RateWindow) -> Self { - self.tertiary = Some(rate); - self - } - - pub fn with_credits_remaining(mut self, credits: f64) -> Self { - self.credits_remaining = Some(credits); - self - } - - pub fn with_code_review_remaining(mut self, percent: f64) -> Self { - self.code_review_remaining_percent = Some(percent); - self - } - - pub fn with_token_usage(mut self, summary: TokenUsageSummary) -> Self { - self.token_usage = Some(summary); - self - } - - pub fn with_daily_usage(mut self, points: Vec) -> Self { - self.daily_usage = points; - self - } - - pub fn with_account_email(mut self, email: impl Into) -> Self { - self.account_email = Some(email.into()); - self - } - - pub fn with_login_method(mut self, method: impl Into) -> Self { - self.login_method = Some(method.into()); - self - } -} - -/// Complete widget snapshot with all provider data -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct WidgetSnapshot { - /// Provider entries - pub entries: Vec, - /// List of enabled providers - pub enabled_providers: Vec, - /// When this snapshot was generated - pub generated_at: DateTime, -} - -impl WidgetSnapshot { - pub fn new(entries: Vec, generated_at: DateTime) -> Self { - let enabled_providers = entries.iter().map(|e| e.provider).collect(); - Self { - entries, - enabled_providers, - generated_at, - } - } - - pub fn with_enabled_providers(mut self, providers: Vec) -> Self { - self.enabled_providers = providers; - self - } - - /// Get entry for a specific provider - pub fn entry_for(&self, provider: ProviderId) -> Option<&WidgetProviderEntry> { - self.entries.iter().find(|e| e.provider == provider) - } - - /// Check if a provider is enabled - pub fn is_enabled(&self, provider: ProviderId) -> bool { - self.enabled_providers.contains(&provider) - } -} - -/// Widget snapshot store -pub struct WidgetSnapshotStore; - -impl WidgetSnapshotStore { - const FILENAME: &'static str = "widget-snapshot.json"; - - /// Load widget snapshot from disk - pub fn load() -> Option { - let path = Self::snapshot_path()?; - let data = fs::read_to_string(&path).ok()?; - serde_json::from_str(&data).ok() - } - - /// Save widget snapshot to disk - pub fn save(snapshot: &WidgetSnapshot) -> Result<(), WidgetSnapshotError> { - let path = Self::snapshot_path().ok_or(WidgetSnapshotError::PathNotAvailable)?; - - if let Some(parent) = path.parent() { - fs::create_dir_all(parent)?; - } - - let json = serde_json::to_string_pretty(snapshot)?; - fs::write(&path, json)?; - - tracing::debug!("Saved widget snapshot to {:?}", path); - Ok(()) - } - - /// Clear widget snapshot - pub fn clear() { - if let Some(path) = Self::snapshot_path() { - // Best-effort cleanup; a missing snapshot file is not an error. - let _cleared_snapshot = fs::remove_file(path); - } - } - - fn snapshot_path() -> Option { - dirs::data_local_dir().map(|d| d.join("CodexBar").join(Self::FILENAME)) - } -} - -/// Widget selection store for tracking user's selected provider -pub struct WidgetSelectionStore; - -impl WidgetSelectionStore { - const FILENAME: &'static str = "widget-selection.json"; - - /// Load selected provider - pub fn load_selected_provider() -> Option { - let path = Self::selection_path()?; - let data = fs::read_to_string(&path).ok()?; - - #[derive(Deserialize)] - struct Selection { - provider: ProviderId, - } - - let selection: Selection = serde_json::from_str(&data).ok()?; - Some(selection.provider) - } - - /// Save selected provider - pub fn save_selected_provider(provider: ProviderId) -> Result<(), WidgetSnapshotError> { - let path = Self::selection_path().ok_or(WidgetSnapshotError::PathNotAvailable)?; - - if let Some(parent) = path.parent() { - fs::create_dir_all(parent)?; - } - - #[derive(Serialize)] - struct Selection { - provider: ProviderId, - } - - let json = serde_json::to_string(&Selection { provider })?; - fs::write(&path, json)?; - Ok(()) - } - - fn selection_path() -> Option { - dirs::data_local_dir().map(|d| d.join("CodexBar").join(Self::FILENAME)) - } -} - -/// Widget snapshot errors -#[derive(Debug, thiserror::Error)] -pub enum WidgetSnapshotError { - #[error("Path not available")] - PathNotAvailable, - #[error("IO error: {0}")] - Io(#[from] std::io::Error), - #[error("JSON error: {0}")] - Json(#[from] serde_json::Error), -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_token_usage_summary() { - let summary = TokenUsageSummary::new() - .with_session(1.50, 10000) - .with_last_30_days(45.00, 300000); - - assert_eq!(summary.session_cost_usd, Some(1.50)); - assert_eq!(summary.session_tokens, Some(10000)); - assert_eq!(summary.last_30_days_cost_usd, Some(45.00)); - assert_eq!(summary.last_30_days_tokens, Some(300000)); - } - - #[test] - fn test_daily_usage_point() { - let point = DailyUsagePoint::new("2026-01-15") - .with_tokens(5000) - .with_cost(0.75); - - assert_eq!(point.day_key, "2026-01-15"); - assert_eq!(point.total_tokens, Some(5000)); - assert_eq!(point.cost_usd, Some(0.75)); - } - - #[test] - fn test_widget_provider_entry() { - let entry = - WidgetProviderEntry::new(ProviderId::Claude, Utc::now()).with_credits_remaining(100.0); - - assert_eq!(entry.provider, ProviderId::Claude); - assert_eq!(entry.credits_remaining, Some(100.0)); - } - - #[test] - fn test_widget_snapshot() { - let entry = WidgetProviderEntry::new(ProviderId::Codex, Utc::now()); - let snapshot = WidgetSnapshot::new(vec![entry], Utc::now()); - - assert_eq!(snapshot.entries.len(), 1); - assert!(snapshot.is_enabled(ProviderId::Codex)); - assert!(!snapshot.is_enabled(ProviderId::Claude)); - } -} diff --git a/rust/src/cost_reporting_period.rs b/rust/src/cost_reporting_period.rs index b85e270475..6f28637a8f 100644 --- a/rust/src/cost_reporting_period.rs +++ b/rust/src/cost_reporting_period.rs @@ -55,9 +55,10 @@ static COST_BUCKET_ZONE: RwLock> = RwLock::new(None); /// The zone local cost history is bucketed in (upstream's pinned bucket /// calendar, `tokenCostUsageBucketTimeZone`). /// -/// The desktop app and the CLI apply the saved -/// `Settings::cost_usage_bucket_time_zone` at startup. Until then, and while -/// nothing is pinned, history is bucketed in the machine zone. +/// The saved `Settings::cost_usage_bucket_time_zone` is stored but not yet +/// applied at runtime: nothing in the desktop app or the CLI calls +/// [`set_cost_bucket_zone`] (the wiring PRs #653 and #665 were closed). Until +/// it is wired, history is bucketed in the machine zone. pub fn cost_bucket_zone() -> CostTimeZone { let pinned = *COST_BUCKET_ZONE .read() diff --git a/rust/src/host/command_runner.rs b/rust/src/host/command_runner.rs index 26a8089e20..f39e680620 100755 --- a/rust/src/host/command_runner.rs +++ b/rust/src/host/command_runner.rs @@ -4,11 +4,6 @@ //! On Windows, uses standard process spawning with output capture. //! Designed for running interactive CLI tools like `codex` and `claude`. -#![allow( - dead_code, - reason = "command runner types reserved for future host management integration" -)] - use std::collections::HashMap; use std::io::{BufRead, BufReader}; #[cfg(windows)] @@ -513,46 +508,6 @@ enum StreamEvent { Closed, } -/// Rolling buffer for substring matching across chunk boundaries -pub struct RollingBuffer { - max_needle_len: usize, - tail: Vec, -} - -impl RollingBuffer { - pub fn new(max_needle_len: usize) -> Self { - Self { - max_needle_len, - tail: Vec::with_capacity(max_needle_len), - } - } - - /// Append data and return the combined buffer for searching - pub fn append(&mut self, data: &[u8]) -> Vec { - if data.is_empty() { - return Vec::new(); - } - - let mut combined = Vec::with_capacity(self.tail.len() + data.len()); - combined.extend_from_slice(&self.tail); - combined.extend_from_slice(data); - - // Keep only the tail for next search - if self.max_needle_len > 1 && combined.len() >= self.max_needle_len - 1 { - self.tail = combined[combined.len() - (self.max_needle_len - 1)..].to_vec(); - } else { - self.tail = combined.clone(); - } - - combined - } - - /// Reset the buffer - pub fn reset(&mut self) { - self.tail.clear(); - } -} - #[cfg(test)] mod tests { use super::*; @@ -565,18 +520,6 @@ mod tests { assert_eq!(opts.timeout, Duration::from_secs(20)); } - #[test] - fn test_rolling_buffer() { - let mut buf = RollingBuffer::new(5); - - let result = buf.append(b"hello"); - assert_eq!(result, b"hello"); - - let result = buf.append(b" world"); - // Should include tail from previous - assert!(result.len() > 6); - } - #[test] fn test_command_runner_new() { let runner = CommandRunner::new(); diff --git a/rust/src/host/mod.rs b/rust/src/host/mod.rs index 7ef6d8f22a..3ff55c3e65 100755 --- a/rust/src/host/mod.rs +++ b/rust/src/host/mod.rs @@ -9,6 +9,4 @@ pub mod session; unused_imports, reason = "command runner API re-exported for future CLI integration" )] -pub use command_runner::{ - CommandError, CommandOptions, CommandResult, CommandRunner, RollingBuffer, -}; +pub use command_runner::{CommandError, CommandOptions, CommandResult, CommandRunner}; diff --git a/rust/src/host/session.rs b/rust/src/host/session.rs index 400752e1b4..0a1a5885e0 100644 --- a/rust/src/host/session.rs +++ b/rust/src/host/session.rs @@ -1,28 +1,8 @@ //! Shared host-session helpers reusable across the egui shell and Tauri shell. //! -//! Keeps detection logic for SSH / RDP / remote sessions and primary-monitor +//! Keeps the SSH / RDP launch-block messages and primary-monitor //! work-area queries in the shared crate so shells don't duplicate the logic. -/// Whether the current process is running inside an SSH session. -pub fn is_ssh_session() -> bool { - std::env::var_os("SSH_CONNECTION").is_some() || std::env::var_os("SSH_CLIENT").is_some() -} - -/// Whether the current process is running inside a Windows Remote Desktop session. -#[cfg(windows)] -pub fn is_remote_session() -> bool { - use windows::Win32::UI::WindowsAndMessaging::{GetSystemMetrics, SM_REMOTESESSION}; - // SAFETY: GetSystemMetrics is a query-only Win32 call taking a constant - // metric index; it reads no caller memory and returns a small int. - unsafe { GetSystemMetrics(SM_REMOTESESSION) != 0 } -} - -/// Non-Windows platforms cannot be in a Windows remote-desktop session. -#[cfg(not(windows))] -pub fn is_remote_session() -> bool { - false -} - /// User-facing message explaining why launch is blocked under SSH. pub fn ssh_session_error_message() -> &'static str { "CodexBar can't render its native window from an SSH session on this machine.\n\nOpen it from the logged-in Windows desktop session instead, or use the installed CLI over SSH:\n\n codexbar-cli usage -p claude\n\nThe startup log is written to %TEMP%\\codexbar_launch_.log." @@ -44,11 +24,6 @@ pub fn launch_block_reason(is_ssh: bool, is_remote: bool) -> Option<&'static str } } -/// Detect current launch-block reason by probing environment + OS APIs directly. -pub fn current_launch_block_reason() -> Option<&'static str> { - launch_block_reason(is_ssh_session(), is_remote_session()) -} - /// Primary-monitor work area in physical pixels (excludes the taskbar on Windows). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct WorkAreaPixels { @@ -108,14 +83,6 @@ pub fn gemini_cli_credentials_path() -> Option { dirs::home_dir().map(|home| home.join(".gemini").join("oauth_creds.json")) } -/// `true` when the Gemini CLI's credentials file exists (i.e. the user has -/// signed in via `gemini auth login` locally). -pub fn gemini_cli_signed_in() -> bool { - gemini_cli_credentials_path() - .map(|p| p.exists()) - .unwrap_or(false) -} - /// Filesystem path to VertexAI application-default credentials. Respects the /// `GOOGLE_APPLICATION_CREDENTIALS` env var when set, otherwise falls back to /// the gcloud well-known location under the OS config dir. @@ -132,13 +99,6 @@ pub fn vertexai_credentials_path() -> Option { }) } -/// `true` when VertexAI application-default credentials exist on disk. -pub fn vertexai_signed_in() -> bool { - vertexai_credentials_path() - .map(|p| p.exists()) - .unwrap_or(false) -} - /// Detect JetBrains / Google-IDE configuration directories under the user's /// config home. Returns an empty list if none are present. pub fn jetbrains_detected_ide_paths() -> Vec { diff --git a/rust/src/login.rs b/rust/src/login.rs index 3d7dab6d35..cedb664b97 100755 --- a/rust/src/login.rs +++ b/rust/src/login.rs @@ -2,11 +2,6 @@ //! //! Runs CLI login commands and captures output/URLs -#![allow( - dead_code, - reason = "login flow types reserved for future session management integration" -)] - use regex_lite::Regex; use std::io::{BufRead, BufReader, Read}; #[cfg(windows)] @@ -41,74 +36,6 @@ pub enum LoginPhase { Complete, } -/// Run Claude CLI login -pub async fn run_claude_login(timeout_secs: u64, on_phase: F) -> LoginResult -where - F: Fn(LoginPhase) + Send + 'static, -{ - run_cli_login( - "claude", - &["/login"], - timeout_secs, - on_phase, - &[ - "Successfully logged in", - "Login successful", - "Logged in successfully", - ], - ) - .await -} - -/// Run Codex CLI login -pub async fn run_codex_login(timeout_secs: u64, on_phase: F) -> LoginResult -where - F: Fn(LoginPhase) + Send + 'static, -{ - run_cli_login( - "codex", - &["auth", "login"], - timeout_secs, - on_phase, - &[ - "Successfully logged in", - "Login successful", - "Logged in successfully", - ], - ) - .await -} - -/// Run Gemini/gcloud login -pub async fn run_gemini_login(timeout_secs: u64, on_phase: F) -> LoginResult -where - F: Fn(LoginPhase) + Send + 'static, -{ - run_cli_login( - "gcloud", - &["auth", "login"], - timeout_secs, - on_phase, - &["You are now logged in", "Credentials saved"], - ) - .await -} - -/// Run Copilot/GitHub device flow login -pub async fn run_copilot_login(timeout_secs: u64, on_phase: F) -> LoginResult -where - F: Fn(LoginPhase) + Send + 'static, -{ - run_cli_login( - "gh", - &["auth", "login", "-w"], - timeout_secs, - on_phase, - &["Logged in as", "Authentication complete"], - ) - .await -} - /// Run Kiro CLI login pub async fn run_kiro_login(timeout_secs: u64, on_phase: F) -> LoginResult where @@ -135,25 +62,6 @@ where .await } -/// Generic CLI login runner (resolves binary via PATH) -async fn run_cli_login( - binary: &str, - args: &[&str], - timeout_secs: u64, - on_phase: F, - success_markers: &[&str], -) -> LoginResult -where - F: Fn(LoginPhase) + Send + 'static, -{ - let binary_path = match which::which(binary) { - Ok(p) => p, - Err(_) => return missing_binary_result(binary), - }; - - run_cli_login_path(&binary_path, args, timeout_secs, on_phase, success_markers).await -} - /// Generic CLI login runner (uses a pre-resolved binary path) async fn run_cli_login_path( binary_path: &std::path::Path, @@ -336,9 +244,3 @@ where Err(e) => state.into_result(LoginOutcome::LaunchFailed(e.to_string())), } } - -/// Open a URL in the default browser -pub fn open_auth_url(url: &str) -> anyhow::Result<()> { - open::that(url)?; - Ok(()) -} diff --git a/rust/src/managed_process.rs b/rust/src/managed_process.rs index 24b4d909b2..cb10359171 100644 --- a/rust/src/managed_process.rs +++ b/rust/src/managed_process.rs @@ -165,12 +165,6 @@ impl ManagedProcess { }) } - /// Candidate IPv4 loopback ports the owned child is currently listening on. - /// Providers use this to decide when the child's local service is ready. - pub fn listening_ports(&self) -> ManagedProcessResult> { - listening_ports_for_pid(self.pid) - } - /// Terminate and reap the owned child, bounded by `cleanup_reserve`. pub async fn shutdown(mut self, cleanup_reserve: Duration) { let Some(resources) = self.take_resources() else { @@ -662,94 +656,9 @@ impl Drop for Attributes { } } -/// NUL-terminated `CreateProcessW` command line: `program` and `args`, each -/// quoted with the MSDN rules. Also used by `host::console_launch`. -pub(crate) fn build_command_line( - program: &Path, - args: &[OsString], -) -> ManagedProcessResult> { - let mut cmdline = Vec::new(); - append_quoted(program.as_os_str(), &mut cmdline)?; - for arg in args { - cmdline.push(b' ' as u16); - append_quoted(arg, &mut cmdline)?; - } - cmdline.push(0); - Ok(cmdline) -} - -fn build_environment_block(overrides: &[(OsString, OsString)]) -> ManagedProcessResult> { - let mut values: Vec<(OsString, OsString)> = std::env::vars_os().collect(); - for (key, value) in overrides { - values.retain(|(existing, _)| { - !existing - .to_string_lossy() - .eq_ignore_ascii_case(&key.to_string_lossy()) - }); - values.push((key.clone(), value.clone())); - } - // CreateProcessW expects Unicode environment blocks sorted case-insensitively. - values.sort_by_cached_key(|(key, _)| key.to_string_lossy().to_uppercase()); - let mut block = Vec::new(); - for (key, value) in values { - let mut entry = key; - entry.push("="); - entry.push(value); - block.extend(encode_wide(&entry)?); - block.push(0); - } - block.push(0); - Ok(block) -} - -fn encode_wide(value: &OsStr) -> ManagedProcessResult> { - let wide: Vec = value.encode_wide().collect(); - if wide.contains(&0) { - return Err(ManagedProcessError::new( - "Process argument contains an embedded NUL".to_string(), - )); - } - Ok(wide) -} - -fn encode_wide_nul(value: &OsStr) -> ManagedProcessResult> { - let mut wide = encode_wide(value)?; - wide.push(0); - Ok(wide) -} - -/// Quote one argument using the MSDN command-line rules. -fn append_quoted(value: &OsStr, output: &mut Vec) -> ManagedProcessResult<()> { - let value = encode_wide(value)?; - let needs_quotes = value.is_empty() - || value - .iter() - .any(|code| matches!(*code, 0x20 | 0x09 | 0x0a | 0x0b | 0x22)); - if !needs_quotes { - output.extend(value); - return Ok(()); - } - - output.push(b'"' as u16); - let mut backslashes = 0_usize; - for code in value { - if code == b'\\' as u16 { - backslashes += 1; - continue; - } - let trailing = if code == b'"' as u16 { - backslashes * 2 + 1 - } else { - backslashes - }; - output.extend(std::iter::repeat_n(b'\\' as u16, trailing)); - output.push(code); - backslashes = 0; - } - output.extend(std::iter::repeat_n(b'\\' as u16, backslashes * 2)); - output.push(b'"' as u16); - Ok(()) -} +mod command_line; +pub(crate) use command_line::build_command_line; +use command_line::{build_environment_block, encode_wide_nul}; fn win_handle(value: &OwnedHandle) -> HANDLE { HANDLE(value.as_raw_handle()) @@ -856,305 +765,4 @@ pub fn listening_ports_for_pid(pid: u32) -> ManagedProcessResult> { } #[cfg(test)] -mod tests { - use super::*; - use std::time::Instant; - - fn test_config() -> ManagedProcessConfig { - ManagedProcessConfig { - program: PathBuf::from("powershell.exe"), - args: vec![ - OsString::from("-NoLogo"), - OsString::from("-NoProfile"), - OsString::from("-NonInteractive"), - OsString::from("-Command"), - OsString::from("Start-Sleep -Seconds 30"), - ], - env: ProcessEnvironment::default(), - cwd: None, - pty_rows: 30, - pty_cols: 120, - label: "test".to_string(), - } - } - - fn start_test_process() -> ManagedProcess { - ManagedProcess::spawn(&test_config()).expect("start a managed test process") - } - - fn wait_for_exit(pid: u32) { - let deadline = Instant::now() + Duration::from_secs(5); - while process_is_alive(pid) { - assert!( - Instant::now() < deadline, - "managed child {pid} was not cleaned up" - ); - std::thread::sleep(Duration::from_millis(25)); - } - } - - fn process_is_alive(pid: u32) -> bool { - use windows::Win32::Foundation::{CloseHandle, WAIT_OBJECT_0}; - use windows::Win32::System::Threading::{ - OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_SYNCHRONIZE, - WaitForSingleObject, - }; - - // SAFETY: OpenProcess returns a handle owned by this function and closed below. - match unsafe { - OpenProcess( - PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_SYNCHRONIZE, - false, - pid, - ) - } { - Ok(handle) => { - // SAFETY: `handle` is a valid process handle. - let exited = unsafe { WaitForSingleObject(handle, 0) } == WAIT_OBJECT_0; - // SAFETY: closing the handle returned by OpenProcess exactly once. - drop(unsafe { CloseHandle(handle) }); - !exited - } - Err(_) => false, - } - } - - #[test] - fn terminal_detects_cursor_request_across_reads() { - let mut tail = Vec::new(); - - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"ready\x1b["), - 0 - ); - assert_eq!(terminal_cursor_position_request_count(&mut tail, b"6n"), 1); - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"plain output"), - 0 - ); - assert_eq!( - terminal_cursor_position_request_count(&mut tail, b"\x1b[6nmore\x1b[6n"), - 2 - ); - } - - #[test] - fn terminal_caps_cursor_replies() { - assert_eq!(terminal_cursor_reply_allowance(0, 2), 2); - assert_eq!(terminal_cursor_reply_allowance(31, 4), 1); - assert_eq!(terminal_cursor_reply_allowance(MAX_CURSOR_REPLIES, 1), 0); - } - - #[test] - fn failed_spawn_closes_pseudoconsole_promptly() { - let mut config = test_config(); - config.args.push(OsString::from("\0")); - let (sender, receiver) = std::sync::mpsc::channel(); - std::thread::spawn(move || { - let failed = ManagedProcess::spawn(&config).is_err(); - sender - .send(failed) - .expect("failure result receiver remains available"); - }); - - assert!( - receiver - .recv_timeout(Duration::from_secs(3)) - .expect("post-ConPTY setup failure should return promptly") - ); - } - - #[test] - fn listener_table_finds_current_process_port() { - let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) - .expect("bind a local IPv4 listener"); - let port = listener.local_addr().expect("listener address").port(); - - let ports = listening_ports_for_pid(std::process::id()) - .expect("read the Windows TCP listener table"); - - assert!( - ports.contains(&port), - "listener table should contain {port}" - ); - } - - #[test] - fn managed_job_terminates_its_owned_process() { - use std::os::windows::io::AsRawHandle as _; - use std::os::windows::process::CommandExt as _; - - const CREATE_NO_WINDOW: u32 = 0x08000000; - let mut command = std::process::Command::new("powershell.exe"); - command - .args([ - "-NoLogo", - "-NoProfile", - "-Command", - "Start-Sleep -Seconds 30", - ]) - .creation_flags(CREATE_NO_WINDOW); - let mut child = command.spawn().expect("spawn an isolated test child"); - let job = create_managed_job("test").expect("create a kill-on-close job"); - if let Err(error) = assign_process_to_job(&job, child.as_raw_handle(), "test") { - drop(child.kill()); - drop(child.wait()); - panic!("assign the test child to its job: {error}"); - } - - // SAFETY: only the isolated test child was assigned to this private job. - unsafe { TerminateJobObject(win_handle(&job), 1) }.expect("terminate the private job"); - let deadline = Instant::now() + Duration::from_secs(3); - loop { - if child.try_wait().expect("inspect the test child").is_some() { - break; - } - if Instant::now() >= deadline { - drop(child.kill()); - drop(child.wait()); - panic!("job termination did not stop the test child"); - } - std::thread::sleep(Duration::from_millis(25)); - } - } - - #[tokio::test] - async fn managed_process_shutdown_stops_owned_child() { - let mut process = start_test_process(); - let pid = process.pid(); - assert!(pid > 0, "managed process exposes its pid"); - assert!( - process.try_wait().expect("poll the child").is_none(), - "managed child is running before shutdown" - ); - - process.shutdown(Duration::from_secs(5)).await; - assert!(!process_is_alive(pid), "shutdown stops the owned child"); - } - - #[test] - fn managed_process_drop_terminates_owned_child() { - let process = start_test_process(); - let pid = process.pid(); - assert!(pid > 0, "managed process exposes its pid"); - - drop(process); - - // Drop terminates synchronously and detaches reaping; wait for the OS - // to report the process gone so the test does not race the cleanup thread. - wait_for_exit(pid); - } - - #[test] - fn managed_process_restart_replaces_without_leaking() { - let mut process = start_test_process(); - let first = process.pid(); - - process.restart().expect("restart the managed child"); - let second = process.pid(); - - assert_ne!(first, second, "restart launches a fresh child"); - assert!(!process_is_alive(first), "restart reaps the previous child"); - assert!( - process.try_wait().expect("poll the replacement").is_none(), - "replacement is running after restart" - ); - - drop(process); - wait_for_exit(second); - } - - #[test] - fn managed_process_reports_exit_code_259_as_exited() { - let mut config = test_config(); - config.program = PathBuf::from("cmd.exe"); - config.args = vec![OsString::from("/C"), OsString::from("exit /B 259")]; - let mut process = ManagedProcess::spawn(&config).expect("start an exit-code test child"); - let deadline = Instant::now() + Duration::from_secs(5); - loop { - if let Some(status) = process.try_wait().expect("poll the exit-code test child") { - assert_eq!(status.exit_code(), 259); - break; - } - assert!( - Instant::now() < deadline, - "the child with exit code 259 was reported as running" - ); - std::thread::sleep(Duration::from_millis(25)); - } - drop(process); - } - - #[tokio::test] - async fn managed_process_contains_descendants_of_the_pty_child() { - let marker = std::env::temp_dir().join(format!( - "codexbar-managed-descendant-{}.txt", - std::process::id() - )); - drop(std::fs::remove_file(&marker)); - let script = format!( - "$psi = New-Object System.Diagnostics.ProcessStartInfo; \ - $psi.FileName = 'powershell.exe'; \ - $psi.Arguments = '-NoLogo -NoProfile -Command \"Start-Sleep -Seconds 120\"'; \ - $psi.UseShellExecute = $false; \ - $p = [System.Diagnostics.Process]::Start($psi); \ - Set-Content -LiteralPath '{}' -Value $p.Id; \ - Start-Sleep -Seconds 120", - marker.display() - ); - let config = ManagedProcessConfig { - program: PathBuf::from("powershell.exe"), - args: vec![ - OsString::from("-NoLogo"), - OsString::from("-NoProfile"), - OsString::from("-NonInteractive"), - OsString::from("-Command"), - OsString::from(script), - ], - env: ProcessEnvironment::default(), - cwd: None, - pty_rows: 30, - pty_cols: 120, - label: "test-descendant".to_string(), - }; - let process = ManagedProcess::spawn(&config).expect("start a managed test process"); - - let descendant = wait_for_descendant_pid(&marker); - assert!( - process_is_alive(descendant), - "the descendant of the PTY child should be running before shutdown" - ); - - process.shutdown(Duration::from_secs(10)).await; - - wait_for_exit_within(descendant, Duration::from_secs(10)); - drop(std::fs::remove_file(&marker)); - } - - fn wait_for_descendant_pid(marker: &Path) -> u32 { - let deadline = Instant::now() + Duration::from_secs(15); - loop { - if let Ok(contents) = std::fs::read_to_string(marker) - && let Ok(pid) = contents.trim().parse::() - { - return pid; - } - assert!( - Instant::now() < deadline, - "the PTY child did not report a descendant pid" - ); - std::thread::sleep(Duration::from_millis(50)); - } - } - - fn wait_for_exit_within(pid: u32, budget: Duration) { - let deadline = Instant::now() + budget; - while process_is_alive(pid) { - assert!( - Instant::now() < deadline, - "descendant {pid} survived the managed job cleanup" - ); - std::thread::sleep(Duration::from_millis(25)); - } - } -} +mod tests; diff --git a/rust/src/managed_process/command_line.rs b/rust/src/managed_process/command_line.rs new file mode 100644 index 0000000000..fb5bd6c5cf --- /dev/null +++ b/rust/src/managed_process/command_line.rs @@ -0,0 +1,94 @@ +//! UTF-16 command-line and environment-block encoding for `CreateProcessW`. + +use super::*; + +/// NUL-terminated `CreateProcessW` command line: `program` and `args`, each +/// quoted with the MSDN rules. Also used by `host::console_launch`. +pub(crate) fn build_command_line( + program: &Path, + args: &[OsString], +) -> ManagedProcessResult> { + let mut cmdline = Vec::new(); + append_quoted(program.as_os_str(), &mut cmdline)?; + for arg in args { + cmdline.push(b' ' as u16); + append_quoted(arg, &mut cmdline)?; + } + cmdline.push(0); + Ok(cmdline) +} + +pub(super) fn build_environment_block( + overrides: &[(OsString, OsString)], +) -> ManagedProcessResult> { + let mut values: Vec<(OsString, OsString)> = std::env::vars_os().collect(); + for (key, value) in overrides { + values.retain(|(existing, _)| { + !existing + .to_string_lossy() + .eq_ignore_ascii_case(&key.to_string_lossy()) + }); + values.push((key.clone(), value.clone())); + } + // CreateProcessW expects Unicode environment blocks sorted case-insensitively. + values.sort_by_cached_key(|(key, _)| key.to_string_lossy().to_uppercase()); + let mut block = Vec::new(); + for (key, value) in values { + let mut entry = key; + entry.push("="); + entry.push(value); + block.extend(encode_wide(&entry)?); + block.push(0); + } + block.push(0); + Ok(block) +} + +fn encode_wide(value: &OsStr) -> ManagedProcessResult> { + let wide: Vec = value.encode_wide().collect(); + if wide.contains(&0) { + return Err(ManagedProcessError::new( + "Process argument contains an embedded NUL".to_string(), + )); + } + Ok(wide) +} + +pub(super) fn encode_wide_nul(value: &OsStr) -> ManagedProcessResult> { + let mut wide = encode_wide(value)?; + wide.push(0); + Ok(wide) +} + +/// Quote one argument using the MSDN command-line rules. +fn append_quoted(value: &OsStr, output: &mut Vec) -> ManagedProcessResult<()> { + let value = encode_wide(value)?; + let needs_quotes = value.is_empty() + || value + .iter() + .any(|code| matches!(*code, 0x20 | 0x09 | 0x0a | 0x0b | 0x22)); + if !needs_quotes { + output.extend(value); + return Ok(()); + } + + output.push(b'"' as u16); + let mut backslashes = 0_usize; + for code in value { + if code == b'\\' as u16 { + backslashes += 1; + continue; + } + let trailing = if code == b'"' as u16 { + backslashes * 2 + 1 + } else { + backslashes + }; + output.extend(std::iter::repeat_n(b'\\' as u16, trailing)); + output.push(code); + backslashes = 0; + } + output.extend(std::iter::repeat_n(b'\\' as u16, backslashes * 2)); + output.push(b'"' as u16); + Ok(()) +} diff --git a/rust/src/managed_process/tests.rs b/rust/src/managed_process/tests.rs new file mode 100644 index 0000000000..8c4ed000e5 --- /dev/null +++ b/rust/src/managed_process/tests.rs @@ -0,0 +1,299 @@ +use super::*; +use std::time::Instant; + +fn test_config() -> ManagedProcessConfig { + ManagedProcessConfig { + program: PathBuf::from("powershell.exe"), + args: vec![ + OsString::from("-NoLogo"), + OsString::from("-NoProfile"), + OsString::from("-NonInteractive"), + OsString::from("-Command"), + OsString::from("Start-Sleep -Seconds 30"), + ], + env: ProcessEnvironment::default(), + cwd: None, + pty_rows: 30, + pty_cols: 120, + label: "test".to_string(), + } +} + +fn start_test_process() -> ManagedProcess { + ManagedProcess::spawn(&test_config()).expect("start a managed test process") +} + +fn wait_for_exit(pid: u32) { + let deadline = Instant::now() + Duration::from_secs(5); + while process_is_alive(pid) { + assert!( + Instant::now() < deadline, + "managed child {pid} was not cleaned up" + ); + std::thread::sleep(Duration::from_millis(25)); + } +} + +fn process_is_alive(pid: u32) -> bool { + use windows::Win32::Foundation::{CloseHandle, WAIT_OBJECT_0}; + use windows::Win32::System::Threading::{ + OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_SYNCHRONIZE, WaitForSingleObject, + }; + + // SAFETY: OpenProcess returns a handle owned by this function and closed below. + match unsafe { + OpenProcess( + PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_SYNCHRONIZE, + false, + pid, + ) + } { + Ok(handle) => { + // SAFETY: `handle` is a valid process handle. + let exited = unsafe { WaitForSingleObject(handle, 0) } == WAIT_OBJECT_0; + // SAFETY: closing the handle returned by OpenProcess exactly once. + drop(unsafe { CloseHandle(handle) }); + !exited + } + Err(_) => false, + } +} + +#[test] +fn terminal_detects_cursor_request_across_reads() { + let mut tail = Vec::new(); + + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"ready\x1b["), + 0 + ); + assert_eq!(terminal_cursor_position_request_count(&mut tail, b"6n"), 1); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"plain output"), + 0 + ); + assert_eq!( + terminal_cursor_position_request_count(&mut tail, b"\x1b[6nmore\x1b[6n"), + 2 + ); +} + +#[test] +fn terminal_caps_cursor_replies() { + assert_eq!(terminal_cursor_reply_allowance(0, 2), 2); + assert_eq!(terminal_cursor_reply_allowance(31, 4), 1); + assert_eq!(terminal_cursor_reply_allowance(MAX_CURSOR_REPLIES, 1), 0); +} + +#[test] +fn failed_spawn_closes_pseudoconsole_promptly() { + let mut config = test_config(); + config.args.push(OsString::from("\0")); + let (sender, receiver) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let failed = ManagedProcess::spawn(&config).is_err(); + sender + .send(failed) + .expect("failure result receiver remains available"); + }); + + assert!( + receiver + .recv_timeout(Duration::from_secs(3)) + .expect("post-ConPTY setup failure should return promptly") + ); +} + +#[test] +fn listener_table_finds_current_process_port() { + let listener = std::net::TcpListener::bind((std::net::Ipv4Addr::LOCALHOST, 0)) + .expect("bind a local IPv4 listener"); + let port = listener.local_addr().expect("listener address").port(); + + let ports = + listening_ports_for_pid(std::process::id()).expect("read the Windows TCP listener table"); + + assert!( + ports.contains(&port), + "listener table should contain {port}" + ); +} + +#[test] +fn managed_job_terminates_its_owned_process() { + use std::os::windows::io::AsRawHandle as _; + use std::os::windows::process::CommandExt as _; + + const CREATE_NO_WINDOW: u32 = 0x08000000; + let mut command = std::process::Command::new("powershell.exe"); + command + .args([ + "-NoLogo", + "-NoProfile", + "-Command", + "Start-Sleep -Seconds 30", + ]) + .creation_flags(CREATE_NO_WINDOW); + let mut child = command.spawn().expect("spawn an isolated test child"); + let job = create_managed_job("test").expect("create a kill-on-close job"); + if let Err(error) = assign_process_to_job(&job, child.as_raw_handle(), "test") { + drop(child.kill()); + drop(child.wait()); + panic!("assign the test child to its job: {error}"); + } + + // SAFETY: only the isolated test child was assigned to this private job. + unsafe { TerminateJobObject(win_handle(&job), 1) }.expect("terminate the private job"); + let deadline = Instant::now() + Duration::from_secs(3); + loop { + if child.try_wait().expect("inspect the test child").is_some() { + break; + } + if Instant::now() >= deadline { + drop(child.kill()); + drop(child.wait()); + panic!("job termination did not stop the test child"); + } + std::thread::sleep(Duration::from_millis(25)); + } +} + +#[tokio::test] +async fn managed_process_shutdown_stops_owned_child() { + let mut process = start_test_process(); + let pid = process.pid(); + assert!(pid > 0, "managed process exposes its pid"); + assert!( + process.try_wait().expect("poll the child").is_none(), + "managed child is running before shutdown" + ); + + process.shutdown(Duration::from_secs(5)).await; + assert!(!process_is_alive(pid), "shutdown stops the owned child"); +} + +#[test] +fn managed_process_drop_terminates_owned_child() { + let process = start_test_process(); + let pid = process.pid(); + assert!(pid > 0, "managed process exposes its pid"); + + drop(process); + + // Drop terminates synchronously and detaches reaping; wait for the OS + // to report the process gone so the test does not race the cleanup thread. + wait_for_exit(pid); +} + +#[test] +fn managed_process_restart_replaces_without_leaking() { + let mut process = start_test_process(); + let first = process.pid(); + + process.restart().expect("restart the managed child"); + let second = process.pid(); + + assert_ne!(first, second, "restart launches a fresh child"); + assert!(!process_is_alive(first), "restart reaps the previous child"); + assert!( + process.try_wait().expect("poll the replacement").is_none(), + "replacement is running after restart" + ); + + drop(process); + wait_for_exit(second); +} + +#[test] +fn managed_process_reports_exit_code_259_as_exited() { + let mut config = test_config(); + config.program = PathBuf::from("cmd.exe"); + config.args = vec![OsString::from("/C"), OsString::from("exit /B 259")]; + let mut process = ManagedProcess::spawn(&config).expect("start an exit-code test child"); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if let Some(status) = process.try_wait().expect("poll the exit-code test child") { + assert_eq!(status.exit_code(), 259); + break; + } + assert!( + Instant::now() < deadline, + "the child with exit code 259 was reported as running" + ); + std::thread::sleep(Duration::from_millis(25)); + } + drop(process); +} + +#[tokio::test] +async fn managed_process_contains_descendants_of_the_pty_child() { + let marker = std::env::temp_dir().join(format!( + "codexbar-managed-descendant-{}.txt", + std::process::id() + )); + drop(std::fs::remove_file(&marker)); + let script = format!( + "$psi = New-Object System.Diagnostics.ProcessStartInfo; \ + $psi.FileName = 'powershell.exe'; \ + $psi.Arguments = '-NoLogo -NoProfile -Command \"Start-Sleep -Seconds 120\"'; \ + $psi.UseShellExecute = $false; \ + $p = [System.Diagnostics.Process]::Start($psi); \ + Set-Content -LiteralPath '{}' -Value $p.Id; \ + Start-Sleep -Seconds 120", + marker.display() + ); + let config = ManagedProcessConfig { + program: PathBuf::from("powershell.exe"), + args: vec![ + OsString::from("-NoLogo"), + OsString::from("-NoProfile"), + OsString::from("-NonInteractive"), + OsString::from("-Command"), + OsString::from(script), + ], + env: ProcessEnvironment::default(), + cwd: None, + pty_rows: 30, + pty_cols: 120, + label: "test-descendant".to_string(), + }; + let process = ManagedProcess::spawn(&config).expect("start a managed test process"); + + let descendant = wait_for_descendant_pid(&marker); + assert!( + process_is_alive(descendant), + "the descendant of the PTY child should be running before shutdown" + ); + + process.shutdown(Duration::from_secs(10)).await; + + wait_for_exit_within(descendant, Duration::from_secs(10)); + drop(std::fs::remove_file(&marker)); +} + +fn wait_for_descendant_pid(marker: &Path) -> u32 { + let deadline = Instant::now() + Duration::from_secs(15); + loop { + if let Ok(contents) = std::fs::read_to_string(marker) + && let Ok(pid) = contents.trim().parse::() + { + return pid; + } + assert!( + Instant::now() < deadline, + "the PTY child did not report a descendant pid" + ); + std::thread::sleep(Duration::from_millis(50)); + } +} + +fn wait_for_exit_within(pid: u32, budget: Duration) { + let deadline = Instant::now() + budget; + while process_is_alive(pid) { + assert!( + Instant::now() < deadline, + "descendant {pid} survived the managed job cleanup" + ); + std::thread::sleep(Duration::from_millis(25)); + } +} diff --git a/rust/src/notifications.rs b/rust/src/notifications.rs index 3a4577ff47..e7e1b2f998 100755 --- a/rust/src/notifications.rs +++ b/rust/src/notifications.rs @@ -2,11 +2,6 @@ //! //! Provides Windows toast notifications for usage alerts -#![allow( - dead_code, - reason = "notification helpers are reserved for future alert integration" -)] - use crate::core::ProviderId; use crate::core::{RateWindow, UsagePace}; use crate::locale::{self, LocaleKey}; @@ -99,17 +94,6 @@ impl NotificationType { ) } - pub fn icon(&self) -> &'static str { - match self { - NotificationType::HighUsage => "⚠️", - NotificationType::CriticalUsage => "🔴", - NotificationType::Exhausted => "🚫", - NotificationType::StatusIssue => "⚡", - NotificationType::SessionDepleted => "🔴", - NotificationType::SessionRestored => "✅", - } - } - fn is_threshold_toast(self) -> bool { matches!( self, @@ -353,35 +337,6 @@ impl NotificationManager { true } - /// Send a notification for a status issue - pub fn notify_status_issue( - &mut self, - provider: ProviderId, - description: &str, - settings: &Settings, - ) { - let key = ( - provider, - String::new(), - String::new(), - NotificationType::StatusIssue, - ); - if !self.sent_notifications.contains(&key) { - self.send_status_notification(provider, description, settings); - self.sent_notifications.insert(key); - } - } - - /// Clear status issue notification (when resolved) - pub fn clear_status_issue(&mut self, provider: ProviderId) { - self.sent_notifications.remove(&( - provider, - String::new(), - String::new(), - NotificationType::StatusIssue, - )); - } - /// Check session quota transitions (depleted/restored) /// Call this with each usage update to detect transitions. /// @@ -530,18 +485,6 @@ impl NotificationManager { } } - fn send_status_notification( - &self, - provider: ProviderId, - description: &str, - settings: &Settings, - ) { - let title = NotificationType::StatusIssue.title(settings.ui_language); - let body = format!("{}: {}", provider.display_name(), description); - self.show_toast(&title, &body); - Self::play_notification_sound(NotificationSoundEvent::StatusIssue, settings); - } - /// Tests record the toast and report it as handed to the OS. #[cfg(test)] fn show_toast(&self, title: &str, body: &str) -> bool { @@ -675,7 +618,7 @@ impl Default for NotificationManager { /// `CreateToastNotifier("CodexBar")` resolves to a valid notifier instead of returning /// null. Must be called at least once before the first toast. Safe to call multiple /// times (idempotent registry write). -#[cfg(target_os = "windows")] +#[cfg(all(target_os = "windows", not(test)))] fn ensure_aumid_registered() { use winreg::RegKey; use winreg::enums::*; @@ -693,545 +636,5 @@ fn ensure_aumid_registered() { } } -/// Simple notification function for one-off notifications -pub fn show_notification(title: &str, body: &str) { - let manager = NotificationManager::new(); - manager.show_toast(title, body); -} - #[cfg(test)] -mod tests { - use super::*; - use crate::core::{PaceStage, RateWindow, UsagePace}; - use chrono::{DateTime, Duration, Utc}; - - #[test] - fn notification_types_map_to_their_sound_events() { - let mappings = [ - ( - NotificationType::HighUsage, - NotificationSoundEvent::HighUsage, - ), - ( - NotificationType::CriticalUsage, - NotificationSoundEvent::CriticalUsage, - ), - ( - NotificationType::Exhausted, - NotificationSoundEvent::Exhausted, - ), - ( - NotificationType::StatusIssue, - NotificationSoundEvent::StatusIssue, - ), - ( - NotificationType::SessionDepleted, - NotificationSoundEvent::SessionDepleted, - ), - ( - NotificationType::SessionRestored, - NotificationSoundEvent::SessionRestored, - ), - ]; - - for (notification_type, sound_event) in mappings { - assert_eq!( - NotificationManager::sound_event_for(notification_type), - sound_event - ); - } - } - - #[test] - fn usage_toasts_follow_ui_language() { - use crate::settings::Language; - - assert_eq!( - NotificationType::HighUsage.title(Language::English), - "High Usage Warning" - ); - assert_eq!( - NotificationManager::notification_body( - ProviderId::Claude, - "weekly", - 86.0, - NotificationType::HighUsage, - Language::English, - ), - "Claude weekly usage at 86% - approaching limit" - ); - assert_eq!( - NotificationType::HighUsage.title(Language::Russian), - "Высокое использование" - ); - assert_eq!( - NotificationManager::notification_body( - ProviderId::Claude, - "weekly", - 86.0, - NotificationType::HighUsage, - Language::Russian, - ), - "Claude (неделя): использовано 86% — лимит почти исчерпан" - ); - // Unknown windows pass through untranslated instead of falling back. - assert_eq!( - NotificationManager::notification_body( - ProviderId::Claude, - "opus", - 100.0, - NotificationType::Exhausted, - Language::English, - ), - "Claude opus usage limit exhausted (100%)" - ); - } - - fn pace(will_last_to_reset: bool, eta_seconds: Option) -> UsagePace { - UsagePace { - stage: PaceStage::Ahead, - delta_percent: 20.0, - expected_used_percent: 40.0, - actual_used_percent: 60.0, - eta_seconds, - will_last_to_reset, - } - } - - fn window(now: DateTime, offset: Duration, minutes: u32) -> RateWindow { - RateWindow::with_details(60.0, Some(minutes), Some(now + offset), None) - } - - #[test] - fn predictive_warning_notifies_once_until_recovery_then_rearms() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let window = window(now, Duration::hours(3), 300); - let risk = pace(false, Some(3600.0)); - let recovery = pace(true, None); - let mut manager = NotificationManager::new(); - - assert!(manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &window, - &risk, - )); - assert!(!manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &window, - &risk, - )); - assert!(!manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &window, - &recovery, - )); - assert!(manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &window, - &risk, - )); - } - - #[test] - fn predictive_warning_reset_jitter_does_not_retrigger() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let mut manager = NotificationManager::new(); - let risk = pace(false, Some(3600.0)); - - assert!(manager.record_predictive_observation( - true, - ProviderId::Codex, - "oauth:account-a", - PredictiveWarningWindow::Weekly, - &window(now, Duration::days(3), 10080), - &risk, - )); - assert!(!manager.record_predictive_observation( - true, - ProviderId::Codex, - "oauth:account-a", - PredictiveWarningWindow::Weekly, - &window(now, Duration::days(3) + Duration::minutes(5), 10080), - &risk, - )); - } - - #[test] - fn predictive_warning_isolates_provider_identity_source_and_window() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let reset = window(now, Duration::hours(3), 300); - let risk = pace(false, Some(3600.0)); - let mut manager = NotificationManager::new(); - - for (provider, identity, warning_window) in [ - ( - ProviderId::Claude, - "cli:person@example.com", - PredictiveWarningWindow::Session, - ), - ( - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - ), - ( - ProviderId::Claude, - "token-account:1", - PredictiveWarningWindow::Session, - ), - ( - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Weekly, - ), - ( - ProviderId::Codex, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - ), - ] { - assert!(manager.record_predictive_observation( - true, - provider, - identity, - warning_window, - &reset, - &risk, - )); - } - } - - #[test] - fn predictive_warning_requires_enabled_confident_positive_risk() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let reset = window(now, Duration::hours(3), 300); - let mut manager = NotificationManager::new(); - - for (enabled, observation) in [ - (false, pace(false, Some(3600.0))), - (true, pace(true, None)), - (true, pace(false, Some(0.0))), - ] { - assert!(!manager.record_predictive_observation( - enabled, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &reset, - &observation, - )); - } - - assert!(manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &reset, - &pace(false, Some(3600.0)), - )); - } - - #[test] - fn unresolved_and_resolved_warning_histories_remain_separate() { - let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); - let reset = window(now, Duration::hours(3), 300); - let risk = pace(false, Some(3600.0)); - let mut manager = NotificationManager::new(); - - assert!(manager.record_predictive_observation( - true, - ProviderId::Claude, - "claude:oauth:unknown", - PredictiveWarningWindow::Session, - &reset, - &risk, - )); - assert!(manager.record_predictive_observation( - true, - ProviderId::Claude, - "oauth:person@example.com", - PredictiveWarningWindow::Session, - &reset, - &risk, - )); - assert_eq!(manager.predictive_warning_keys.len(), 2); - } - - #[test] - fn session_below_high_does_not_rearm_weekly_high_toast() { - // Repro for #198: session cool + weekly hot on every refresh used to - // clear all provider keys on the session call, then re-fire weekly. - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - assert!(settings.show_notifications); - assert!((settings.high_usage_threshold - 70.0).abs() < f64::EPSILON); - - let account = ""; - let weekly_key = ( - ProviderId::Claude, - account.to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - - manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); - assert!(manager.sent_notifications.contains(&weekly_key)); - - // Simulate several refresh cycles: session still cool, weekly still hot. - for _ in 0..5 { - manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); - } - assert_eq!( - manager - .sent_notifications - .iter() - .filter(|key| key == &&weekly_key) - .count(), - 1, - "weekly high toast must arm only once while still above threshold" - ); - - // Drop weekly below high → re-arm allowed on next climb. - manager.check_and_notify(ProviderId::Claude, account, "weekly", 50.0, &settings); - assert!(!manager.sent_notifications.contains(&weekly_key)); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); - assert!(manager.sent_notifications.contains(&weekly_key)); - } - - #[test] - fn missing_session_lane_preserves_warning_history_and_weekly_eligibility() { - let settings = Settings { - high_usage_threshold: 50.0, - ..Settings::default() - }; - let mut manager = NotificationManager::new(); - let account = ""; - - let session_high = ( - ProviderId::Claude, - account.to_string(), - "session".to_string(), - NotificationType::HighUsage, - ); - let weekly_high = ( - ProviderId::Claude, - account.to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - - assert!(manager.check_session_lane(ProviderId::Claude, account, 51.0, false, &settings,)); - assert!(manager.sent_notifications.contains(&session_high)); - - assert!(!manager.check_session_lane(ProviderId::Claude, account, 0.0, true, &settings,)); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); - assert!(manager.sent_notifications.contains(&weekly_high)); - - assert!(manager.check_session_lane(ProviderId::Claude, account, 52.0, false, &settings,)); - assert!(manager.check_session_lane(ProviderId::Claude, account, 81.0, false, &settings,)); - assert_eq!( - manager - .sent_notifications - .iter() - .filter(|key| **key == session_high) - .count(), - 1, - "missing session data must not re-arm the session warning" - ); - assert!(manager.sent_notifications.contains(&weekly_high)); - } - - #[test] - fn threshold_keys_isolate_session_and_weekly() { - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - let account = ""; - - manager.check_and_notify(ProviderId::Claude, account, "session", 75.0, &settings); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 75.0, &settings); - - assert!(manager.sent_notifications.contains(&( - ProviderId::Claude, - account.to_string(), - "session".to_string(), - NotificationType::HighUsage, - ))); - assert!(manager.sent_notifications.contains(&( - ProviderId::Claude, - account.to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ))); - - // Cool only session; weekly stays armed. - manager.check_and_notify(ProviderId::Claude, account, "session", 10.0, &settings); - assert!(!manager.sent_notifications.contains(&( - ProviderId::Claude, - account.to_string(), - "session".to_string(), - NotificationType::HighUsage, - ))); - assert!(manager.sent_notifications.contains(&( - ProviderId::Claude, - account.to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ))); - } - - /// Mirrors `notify_usage_thresholds` in the Tauri shell: each refresh - /// calls session then weekly. Confidence pass for #198 over many cycles. - #[test] - fn refresh_loop_session_cool_weekly_hot_toasts_once() { - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - let account = ""; - let weekly_high = ( - ProviderId::Claude, - account.to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - - let mut weekly_fires = 0usize; - for _ in 0..30 { - let before = manager.sent_notifications.contains(&weekly_high); - // Same order as apps/desktop-tauri/.../providers.rs - manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); - manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); - let after = manager.sent_notifications.contains(&weekly_high); - if after && !before { - weekly_fires += 1; - } - } - - assert_eq!( - weekly_fires, 1, - "weekly high must fire exactly once across 30 refresh cycles" - ); - assert!(manager.sent_notifications.contains(&weekly_high)); - // Session never armed high while cool. - assert!(!manager.sent_notifications.contains(&( - ProviderId::Claude, - account.to_string(), - "session".to_string(), - NotificationType::HighUsage, - ))); - } - - #[test] - fn threshold_keys_isolate_accounts_on_same_provider() { - // Two accounts on the same provider can each fire High once for weekly. - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - - let key_a = ( - ProviderId::Claude, - "account-a".to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - let key_b = ( - ProviderId::Claude, - "account-b".to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - - manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 80.0, &settings); - manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); - - assert!(manager.sent_notifications.contains(&key_a)); - assert!(manager.sent_notifications.contains(&key_b)); - - // Re-poll both still hot: neither re-fires (still armed, no second insert). - manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 80.0, &settings); - manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); - assert_eq!( - manager - .sent_notifications - .iter() - .filter(|k| k.0 == ProviderId::Claude - && k.2 == "weekly" - && k.3 == NotificationType::HighUsage) - .count(), - 2 - ); - } - - #[test] - fn account_a_session_cool_does_not_clear_account_b_weekly() { - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - - let key_b_weekly = ( - ProviderId::Claude, - "account-b".to_string(), - "weekly".to_string(), - NotificationType::HighUsage, - ); - - manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); - assert!(manager.sent_notifications.contains(&key_b_weekly)); - - // Account A session cool must not clear account B weekly armed state. - manager.check_and_notify(ProviderId::Claude, "account-a", "session", 10.0, &settings); - assert!( - manager.sent_notifications.contains(&key_b_weekly), - "account A session cool must not clear account B weekly threshold key" - ); - - // Account A weekly cool must also not clear account B. - manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 10.0, &settings); - assert!(manager.sent_notifications.contains(&key_b_weekly)); - } - - #[test] - fn session_transition_isolates_accounts() { - let mut manager = NotificationManager::new(); - let settings = Settings::default(); - - let depleted_a = ( - ProviderId::Claude, - "account-a".to_string(), - "session".to_string(), - NotificationType::SessionDepleted, - ); - let depleted_b = ( - ProviderId::Claude, - "account-b".to_string(), - "session".to_string(), - NotificationType::SessionDepleted, - ); - - manager.check_session_transition(ProviderId::Claude, "account-a", 50.0, &settings); - manager.check_session_transition(ProviderId::Claude, "account-a", 100.0, &settings); - assert!(manager.sent_notifications.contains(&depleted_a)); - assert!(!manager.sent_notifications.contains(&depleted_b)); - - // Account B still has quota; restoring A must not affect B's lane. - manager.check_session_transition(ProviderId::Claude, "account-b", 40.0, &settings); - manager.check_session_transition(ProviderId::Claude, "account-a", 20.0, &settings); - assert!(!manager.sent_notifications.contains(&depleted_a)); - assert!(!manager.sent_notifications.contains(&depleted_b)); - - // Account B can still fire depleted independently. - manager.check_session_transition(ProviderId::Claude, "account-b", 100.0, &settings); - assert!(manager.sent_notifications.contains(&depleted_b)); - } -} +mod tests; diff --git a/rust/src/notifications/tests.rs b/rust/src/notifications/tests.rs new file mode 100644 index 0000000000..1b9e31b310 --- /dev/null +++ b/rust/src/notifications/tests.rs @@ -0,0 +1,533 @@ +use super::*; +use crate::core::{PaceStage, RateWindow, UsagePace}; +use chrono::{DateTime, Duration, Utc}; + +#[test] +fn notification_types_map_to_their_sound_events() { + let mappings = [ + ( + NotificationType::HighUsage, + NotificationSoundEvent::HighUsage, + ), + ( + NotificationType::CriticalUsage, + NotificationSoundEvent::CriticalUsage, + ), + ( + NotificationType::Exhausted, + NotificationSoundEvent::Exhausted, + ), + ( + NotificationType::StatusIssue, + NotificationSoundEvent::StatusIssue, + ), + ( + NotificationType::SessionDepleted, + NotificationSoundEvent::SessionDepleted, + ), + ( + NotificationType::SessionRestored, + NotificationSoundEvent::SessionRestored, + ), + ]; + + for (notification_type, sound_event) in mappings { + assert_eq!( + NotificationManager::sound_event_for(notification_type), + sound_event + ); + } +} + +#[test] +fn usage_toasts_follow_ui_language() { + use crate::settings::Language; + + assert_eq!( + NotificationType::HighUsage.title(Language::English), + "High Usage Warning" + ); + assert_eq!( + NotificationManager::notification_body( + ProviderId::Claude, + "weekly", + 86.0, + NotificationType::HighUsage, + Language::English, + ), + "Claude weekly usage at 86% - approaching limit" + ); + assert_eq!( + NotificationType::HighUsage.title(Language::Russian), + "Высокое использование" + ); + assert_eq!( + NotificationManager::notification_body( + ProviderId::Claude, + "weekly", + 86.0, + NotificationType::HighUsage, + Language::Russian, + ), + "Claude (неделя): использовано 86% — лимит почти исчерпан" + ); + // Unknown windows pass through untranslated instead of falling back. + assert_eq!( + NotificationManager::notification_body( + ProviderId::Claude, + "opus", + 100.0, + NotificationType::Exhausted, + Language::English, + ), + "Claude opus usage limit exhausted (100%)" + ); +} + +fn pace(will_last_to_reset: bool, eta_seconds: Option) -> UsagePace { + UsagePace { + stage: PaceStage::Ahead, + delta_percent: 20.0, + expected_used_percent: 40.0, + actual_used_percent: 60.0, + eta_seconds, + will_last_to_reset, + } +} + +fn window(now: DateTime, offset: Duration, minutes: u32) -> RateWindow { + RateWindow::with_details(60.0, Some(minutes), Some(now + offset), None) +} + +#[test] +fn predictive_warning_notifies_once_until_recovery_then_rearms() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let window = window(now, Duration::hours(3), 300); + let risk = pace(false, Some(3600.0)); + let recovery = pace(true, None); + let mut manager = NotificationManager::new(); + + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &recovery, + )); + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &window, + &risk, + )); +} + +#[test] +fn predictive_warning_reset_jitter_does_not_retrigger() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let mut manager = NotificationManager::new(); + let risk = pace(false, Some(3600.0)); + + assert!(manager.record_predictive_observation( + true, + ProviderId::Codex, + "oauth:account-a", + PredictiveWarningWindow::Weekly, + &window(now, Duration::days(3), 10080), + &risk, + )); + assert!(!manager.record_predictive_observation( + true, + ProviderId::Codex, + "oauth:account-a", + PredictiveWarningWindow::Weekly, + &window(now, Duration::days(3) + Duration::minutes(5), 10080), + &risk, + )); +} + +#[test] +fn predictive_warning_isolates_provider_identity_source_and_window() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let reset = window(now, Duration::hours(3), 300); + let risk = pace(false, Some(3600.0)); + let mut manager = NotificationManager::new(); + + for (provider, identity, warning_window) in [ + ( + ProviderId::Claude, + "cli:person@example.com", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "token-account:1", + PredictiveWarningWindow::Session, + ), + ( + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Weekly, + ), + ( + ProviderId::Codex, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + ), + ] { + assert!(manager.record_predictive_observation( + true, + provider, + identity, + warning_window, + &reset, + &risk, + )); + } +} + +#[test] +fn predictive_warning_requires_enabled_confident_positive_risk() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let reset = window(now, Duration::hours(3), 300); + let mut manager = NotificationManager::new(); + + for (enabled, observation) in [ + (false, pace(false, Some(3600.0))), + (true, pace(true, None)), + (true, pace(false, Some(0.0))), + ] { + assert!(!manager.record_predictive_observation( + enabled, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &reset, + &observation, + )); + } + + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &reset, + &pace(false, Some(3600.0)), + )); +} + +#[test] +fn unresolved_and_resolved_warning_histories_remain_separate() { + let now = DateTime::from_timestamp(1_800_000_000, 0).unwrap(); + let reset = window(now, Duration::hours(3), 300); + let risk = pace(false, Some(3600.0)); + let mut manager = NotificationManager::new(); + + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "claude:oauth:unknown", + PredictiveWarningWindow::Session, + &reset, + &risk, + )); + assert!(manager.record_predictive_observation( + true, + ProviderId::Claude, + "oauth:person@example.com", + PredictiveWarningWindow::Session, + &reset, + &risk, + )); + assert_eq!(manager.predictive_warning_keys.len(), 2); +} + +#[test] +fn session_below_high_does_not_rearm_weekly_high_toast() { + // Repro for #198: session cool + weekly hot on every refresh used to + // clear all provider keys on the session call, then re-fire weekly. + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + assert!(settings.show_notifications); + assert!((settings.high_usage_threshold - 70.0).abs() < f64::EPSILON); + + let account = ""; + let weekly_key = ( + ProviderId::Claude, + account.to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + + manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); + assert!(manager.sent_notifications.contains(&weekly_key)); + + // Simulate several refresh cycles: session still cool, weekly still hot. + for _ in 0..5 { + manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); + } + assert_eq!( + manager + .sent_notifications + .iter() + .filter(|key| key == &&weekly_key) + .count(), + 1, + "weekly high toast must arm only once while still above threshold" + ); + + // Drop weekly below high → re-arm allowed on next climb. + manager.check_and_notify(ProviderId::Claude, account, "weekly", 50.0, &settings); + assert!(!manager.sent_notifications.contains(&weekly_key)); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); + assert!(manager.sent_notifications.contains(&weekly_key)); +} + +#[test] +fn missing_session_lane_preserves_warning_history_and_weekly_eligibility() { + let settings = Settings { + high_usage_threshold: 50.0, + ..Settings::default() + }; + let mut manager = NotificationManager::new(); + let account = ""; + + let session_high = ( + ProviderId::Claude, + account.to_string(), + "session".to_string(), + NotificationType::HighUsage, + ); + let weekly_high = ( + ProviderId::Claude, + account.to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + + assert!(manager.check_session_lane(ProviderId::Claude, account, 51.0, false, &settings,)); + assert!(manager.sent_notifications.contains(&session_high)); + + assert!(!manager.check_session_lane(ProviderId::Claude, account, 0.0, true, &settings,)); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); + assert!(manager.sent_notifications.contains(&weekly_high)); + + assert!(manager.check_session_lane(ProviderId::Claude, account, 52.0, false, &settings,)); + assert!(manager.check_session_lane(ProviderId::Claude, account, 81.0, false, &settings,)); + assert_eq!( + manager + .sent_notifications + .iter() + .filter(|key| **key == session_high) + .count(), + 1, + "missing session data must not re-arm the session warning" + ); + assert!(manager.sent_notifications.contains(&weekly_high)); +} + +#[test] +fn threshold_keys_isolate_session_and_weekly() { + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + let account = ""; + + manager.check_and_notify(ProviderId::Claude, account, "session", 75.0, &settings); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 75.0, &settings); + + assert!(manager.sent_notifications.contains(&( + ProviderId::Claude, + account.to_string(), + "session".to_string(), + NotificationType::HighUsage, + ))); + assert!(manager.sent_notifications.contains(&( + ProviderId::Claude, + account.to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ))); + + // Cool only session; weekly stays armed. + manager.check_and_notify(ProviderId::Claude, account, "session", 10.0, &settings); + assert!(!manager.sent_notifications.contains(&( + ProviderId::Claude, + account.to_string(), + "session".to_string(), + NotificationType::HighUsage, + ))); + assert!(manager.sent_notifications.contains(&( + ProviderId::Claude, + account.to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ))); +} + +/// Mirrors `notify_usage_thresholds` in the Tauri shell: each refresh +/// calls session then weekly. Confidence pass for #198 over many cycles. +#[test] +fn refresh_loop_session_cool_weekly_hot_toasts_once() { + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + let account = ""; + let weekly_high = ( + ProviderId::Claude, + account.to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + + let mut weekly_fires = 0usize; + for _ in 0..30 { + let before = manager.sent_notifications.contains(&weekly_high); + // Same order as apps/desktop-tauri/.../providers.rs + manager.check_and_notify(ProviderId::Claude, account, "session", 20.0, &settings); + manager.check_and_notify(ProviderId::Claude, account, "weekly", 76.0, &settings); + let after = manager.sent_notifications.contains(&weekly_high); + if after && !before { + weekly_fires += 1; + } + } + + assert_eq!( + weekly_fires, 1, + "weekly high must fire exactly once across 30 refresh cycles" + ); + assert!(manager.sent_notifications.contains(&weekly_high)); + // Session never armed high while cool. + assert!(!manager.sent_notifications.contains(&( + ProviderId::Claude, + account.to_string(), + "session".to_string(), + NotificationType::HighUsage, + ))); +} + +#[test] +fn threshold_keys_isolate_accounts_on_same_provider() { + // Two accounts on the same provider can each fire High once for weekly. + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + + let key_a = ( + ProviderId::Claude, + "account-a".to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + let key_b = ( + ProviderId::Claude, + "account-b".to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + + manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 80.0, &settings); + manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); + + assert!(manager.sent_notifications.contains(&key_a)); + assert!(manager.sent_notifications.contains(&key_b)); + + // Re-poll both still hot: neither re-fires (still armed, no second insert). + manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 80.0, &settings); + manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); + assert_eq!( + manager + .sent_notifications + .iter() + .filter(|k| k.0 == ProviderId::Claude + && k.2 == "weekly" + && k.3 == NotificationType::HighUsage) + .count(), + 2 + ); +} + +#[test] +fn account_a_session_cool_does_not_clear_account_b_weekly() { + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + + let key_b_weekly = ( + ProviderId::Claude, + "account-b".to_string(), + "weekly".to_string(), + NotificationType::HighUsage, + ); + + manager.check_and_notify(ProviderId::Claude, "account-b", "weekly", 80.0, &settings); + assert!(manager.sent_notifications.contains(&key_b_weekly)); + + // Account A session cool must not clear account B weekly armed state. + manager.check_and_notify(ProviderId::Claude, "account-a", "session", 10.0, &settings); + assert!( + manager.sent_notifications.contains(&key_b_weekly), + "account A session cool must not clear account B weekly threshold key" + ); + + // Account A weekly cool must also not clear account B. + manager.check_and_notify(ProviderId::Claude, "account-a", "weekly", 10.0, &settings); + assert!(manager.sent_notifications.contains(&key_b_weekly)); +} + +#[test] +fn session_transition_isolates_accounts() { + let mut manager = NotificationManager::new(); + let settings = Settings::default(); + + let depleted_a = ( + ProviderId::Claude, + "account-a".to_string(), + "session".to_string(), + NotificationType::SessionDepleted, + ); + let depleted_b = ( + ProviderId::Claude, + "account-b".to_string(), + "session".to_string(), + NotificationType::SessionDepleted, + ); + + manager.check_session_transition(ProviderId::Claude, "account-a", 50.0, &settings); + manager.check_session_transition(ProviderId::Claude, "account-a", 100.0, &settings); + assert!(manager.sent_notifications.contains(&depleted_a)); + assert!(!manager.sent_notifications.contains(&depleted_b)); + + // Account B still has quota; restoring A must not affect B's lane. + manager.check_session_transition(ProviderId::Claude, "account-b", 40.0, &settings); + manager.check_session_transition(ProviderId::Claude, "account-a", 20.0, &settings); + assert!(!manager.sent_notifications.contains(&depleted_a)); + assert!(!manager.sent_notifications.contains(&depleted_b)); + + // Account B can still fire depleted independently. + manager.check_session_transition(ProviderId::Claude, "account-b", 100.0, &settings); + assert!(manager.sent_notifications.contains(&depleted_b)); +} diff --git a/rust/src/settings.rs b/rust/src/settings.rs index e3ee9f6883..76e0503da0 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -6,11 +6,6 @@ //! - Manual cookies //! - Other user preferences -#![allow( - dead_code, - reason = "settings types mirror the full config schema; some fields are not yet consumed" -)] - use serde::{Deserialize, Serialize}; use std::collections::{BTreeMap, HashMap, HashSet}; use std::path::{Path, PathBuf}; @@ -33,10 +28,12 @@ mod config_path; mod cost_time_zone; mod manual_cookies; mod optional_details; +mod persistence; mod preferences_document; +mod provider_accessors; mod provider_workspace; mod raw; -mod status; +mod start_at_login; mod types; pub use api_keys::*; @@ -50,7 +47,6 @@ pub use optional_details::provider_has_optional_details; pub use preferences_document::*; pub use provider_workspace::*; use raw::RawSettings; -pub use status::*; pub use types::*; #[cfg(test)] @@ -682,244 +678,6 @@ pub fn normalize_overview_layout(value: &str) -> String { } impl Settings { - /// Get the settings file path (`CODEXBAR_CONFIG` when the CLI applied it). - pub fn settings_path() -> Option { - config_path::settings_file_for(settings_file_override(), crate::logging::config_root()) - } - - /// Load settings from disk - pub fn load() -> Self { - let path = Self::settings_path(); - #[allow( - unused_mut, - reason = "mutability is needed for conditional initialization paths that the compiler cannot prove" - )] - let mut settings = Self::load_from_path(path.as_deref()); - - // Sync autostart toggle with actual registry state and repair stale commands from older builds. - // A CLI run on a `CODEXBAR_CONFIG` file skips these desktop integrations, so it touches - // neither the Run key nor the default config root. - #[cfg(target_os = "windows")] - if settings_file_override().is_none() { - settings.start_at_login = Self::sync_start_at_login_registry(); - settings.apply_promote_tray_default_migration(); - } - - // One-shot migration: materialize the persisted hidden usage-item list - // from the pre-0.62 per-provider visibility flags. After this the list - // is the sole source of truth and the flags stay untouched. - settings.migrate_legacy_usage_item_flags(); - - settings - } - - fn load_from_path(path: Option<&Path>) -> Self { - match path { - Some(path) if path.exists() => match crate::secure_file::read_string(path) { - Ok(content) => { - serde_json::from_str(content.trim_start_matches('\u{feff}')).unwrap_or_default() - } - Err(_) => Self::default(), - }, - _ => Self::default(), - } - } - - /// Materialize `hidden_usage_item_ids` from the pre-0.62 per-provider - /// visibility flags where the list was never persisted. Idempotent: a - /// provider with an explicit list is left alone. - fn migrate_legacy_usage_item_flags(&mut self) { - if self - .provider_config(ProviderId::Codex) - .and_then(|config| config.hidden_usage_item_ids.as_ref()) - .is_none() - && !self.spark_usage_visible(ProviderId::Codex) - { - self.toggle_hidden_items(ProviderId::Codex, &CODEX_SPARK_USAGE_ITEM_IDS, false); - } - if self - .provider_config(ProviderId::Claude) - .and_then(|config| config.hidden_usage_item_ids.as_ref()) - .is_none() - && !self.claude_daily_routines_usage_visible - { - self.toggle_hidden_items( - ProviderId::Claude, - &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], - false, - ); - } - } - - /// Marker written after the one-shot "pin tray by default" migration (issue #237). - fn promote_tray_default_marker_path() -> Option { - crate::logging::config_root().map(|p| p.join(".tray-pin-default-v1")) - } - - /// Old builds defaulted `promote_tray_icon` to false and persisted that on any - /// settings save. Flip those installs to the new default once; later opt-outs - /// are preserved because the marker file remains. - fn should_migrate_promote_tray_default( - promote_tray_icon: bool, - already_migrated: bool, - ) -> bool { - !already_migrated && !promote_tray_icon - } - - fn apply_promote_tray_default_migration(&mut self) { - let Some(marker) = Self::promote_tray_default_marker_path() else { - return; - }; - let already_migrated = marker.exists(); - if Self::should_migrate_promote_tray_default(self.promote_tray_icon, already_migrated) { - self.promote_tray_icon = true; - if let Err(error) = self.save() { - tracing::warn!("Failed to persist promote_tray_icon default migration: {error}"); - } - } - if !already_migrated && let Some(parent) = marker.parent() { - // Best-effort marker dir creation; the write below reports failure. - let _created_dir = std::fs::create_dir_all(parent); - if let Err(error) = std::fs::write(&marker, b"1") { - tracing::warn!("Failed to write promote_tray_icon migration marker: {error}"); - } - } - } - - /// Save settings to disk - pub fn save(&self) -> anyhow::Result<()> { - let path = Self::settings_path() - .ok_or_else(|| anyhow::anyhow!("Could not determine settings path"))?; - - self.save_to_path(&path) - } - - fn save_to_path(&self, path: &Path) -> anyhow::Result<()> { - // Ensure parent directory exists - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - - let json = serde_json::to_string_pretty(self)?; - crate::secure_file::write_string(path, &json)?; - - Ok(()) - } - - fn start_at_login_exe_path(current_exe: &std::path::Path) -> std::path::PathBuf { - let file_name = current_exe.file_name().and_then(|name| name.to_str()); - if file_name.is_some_and(|name| { - name.eq_ignore_ascii_case("codexbar-cli.exe") - || name.eq_ignore_ascii_case("codexbar-desktop.exe") - }) && let Some(desktop_exe) = current_exe - .parent() - .map(|dir| dir.join("codexbar.exe")) - .filter(|path| path.exists()) - { - return desktop_exe; - } - - current_exe.to_path_buf() - } - - fn start_at_login_command(current_exe: &std::path::Path) -> String { - let exe_path = Self::start_at_login_exe_path(current_exe); - format!("\"{}\"", exe_path.display()) - } - - fn start_at_login_command_needs_repair(existing: &str, current_exe: &std::path::Path) -> bool { - existing != Self::start_at_login_command(current_exe) - } - - #[cfg(target_os = "windows")] - pub fn apply_start_at_login_registry(enabled: bool) -> anyhow::Result<()> { - use winreg::RegKey; - use winreg::enums::*; - - let hkcu = RegKey::predef(HKEY_CURRENT_USER); - let run_key = hkcu.open_subkey_with_flags( - r"Software\Microsoft\Windows\CurrentVersion\Run", - KEY_READ | KEY_WRITE, - )?; - - if enabled { - let exe_path = std::env::current_exe()?; - let command = Self::start_at_login_command(&exe_path); - run_key.set_value("CodexBar", &command)?; - } else { - // Best-effort removal; a missing value means the desired state already. - let _removed_value = run_key.delete_value("CodexBar"); - } - - Ok(()) - } - - #[cfg(target_os = "windows")] - fn sync_start_at_login_registry() -> bool { - use winreg::RegKey; - use winreg::enums::*; - - let hkcu = RegKey::predef(HKEY_CURRENT_USER); - let Ok(run_key) = hkcu.open_subkey_with_flags( - r"Software\Microsoft\Windows\CurrentVersion\Run", - KEY_READ | KEY_WRITE, - ) else { - return false; - }; - - let Ok(existing) = run_key.get_value::("CodexBar") else { - return false; - }; - - match std::env::current_exe() { - Ok(exe_path) if Self::start_at_login_command_needs_repair(&existing, &exe_path) => { - let command = Self::start_at_login_command(&exe_path); - if let Err(error) = run_key.set_value("CodexBar", &command) { - tracing::warn!("Failed to repair CodexBar start-at-login command: {error}"); - } - } - Err(error) => { - tracing::warn!( - "Failed to resolve current executable for start-at-login sync: {error}" - ); - } - _ => {} - } - - true - } - - #[cfg(not(target_os = "windows"))] - pub fn apply_start_at_login_registry(_enabled: bool) -> anyhow::Result<()> { - Ok(()) - } - - /// Set start at login (updates Windows registry) - pub fn set_start_at_login(&mut self, enabled: bool) -> anyhow::Result<()> { - self.start_at_login = enabled; - Self::apply_start_at_login_registry(enabled)?; - Ok(()) - } - - /// Check if start at login is actually enabled in registry - #[cfg(target_os = "windows")] - pub fn is_start_at_login_enabled() -> bool { - use winreg::RegKey; - use winreg::enums::*; - - let hkcu = RegKey::predef(HKEY_CURRENT_USER); - if let Ok(run_key) = hkcu.open_subkey(r"Software\Microsoft\Windows\CurrentVersion\Run") { - run_key.get_value::("CodexBar").is_ok() - } else { - false - } - } - - #[cfg(not(target_os = "windows"))] - pub fn is_start_at_login_enabled() -> bool { - false - } - /// Check if a provider is enabled pub fn is_provider_enabled(&self, id: ProviderId) -> bool { self.enabled_providers.contains(id.cli_name()) @@ -961,18 +719,6 @@ impl Settings { .collect() } - /// Get all available providers with their enabled status - pub fn get_all_providers_status(&self) -> Vec { - self.provider_display_order() - .into_iter() - .map(|id| ProviderStatus { - id: id.cli_name().to_string(), - name: id.display_name().to_string(), - enabled: self.is_provider_enabled(id), - }) - .collect() - } - /// Provider display order as typed IDs, falling back to canonical order /// when no custom order has been persisted. pub fn provider_display_order(&self) -> Vec { @@ -1000,530 +746,4 @@ impl Settings { self.provider_metrics .insert(id.cli_name().to_string(), metric); } - - // ── Per-provider configuration accessors ───────────────────────── - // - // These thin wrappers around `provider_configs` apply provider-specific - // defaults (e.g. cookie/usage source defaults to `"auto"`) so callers - // never have to reach into the raw `Option` fields. The - // `*_str` / boolean / setter pairs intentionally mirror the names of - // the legacy flat fields so call-site migration is mechanical. - - /// Read-only access to a provider's stored config, if any. - pub fn provider_config(&self, id: ProviderId) -> Option<&ProviderConfig> { - self.provider_configs.get(&id) - } - - /// Mutable access to a provider's config, lazily creating an empty - /// entry if none exists. - pub fn provider_config_mut(&mut self, id: ProviderId) -> &mut ProviderConfig { - self.provider_configs.entry(id).or_default() - } - - /// Cookie source for `id`. Kimi, Charm Hyper and Groq follow upstream's - /// automatic default; providers with no specific default retain the - /// legacy manual default. - pub fn cookie_source(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.cookie_source.as_deref()) - .unwrap_or(match id { - ProviderId::Kimi | ProviderId::Hyper | ProviderId::Groq => "auto", - _ => DEFAULT_COOKIE_SOURCE, - }) - } - - pub fn set_cookie_source(&mut self, id: ProviderId, source: impl Into) { - self.provider_config_mut(id).cookie_source = Some(source.into()); - } - - /// Usage source for `id`, or the default `"auto"` if unset. - pub fn usage_source(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.usage_source.as_deref()) - .unwrap_or(DEFAULT_PROVIDER_SOURCE) - } - - pub fn set_usage_source(&mut self, id: ProviderId, source: impl Into) { - self.provider_config_mut(id).usage_source = Some(source.into()); - } - - /// API region for `id`, or the provider-specific default if unset. - pub fn api_region(&self, id: ProviderId) -> &str { - if id == ProviderId::AlibabaTokenPlan && !self.alibaba_token_plan_region.trim().is_empty() { - return self.alibaba_token_plan_region.as_str(); - } - self.provider_configs - .get(&id) - .and_then(|c| c.api_region.as_deref()) - .unwrap_or_else(|| default_api_region(id)) - } - - pub fn set_api_region(&mut self, id: ProviderId, region: impl Into) { - let region = region.into(); - if id == ProviderId::AlibabaTokenPlan { - self.alibaba_token_plan_region = region.clone(); - } - self.provider_config_mut(id).api_region = Some(region); - } - - /// Manual cookie header for `id`, or `""` if unset. - pub fn manual_cookie_header(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.manual_cookie_header.as_deref()) - .unwrap_or("") - } - - pub fn set_manual_cookie_header(&mut self, id: ProviderId, header: impl Into) { - self.provider_config_mut(id).manual_cookie_header = Some(header.into()); - } - - /// API token for `id`, or `""` if unset. - pub fn api_token(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.api_token.as_deref()) - .unwrap_or("") - } - - pub fn set_api_token(&mut self, id: ProviderId, token: impl Into) { - self.provider_config_mut(id).api_token = Some(token.into()); - } - - pub fn management_api_token(&self, id: ProviderId) -> Option<&str> { - self.provider_configs - .get(&id) - .and_then(|config| config.management_api_token.as_deref()) - .map(str::trim) - .filter(|value| !value.is_empty()) - } - - pub fn set_management_api_token(&mut self, id: ProviderId, token: Option) { - let token = token - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()); - self.provider_config_mut(id).management_api_token = token; - } - - /// Workspace ID override for `id`, or `""` if unset. - pub fn workspace_id(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.workspace_id.as_deref()) - .unwrap_or("") - } - - pub fn set_workspace_id(&mut self, id: ProviderId, value: impl Into) { - self.provider_config_mut(id).workspace_id = Some(value.into()); - } - - /// Optional user-entered allowance for Copilot seat AI credits. - /// - /// GitHub reports the absolute `credits_used` counter but does not expose - /// a documented included-credit ceiling, so callers must keep an absent - /// or non-positive value as unknown rather than inventing a denominator. - /// - /// This setter is the single owner of the positive-finite invariant: - /// invalid values are rejected instead of silently dropped, while the - /// getter keeps defensively filtering values persisted by older builds. - pub fn seat_credit_entitlement(&self, id: ProviderId) -> Option { - self.provider_configs - .get(&id) - .and_then(|config| config.seat_credit_entitlement) - .filter(|value| value.is_finite() && *value > 0.0) - } - - pub fn set_seat_credit_entitlement( - &mut self, - id: ProviderId, - value: Option, - ) -> Result<(), String> { - if let Some(value) = value - && (!value.is_finite() || value <= 0.0) - { - return Err( - "Copilot seat AI-credit allowance must be a finite number greater than zero" - .to_string(), - ); - } - self.provider_config_mut(id).seat_credit_entitlement = value; - Ok(()) - } - - /// Wayfinder gateway URL, defaulting to the local loopback gateway. - pub fn gateway_url(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.gateway_url.as_deref()) - .unwrap_or_else(|| { - if id == ProviderId::Wayfinder { - crate::providers::wayfinder::DEFAULT_GATEWAY_URL - } else { - "" - } - }) - } - - pub fn set_gateway_url(&mut self, id: ProviderId, value: impl Into) { - self.provider_config_mut(id).gateway_url = Some(value.into()); - } - - /// IDE base path override for `id`, or `""` if unset. - pub fn ide_base_path(&self, id: ProviderId) -> &str { - self.provider_configs - .get(&id) - .and_then(|c| c.ide_base_path.as_deref()) - .unwrap_or("") - } - - pub fn set_ide_base_path(&mut self, id: ProviderId, value: impl Into) { - self.provider_config_mut(id).ide_base_path = Some(value.into()); - } - - /// Codex `openai_web_extras` toggle, default `true`. - pub fn openai_web_extras(&self, id: ProviderId) -> bool { - self.provider_configs - .get(&id) - .and_then(|c| c.openai_web_extras) - .unwrap_or(DEFAULT_CODEX_OPENAI_WEB_EXTRAS) - } - - pub fn set_openai_web_extras(&mut self, id: ProviderId, value: bool) { - self.provider_config_mut(id).openai_web_extras = Some(value); - } - - /// Codex Spark rows are visible by default. - pub fn spark_usage_visible(&self, id: ProviderId) -> bool { - self.provider_configs - .get(&id) - .and_then(|c| c.spark_usage_visible) - .unwrap_or(DEFAULT_CODEX_SPARK_USAGE_VISIBLE) - } - - pub fn set_spark_usage_visible(&mut self, id: ProviderId, value: bool) { - self.provider_config_mut(id).spark_usage_visible = Some(value); - } - - /// Return the persisted hidden usage-item IDs for `id`. - pub fn hidden_usage_item_ids(&self, id: ProviderId) -> Vec { - self.provider_configs - .get(&id) - .and_then(|config| config.hidden_usage_item_ids.as_ref()) - .map_or_else(Vec::new, |ids| normalize_hidden_usage_item_ids(ids.clone())) - } - - /// Persist an explicit presentation-only usage-item visibility list. - pub fn set_hidden_usage_item_ids(&mut self, id: ProviderId, ids: Vec) { - let hidden = normalize_hidden_usage_item_ids(ids); - self.provider_config_mut(id).hidden_usage_item_ids = Some(hidden); - } - - /// Add or remove `items` from the provider's hidden usage-item list. - /// - /// `visible = false` hides the items; `visible = true` un-hides them. This - /// is the single write path for usage-item visibility; it only touches the - /// presentation list and never the legacy per-provider boolean flags. - pub fn toggle_hidden_items(&mut self, id: ProviderId, items: &[&str], visible: bool) { - let mut hidden = self.hidden_usage_item_ids(id); - if visible { - hidden.retain(|item| !items.contains(&item.as_str())); - } else { - hidden.extend(items.iter().map(|item| (*item).to_string())); - } - self.set_hidden_usage_item_ids(id, hidden); - } - - /// Update the old Claude Daily Routines flag without touching the - /// usage-item list; the flag is presentation-only and kept for callers - /// that still read the boolean directly. - pub fn set_claude_daily_routines_usage_visible(&mut self, value: bool) { - self.claude_daily_routines_usage_visible = value; - } - - /// Per-provider historical-tracking toggle (currently codex-only). - pub fn historical_tracking(&self, id: ProviderId) -> bool { - self.provider_configs - .get(&id) - .map(|c| c.historical_tracking) - .unwrap_or(false) - } - - pub fn set_historical_tracking(&mut self, id: ProviderId, value: bool) { - self.provider_config_mut(id).historical_tracking = value; - } - - /// Per-provider "avoid keychain prompts" toggle (currently claude-only). - pub fn avoid_keychain_prompts(&self, id: ProviderId) -> bool { - self.provider_configs - .get(&id) - .map(|c| c.avoid_keychain_prompts) - .unwrap_or(false) - } - - pub fn set_avoid_keychain_prompts(&mut self, id: ProviderId, value: bool) { - self.provider_config_mut(id).avoid_keychain_prompts = value; - } - - /// Whether the desktop shell may reopen a captured CLI session after its - /// provider quota becomes available again. This is intentionally opt-in. - pub fn auto_resume_after_quota_reset(&self, id: ProviderId) -> bool { - self.provider_configs - .get(&id) - .map(|config| config.auto_resume_after_quota_reset) - .unwrap_or(false) - } - - pub fn set_auto_resume_after_quota_reset(&mut self, id: ProviderId, value: bool) { - self.provider_config_mut(id).auto_resume_after_quota_reset = value; - } - - // ── Legacy field-name aliases ──────────────────────────────────── - // - // Keep the names of the old flat per-provider fields available as - // accessor methods so existing call sites only need a `()` (read) or - // `set_` prefix (write). New code should prefer the typed accessors - // above. - - pub fn codex_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Codex) - } - pub fn set_codex_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Codex, v) - } - pub fn claude_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Claude) - } - pub fn set_claude_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Claude, v) - } - pub fn cursor_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Cursor) - } - pub fn set_cursor_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Cursor, v) - } - pub fn opencode_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::OpenCode) - } - pub fn set_opencode_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::OpenCode, v) - } - pub fn factory_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Factory) - } - pub fn set_factory_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Factory, v) - } - pub fn alibaba_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Alibaba) - } - pub fn set_alibaba_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Alibaba, v) - } - pub fn kimi_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Kimi) - } - pub fn set_kimi_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Kimi, v) - } - pub fn minimax_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::MiniMax) - } - pub fn set_minimax_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::MiniMax, v) - } - pub fn augment_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Augment) - } - pub fn set_augment_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Augment, v) - } - pub fn amp_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Amp) - } - pub fn set_amp_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Amp, v) - } - pub fn ollama_cookie_source(&self) -> &str { - self.cookie_source(ProviderId::Ollama) - } - pub fn set_ollama_cookie_source(&mut self, v: impl Into) { - self.set_cookie_source(ProviderId::Ollama, v) - } - - pub fn claude_usage_source(&self) -> &str { - self.usage_source(ProviderId::Claude) - } - pub fn set_claude_usage_source(&mut self, v: impl Into) { - self.set_usage_source(ProviderId::Claude, v) - } - pub fn codex_usage_source(&self) -> &str { - self.usage_source(ProviderId::Codex) - } - pub fn set_codex_usage_source(&mut self, v: impl Into) { - self.set_usage_source(ProviderId::Codex, v) - } - - pub fn alibaba_api_region(&self) -> &str { - self.api_region(ProviderId::Alibaba) - } - pub fn set_alibaba_api_region(&mut self, v: impl Into) { - self.set_api_region(ProviderId::Alibaba, v) - } - pub fn zai_api_region(&self) -> &str { - self.api_region(ProviderId::Zai) - } - pub fn set_zai_api_region(&mut self, v: impl Into) { - self.set_api_region(ProviderId::Zai, v) - } - pub fn minimax_api_region(&self) -> &str { - self.api_region(ProviderId::MiniMax) - } - pub fn set_minimax_api_region(&mut self, v: impl Into) { - self.set_api_region(ProviderId::MiniMax, v) - } - - pub fn alibaba_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::Alibaba) - } - pub fn set_alibaba_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::Alibaba, v) - } - pub fn kimi_manual_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::Kimi) - } - pub fn set_kimi_manual_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::Kimi, v) - } - pub fn augment_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::Augment) - } - pub fn set_augment_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::Augment, v) - } - pub fn amp_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::Amp) - } - pub fn set_amp_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::Amp, v) - } - pub fn ollama_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::Ollama) - } - pub fn set_ollama_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::Ollama, v) - } - pub fn minimax_cookie_header(&self) -> &str { - self.manual_cookie_header(ProviderId::MiniMax) - } - pub fn set_minimax_cookie_header(&mut self, v: impl Into) { - self.set_manual_cookie_header(ProviderId::MiniMax, v) - } - - pub fn opencode_workspace_id(&self) -> &str { - self.workspace_id(ProviderId::OpenCode) - } - pub fn set_opencode_workspace_id(&mut self, v: impl Into) { - self.set_workspace_id(ProviderId::OpenCode, v) - } - pub fn minimax_api_token(&self) -> &str { - self.api_token(ProviderId::MiniMax) - } - pub fn set_minimax_api_token(&mut self, v: impl Into) { - self.set_api_token(ProviderId::MiniMax, v) - } - pub fn jetbrains_ide_base_path(&self) -> &str { - self.ide_base_path(ProviderId::JetBrains) - } - pub fn set_jetbrains_ide_base_path(&mut self, v: impl Into) { - self.set_ide_base_path(ProviderId::JetBrains, v) - } - - pub fn codex_openai_web_extras(&self) -> bool { - self.openai_web_extras(ProviderId::Codex) - } - pub fn set_codex_openai_web_extras(&mut self, v: bool) { - self.set_openai_web_extras(ProviderId::Codex, v) - } - pub fn codex_spark_usage_visible(&self) -> bool { - self.spark_usage_visible(ProviderId::Codex) - } - pub fn set_codex_spark_usage_visible(&mut self, v: bool) { - self.set_spark_usage_visible(ProviderId::Codex, v) - } - pub fn codex_historical_tracking(&self) -> bool { - self.historical_tracking(ProviderId::Codex) - } - pub fn set_codex_historical_tracking(&mut self, v: bool) { - self.set_historical_tracking(ProviderId::Codex, v) - } - pub fn claude_avoid_keychain_prompts(&self) -> bool { - self.avoid_keychain_prompts(ProviderId::Claude) - } - pub fn set_claude_avoid_keychain_prompts(&mut self, v: bool) { - self.set_avoid_keychain_prompts(ProviderId::Claude, v) - } - - /// Claude-only: whether the external claude-swap (`cswap`) adapter is - /// enabled. Disabled by default. - pub fn claude_swap_enabled(&self) -> bool { - self.provider_configs - .get(&ProviderId::Claude) - .map(|config| config.claude_swap_enabled) - .unwrap_or(false) - } - - pub fn set_claude_swap_enabled(&mut self, value: bool) { - self.provider_config_mut(ProviderId::Claude) - .claude_swap_enabled = value; - } - - /// Claude-only: configured claude-swap executable path, or `""` when unset. - pub fn claude_swap_executable_path(&self) -> &str { - self.provider_configs - .get(&ProviderId::Claude) - .and_then(|config| config.claude_swap_executable_path.as_deref()) - .unwrap_or("") - } - - pub fn set_claude_swap_executable_path(&mut self, value: impl Into) { - let trimmed = value.into().trim().to_string(); - let config = self.provider_config_mut(ProviderId::Claude); - config.claude_swap_executable_path = if trimmed.is_empty() { - None - } else { - Some(trimmed) - } - } - - // ── Per-provider accent color override (#2972) ────────────────── - - /// The user-overridden accent color for `id`, or `None` to use the - /// shipped brand color. - pub fn accent_color(&self, id: ProviderId) -> Option<&str> { - self.provider_configs - .get(&id) - .and_then(|c| c.accent_color.as_deref()) - .filter(|s| !s.trim().is_empty()) - } - - /// Set the accent color override for `id`. Pass an empty string or - /// `None` to clear the override and revert to the shipped brand color. - pub fn set_accent_color(&mut self, id: ProviderId, color: Option>) { - let entry = self.provider_config_mut(id); - entry.accent_color = color - .map(Into::into) - .filter(|s: &String| !s.trim().is_empty()); - } - - /// Resolve the effective accent color for `id`: the user override if - /// set, otherwise the shipped brand color from the provider registry. - pub fn effective_accent_color(&self, id: ProviderId) -> String { - if let Some(override_color) = self.accent_color(id) { - return override_color.trim().to_string(); - } - crate::core::brand_color(id).to_string() - } } diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index 3ea5048ca4..1eeb3732fd 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -185,10 +185,8 @@ pub struct SavedApiKeyInfo { pub struct ProviderConfigInfo { pub id: ProviderId, pub name: &'static str, - pub requires_api_key: bool, pub api_key_env_var: Option<&'static str>, pub api_key_help: Option<&'static str>, - pub config_file_path: Option<&'static str>, pub dashboard_url: Option<&'static str>, } @@ -198,10 +196,8 @@ pub fn get_api_key_providers() -> Vec { ProviderConfigInfo { id: ProviderId::Alibaba, name: "Alibaba Coding Plan", - requires_api_key: true, api_key_env_var: Some("ALIBABA_CODING_PLAN_API_KEY"), api_key_help: Some("Get your Coding Plan API key from Alibaba Model Studio / Bailian"), - config_file_path: Some("~/.codexbar/config.json"), dashboard_url: Some( "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=coding-plan#/efm/detail", ), @@ -209,65 +205,53 @@ pub fn get_api_key_providers() -> Vec { ProviderConfigInfo { id: ProviderId::Amp, name: "Amp (Sourcegraph)", - requires_api_key: true, api_key_env_var: Some("SRC_ACCESS_TOKEN"), api_key_help: Some("Get your token from Sourcegraph → Settings → Access Tokens"), - config_file_path: Some("~/.amp/config.json"), dashboard_url: Some("https://sourcegraph.com/cody/manage"), }, ProviderConfigInfo { id: ProviderId::Copilot, name: "GitHub Copilot (legacy token)", - requires_api_key: true, api_key_env_var: Some("GITHUB_TOKEN"), api_key_help: Some( "Optional fallback. Prefer Providers → Copilot → Sign in with GitHub.", ), - config_file_path: None, dashboard_url: Some("https://github.com/settings/copilot"), }, ProviderConfigInfo { id: ProviderId::Zai, name: "z.ai", - requires_api_key: true, api_key_env_var: Some("Z_AI_API_KEY or ZAI_API_TOKEN"), api_key_help: Some( "Get your API token from z.ai Dashboard. BigModel team usage can set Z_AI_BIGMODEL_ORGANIZATION + Z_AI_BIGMODEL_PROJECT, or provider workspace_id as organization|project.", ), - config_file_path: None, dashboard_url: Some("https://z.ai/manage-apikey/coding-plan/personal/my-plan"), }, ProviderConfigInfo { id: ProviderId::Warp, name: "Warp", - requires_api_key: true, api_key_env_var: Some("WARP_API_KEY"), api_key_help: Some( "Get your API key from Warp → Settings → API Keys (docs.warp.dev/reference/cli/api-keys)", ), - config_file_path: None, dashboard_url: Some("https://docs.warp.dev/reference/cli/api-keys"), }, ProviderConfigInfo { id: ProviderId::Ollama, name: "Ollama", - requires_api_key: false, api_key_env_var: Some("OLLAMA_API_KEY / OLLAMA_KEY"), api_key_help: Some( "Optional: use an Ollama API key for Cloud validation, or browser cookies for usage.", ), - config_file_path: None, dashboard_url: Some("https://ollama.com/settings"), }, ProviderConfigInfo { id: ProviderId::MiniMax, name: "MiniMax", - requires_api_key: false, api_key_env_var: Some("MINIMAX_API_KEY"), api_key_help: Some( "Optional: a MiniMax API key reads real coding-plan quota via the console's remains API, bypassing the client-rendered usage/plan pages that browser cookies alone cannot scrape.", ), - config_file_path: None, dashboard_url: Some( "https://platform.minimax.io/user-center/basic-information/interface-key", ), @@ -275,480 +259,390 @@ pub fn get_api_key_providers() -> Vec { ProviderConfigInfo { id: ProviderId::AzureOpenAI, name: "Azure OpenAI", - requires_api_key: true, api_key_env_var: Some( "AZURE_OPENAI_API_KEY + AZURE_OPENAI_ENDPOINT + AZURE_OPENAI_DEPLOYMENT", ), api_key_help: Some( "Use env vars, or save JSON {api_key, endpoint, deployment, api_version}; composite api_key|endpoint|deployment[|api_version] also works.", ), - config_file_path: Some("~/.codexbar/api_keys.json"), dashboard_url: Some("https://ai.azure.com"), }, ProviderConfigInfo { id: ProviderId::OpenRouter, name: "OpenRouter", - requires_api_key: true, api_key_env_var: Some("OPENROUTER_API_KEY"), api_key_help: Some( "Required. Enter a regular API key or a Management API key here. Management keys also enable account Activity on the official OpenRouter API.", ), - config_file_path: None, dashboard_url: Some("https://openrouter.ai/settings/credits"), }, ProviderConfigInfo { id: ProviderId::NanoGPT, name: "NanoGPT", - requires_api_key: true, api_key_env_var: Some("NANOGPT_API_KEY"), api_key_help: Some("Get your API key from nano-gpt.com/api"), - config_file_path: None, dashboard_url: Some("https://nano-gpt.com/api"), }, ProviderConfigInfo { id: ProviderId::Infini, name: "Infini AI", - requires_api_key: true, api_key_env_var: Some("INFINI_API_KEY"), api_key_help: Some("Get your API key from Infini Cloud → Settings → API Keys"), - config_file_path: None, dashboard_url: Some("https://cloud.infini-ai.com"), }, ProviderConfigInfo { id: ProviderId::Kimi, name: "Kimi Code API", - requires_api_key: true, api_key_env_var: Some("KIMI_CODE_API_KEY"), api_key_help: Some( "Get your Kimi Code API key from Kimi. Optional HTTPS proxy base URL: KIMI_CODE_BASE_URL.", ), - config_file_path: None, dashboard_url: Some("https://platform.moonshot.cn/console/api-keys"), }, ProviderConfigInfo { id: ProviderId::Kilo, name: "Kilo", - requires_api_key: true, api_key_env_var: Some("KILO_API_KEY"), api_key_help: Some("Get your API key from Kilo, or run `kilo auth login`."), - config_file_path: Some("~/.local/share/kilo/auth.json"), dashboard_url: Some("https://app.kilo.ai/usage"), }, ProviderConfigInfo { id: ProviderId::Bedrock, name: "AWS Bedrock", - requires_api_key: true, api_key_env_var: Some( "AWS_ACCESS_KEY_ID:AWS_SECRET_ACCESS_KEY[:AWS_SESSION_TOKEN] or AWS_PROFILE", ), api_key_help: Some( "Paste access_key:secret_key[:session_token], JSON credentials, profile:name, or use AWS env vars/AWS CLI profiles.", ), - config_file_path: None, dashboard_url: Some("https://console.aws.amazon.com/bedrock"), }, ProviderConfigInfo { id: ProviderId::Codebuff, name: "Codebuff", - requires_api_key: true, api_key_env_var: Some("CODEBUFF_API_KEY"), api_key_help: Some( "Get your API key from Codebuff, or sign in with Codebuff/Manicode.", ), - config_file_path: Some("~/.config/manicode/credentials.json"), dashboard_url: Some("https://www.codebuff.com/usage"), }, ProviderConfigInfo { id: ProviderId::DeepSeek, name: "DeepSeek", - requires_api_key: true, api_key_env_var: Some("DEEPSEEK_API_KEY"), api_key_help: Some("Get your API key from platform.deepseek.com."), - config_file_path: None, dashboard_url: Some("https://platform.deepseek.com/usage"), }, ProviderConfigInfo { id: ProviderId::DeepInfra, name: "DeepInfra", - requires_api_key: true, api_key_env_var: Some("DEEPINFRA_API_KEY"), api_key_help: Some( "Get your API key from deepinfra.com/dash. Also accepts DEEPINFRA_TOKEN.", ), - config_file_path: None, dashboard_url: Some("https://deepinfra.com/dash"), }, ProviderConfigInfo { id: ProviderId::HuggingFace, name: "Hugging Face", - requires_api_key: true, api_key_env_var: Some( "CODEXBAR_HUGGINGFACE_API_KEY / HF_TOKEN / HUGGING_FACE_HUB_TOKEN", ), api_key_help: Some( "Add a Hugging Face access token here, set HF_TOKEN, or run `hf auth login`.", ), - config_file_path: None, dashboard_url: Some("https://huggingface.co/settings/billing"), }, ProviderConfigInfo { id: ProviderId::V0, name: "v0", - requires_api_key: true, api_key_env_var: Some("V0_API_KEY"), api_key_help: Some( "Add a v0 Platform API key. An optional scope can use the provider workspace field or V0_SCOPE.", ), - config_file_path: None, dashboard_url: Some("https://v0.app/settings/billing"), }, ProviderConfigInfo { id: ProviderId::Fireworks, name: "Fireworks", - requires_api_key: true, api_key_env_var: Some("FIREWORKS_API_KEY"), api_key_help: Some( "Get your API key from app.fireworks.ai. Also set the account slug from \ app.fireworks.ai/accounts/ (FIREWORKS_ACCOUNT_SLUG).", ), - config_file_path: None, dashboard_url: Some("https://app.fireworks.ai"), }, ProviderConfigInfo { id: ProviderId::AiAnd, name: "ai&", - requires_api_key: true, api_key_env_var: Some("AIAND_API_KEY"), api_key_help: Some("Get your API key from console.aiand.com."), - config_file_path: None, dashboard_url: Some("https://console.aiand.com"), }, ProviderConfigInfo { id: ProviderId::AtlasCloud, name: "Atlas Cloud", - requires_api_key: true, api_key_env_var: Some("ATLASCLOUD_API_KEY"), api_key_help: Some( "Get an API key from Atlas Cloud and set it in Preferences or ATLASCLOUD_API_KEY.", ), - config_file_path: None, dashboard_url: Some(crate::providers::atlascloud::DASHBOARD_URL), }, ProviderConfigInfo { id: ProviderId::ZenMux, name: "ZenMux", - requires_api_key: true, api_key_env_var: Some("ZENMUX_MANAGEMENT_API_KEY"), api_key_help: Some( "Use a ZenMux Management API key (not an inference key). Also accepts ZENMUX_API_KEY.", ), - config_file_path: None, dashboard_url: Some("https://zenmux.ai/platform/management"), }, ProviderConfigInfo { id: ProviderId::ClinePass, name: "ClinePass", - requires_api_key: false, api_key_env_var: Some("CLINE_API_KEY"), api_key_help: Some( "Paste an API key, or run cline auth. Reads the existing Cline session without copying it. Also accepts CLINEPASS_API_KEY.", ), - config_file_path: None, dashboard_url: Some("https://app.cline.bot/dashboard/subscription?personal=true"), }, ProviderConfigInfo { id: ProviderId::Neuralwatt, name: "Neuralwatt", - requires_api_key: true, api_key_env_var: Some("NEURALWATT_API_KEY"), api_key_help: Some("Get your API key from portal.neuralwatt.com."), - config_file_path: None, dashboard_url: Some("https://portal.neuralwatt.com/dashboard"), }, ProviderConfigInfo { id: ProviderId::DevPass, name: "DevPass", - requires_api_key: true, api_key_env_var: Some("DEVPASS_API_KEY"), api_key_help: Some( "Use a regular LLM Gateway API key. Publishable keys and end-user sessions cannot read plan state.", ), - config_file_path: None, dashboard_url: Some("https://devpass.llmgateway.io/dashboard"), }, ProviderConfigInfo { id: ProviderId::XKiro, name: "xKiro", - requires_api_key: true, api_key_env_var: Some("XKIRO_API_KEY"), api_key_help: Some( "Create an API key at xkiro.com. It is sent only to api.xkiro.com and reads the free usage endpoint.", ), - config_file_path: None, dashboard_url: Some("https://xkiro.com"), }, ProviderConfigInfo { id: ProviderId::Vercel, name: "Vercel AI Gateway", - requires_api_key: true, api_key_env_var: Some("AI_GATEWAY_API_KEY"), api_key_help: Some( "Add a Vercel AI Gateway API key to show the team's credit balance and lifetime spend.", ), - config_file_path: None, dashboard_url: Some("https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai-gateway"), }, ProviderConfigInfo { id: ProviderId::Doubao, name: "Doubao / Volcengine Ark", - requires_api_key: true, api_key_env_var: Some( "ARK_API_KEY or VOLCENGINE_ACCESS_KEY_ID + VOLCENGINE_SECRET_ACCESS_KEY", ), api_key_help: Some( "Use ARK_API_KEY for chat probe fallback, or paste Coding Plan credentials as access_key|secret_key|region (region defaults to cn-beijing).", ), - config_file_path: None, dashboard_url: Some("https://console.volcengine.com/ark/region:ark+cn-beijing/usage"), }, ProviderConfigInfo { id: ProviderId::StepFun, name: "StepFun", - requires_api_key: true, api_key_env_var: Some("STEPFUN_OASIS_TOKEN"), api_key_help: Some("Paste an existing Oasis-Token from StepFun login."), - config_file_path: None, dashboard_url: Some("https://platform.stepfun.com/dashboard"), }, ProviderConfigInfo { id: ProviderId::Venice, name: "Venice", - requires_api_key: true, api_key_env_var: Some("VENICE_API_KEY"), api_key_help: Some("Get your API key from Venice settings."), - config_file_path: None, dashboard_url: Some("https://venice.ai/settings/api"), }, ProviderConfigInfo { id: ProviderId::OpenAIApi, name: "OpenAI", - requires_api_key: true, api_key_env_var: Some("OPENAI_ADMIN_KEY / OPENAI_API_KEY"), api_key_help: Some( "Use an OpenAI Admin API key for usage, or a platform key for legacy billing balance.", ), - config_file_path: None, dashboard_url: Some("https://platform.openai.com/usage"), }, ProviderConfigInfo { id: ProviderId::Grok, name: "Grok", - requires_api_key: false, api_key_env_var: None, api_key_help: Some("Uses Grok browser cookies or ~/.grok/auth.json."), - config_file_path: Some("~/.grok/auth.json"), dashboard_url: Some("https://grok.com/settings/subscription"), }, ProviderConfigInfo { id: ProviderId::Xai, name: "xAI", - requires_api_key: true, api_key_env_var: Some("XAI_MANAGEMENT_API_KEY"), api_key_help: Some( "Create a Management API key at console.x.ai under Settings > Management Keys (inference keys are rejected). Team ID goes in provider workspace settings or XAI_TEAM_ID.", ), - config_file_path: Some("~/.codexbar/config.json"), dashboard_url: Some("https://console.x.ai"), }, ProviderConfigInfo { id: ProviderId::ElevenLabs, name: "ElevenLabs", - requires_api_key: true, api_key_env_var: Some("ELEVENLABS_API_KEY / XI_API_KEY"), api_key_help: Some("Get your API key from ElevenLabs Settings > API Keys."), - config_file_path: None, dashboard_url: Some("https://elevenlabs.io/app/settings/api-keys"), }, ProviderConfigInfo { id: ProviderId::Deepgram, name: "Deepgram", - requires_api_key: true, api_key_env_var: Some("DEEPGRAM_API_KEY"), api_key_help: Some("Use a Deepgram API key with Management API access."), - config_file_path: None, dashboard_url: Some("https://console.deepgram.com/usage"), }, ProviderConfigInfo { id: ProviderId::Groq, name: "Groq", - requires_api_key: true, api_key_env_var: Some("GROQ_API_KEY"), api_key_help: Some( "Usage & spend come from your console.groq.com browser session automatically. \ An API key is optional and only adds Enterprise Prometheus metrics.", ), - config_file_path: None, dashboard_url: Some("https://console.groq.com/settings/metrics"), }, ProviderConfigInfo { id: ProviderId::LLMProxy, name: "LLM Proxy", - requires_api_key: true, api_key_env_var: Some("LLM_PROXY_API_KEY + LLM_PROXY_BASE_URL"), api_key_help: Some( "Set an LLM Proxy API key and base URL (Settings or LLM_PROXY_BASE_URL) for quota-stats.", ), - config_file_path: None, dashboard_url: None, }, ProviderConfigInfo { id: ProviderId::Chutes, name: "Chutes", - requires_api_key: true, api_key_env_var: Some("CHUTES_API_KEY"), api_key_help: Some( "Paste a Chutes API key. Optional API URL override: CHUTES_API_URL.", ), - config_file_path: None, dashboard_url: Some("https://chutes.ai"), }, ProviderConfigInfo { id: ProviderId::LiteLLM, name: "LiteLLM", - requires_api_key: true, api_key_env_var: Some("LITELLM_API_KEY + LITELLM_BASE_URL"), api_key_help: Some( "Paste a LiteLLM key and set the base URL in provider extras or LITELLM_BASE_URL.", ), - config_file_path: None, dashboard_url: None, }, ProviderConfigInfo { id: ProviderId::LLMMan, name: "llmman", - requires_api_key: false, api_key_env_var: Some("LLMMAN_API_KEY + LLMMAN_HOST"), api_key_help: Some( "Optional: an open local daemon needs no key. Set the base URL in provider extras or LLMMAN_HOST (default http://127.0.0.1:17434).", ), - config_file_path: None, dashboard_url: None, }, ProviderConfigInfo { id: ProviderId::Poe, name: "Poe", - requires_api_key: true, api_key_env_var: Some("POE_API_KEY"), api_key_help: Some("Get your API key from Poe API settings."), - config_file_path: None, dashboard_url: Some("https://poe.com/settings/subscription"), }, ProviderConfigInfo { id: ProviderId::Devin, name: "Devin", - requires_api_key: true, api_key_env_var: Some("DEVIN_BEARER_TOKEN + DEVIN_ORG"), api_key_help: Some( "Paste a Devin bearer token and set the organization in provider extras or DEVIN_ORG.", ), - config_file_path: None, dashboard_url: Some("https://app.devin.ai/settings/billing"), }, ProviderConfigInfo { id: ProviderId::Zed, name: "Zed", - requires_api_key: true, api_key_env_var: Some("ZED_CREDENTIALS"), api_key_help: Some( "Paste Zed credentials as `user_id access_token`; optional API URL in provider extras.", ), - config_file_path: Some("~/.config/zed/settings.json"), dashboard_url: Some("https://zed.dev/account"), }, ProviderConfigInfo { id: ProviderId::CrossModel, name: "CrossModel", - requires_api_key: true, api_key_env_var: Some("CROSSMODEL_API_KEY"), api_key_help: Some( "Paste a CrossModel API key. Optional API URL override: CROSSMODEL_API_URL.", ), - config_file_path: None, dashboard_url: Some("https://crossmodel.ai"), }, ProviderConfigInfo { id: ProviderId::Sub2Api, name: "sub2api", - requires_api_key: true, api_key_env_var: Some("SUB2API_API_KEY"), api_key_help: Some( "Paste a group API key and set the base URL in provider extras or SUB2API_BASE_URL. HTTPS required (loopback HTTP allowed for local dev).", ), - config_file_path: None, dashboard_url: None, }, ProviderConfigInfo { id: ProviderId::Factory, name: "Droid (Factory)", - requires_api_key: true, api_key_env_var: Some("FACTORY_API_KEY"), api_key_help: Some( "Get your API key from Factory → Settings → API Keys. Optional fallback: %USERPROFILE%\\.factory\\.env. Auto mode tries the key first, then browser cookies.", ), - config_file_path: Some("%USERPROFILE%\\.factory\\.env"), dashboard_url: Some("https://app.factory.ai/settings/api-keys"), }, ProviderConfigInfo { id: ProviderId::Meta, name: "Meta", - requires_api_key: true, api_key_env_var: Some("MODEL_API_KEY / META_API_KEY"), api_key_help: Some("Create key in Meta Model API dashboard"), - config_file_path: None, dashboard_url: Some("https://dev.meta.ai/docs"), }, ProviderConfigInfo { id: ProviderId::Hyper, name: "Charm Hyper", - requires_api_key: false, api_key_env_var: Some("HYPER_API_KEY"), api_key_help: Some( "Fallback when no hyper.charm.land session is available. Save an API key here or set HYPER_API_KEY.", ), - config_file_path: None, dashboard_url: Some("https://hyper.charm.land"), }, ProviderConfigInfo { id: ProviderId::GitKraken, name: "GitKraken AI", - requires_api_key: true, api_key_env_var: Some("GITKRAKEN_API_TOKEN"), api_key_help: Some( "Save a GitKraken access token. Optional organization ID: provider extras or GITKRAKEN_ORG_ID.", ), - config_file_path: None, dashboard_url: Some("https://gitkraken.dev/account#ai-usage"), }, ProviderConfigInfo { id: ProviderId::Bifrost, name: "Bifrost", - requires_api_key: true, api_key_env_var: Some("BIFROST_API_KEY"), api_key_help: Some( "Save a Bifrost virtual key and configure the gateway base URL in provider settings. Or set BIFROST_API_KEY and BIFROST_BASE_URL.", ), - config_file_path: None, dashboard_url: None, }, ProviderConfigInfo { id: ProviderId::Aixy, name: "Aixy", - requires_api_key: true, api_key_env_var: Some("AIXY_API_KEY"), api_key_help: Some( "Save an Aixy API key. Leave the Base URL empty for the hosted gateway, or set it for a self-hosted one. Or set AIXY_API_KEY and AIXY_BASE_URL.", ), - config_file_path: None, dashboard_url: Some("https://dash.aixy-gateway.com"), }, ] diff --git a/rust/src/settings/cost_time_zone.rs b/rust/src/settings/cost_time_zone.rs index f50e43a466..af2edaece6 100644 --- a/rust/src/settings/cost_time_zone.rs +++ b/rust/src/settings/cost_time_zone.rs @@ -4,10 +4,13 @@ //! Day buckets, month-to-date bounds, and cost caches use one saved IANA zone //! so history keeps its day boundaries when the machine's zone changes. An //! empty value means the machine zone. There is no settings control for it, -//! as upstream: the desktop app pins the machine zone on first launch. +//! as upstream. The saved zone is stored but not yet applied at runtime: +//! nothing in the desktop app or the CLI pins it on first launch or applies it +//! to bucketing (the wiring PRs #653 and #665 were closed), so history is +//! bucketed in the machine zone. use super::Settings; -use crate::cost_reporting_period::{CostTimeZone, set_cost_bucket_zone}; +use crate::cost_reporting_period::CostTimeZone; /// Trim a saved bucket zone; anything but an IANA zone name reads as unpinned /// (`""`). @@ -43,11 +46,6 @@ impl Settings { None => false, } } - - /// Bucket this process's local cost history in the saved zone. - pub fn apply_cost_usage_bucket_zone(&self) -> CostTimeZone { - set_cost_bucket_zone(&self.cost_usage_bucket_time_zone) - } } #[cfg(test)] diff --git a/rust/src/settings/persistence.rs b/rust/src/settings/persistence.rs new file mode 100644 index 0000000000..ad672bd2a1 --- /dev/null +++ b/rust/src/settings/persistence.rs @@ -0,0 +1,129 @@ +//! Loading, migrating and saving `settings.json`. + +use super::*; + +impl Settings { + /// Get the settings file path (`CODEXBAR_CONFIG` when the CLI applied it). + pub fn settings_path() -> Option { + config_path::settings_file_for(settings_file_override(), crate::logging::config_root()) + } + + /// Load settings from disk + pub fn load() -> Self { + let path = Self::settings_path(); + #[allow( + unused_mut, + reason = "mutability is needed for conditional initialization paths that the compiler cannot prove" + )] + let mut settings = Self::load_from_path(path.as_deref()); + + // Sync autostart toggle with actual registry state and repair stale commands from older builds. + // A CLI run on a `CODEXBAR_CONFIG` file skips these desktop integrations, so it touches + // neither the Run key nor the default config root. + #[cfg(target_os = "windows")] + if settings_file_override().is_none() { + settings.start_at_login = Self::sync_start_at_login_registry(); + settings.apply_promote_tray_default_migration(); + } + + // One-shot migration: materialize the persisted hidden usage-item list + // from the pre-0.62 per-provider visibility flags. After this the list + // is the sole source of truth and the flags stay untouched. + settings.migrate_legacy_usage_item_flags(); + + settings + } + + pub(super) fn load_from_path(path: Option<&Path>) -> Self { + match path { + Some(path) if path.exists() => match crate::secure_file::read_string(path) { + Ok(content) => { + serde_json::from_str(content.trim_start_matches('\u{feff}')).unwrap_or_default() + } + Err(_) => Self::default(), + }, + _ => Self::default(), + } + } + + /// Materialize `hidden_usage_item_ids` from the pre-0.62 per-provider + /// visibility flags where the list was never persisted. Idempotent: a + /// provider with an explicit list is left alone. + pub(super) fn migrate_legacy_usage_item_flags(&mut self) { + if self + .provider_config(ProviderId::Codex) + .and_then(|config| config.hidden_usage_item_ids.as_ref()) + .is_none() + && !self.spark_usage_visible(ProviderId::Codex) + { + self.toggle_hidden_items(ProviderId::Codex, &CODEX_SPARK_USAGE_ITEM_IDS, false); + } + if self + .provider_config(ProviderId::Claude) + .and_then(|config| config.hidden_usage_item_ids.as_ref()) + .is_none() + && !self.claude_daily_routines_usage_visible + { + self.toggle_hidden_items( + ProviderId::Claude, + &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], + false, + ); + } + } + + /// Marker written after the one-shot "pin tray by default" migration (issue #237). + pub(super) fn promote_tray_default_marker_path() -> Option { + crate::logging::config_root().map(|p| p.join(".tray-pin-default-v1")) + } + + /// Old builds defaulted `promote_tray_icon` to false and persisted that on any + /// settings save. Flip those installs to the new default once; later opt-outs + /// are preserved because the marker file remains. + pub(super) fn should_migrate_promote_tray_default( + promote_tray_icon: bool, + already_migrated: bool, + ) -> bool { + !already_migrated && !promote_tray_icon + } + + pub(super) fn apply_promote_tray_default_migration(&mut self) { + let Some(marker) = Self::promote_tray_default_marker_path() else { + return; + }; + let already_migrated = marker.exists(); + if Self::should_migrate_promote_tray_default(self.promote_tray_icon, already_migrated) { + self.promote_tray_icon = true; + if let Err(error) = self.save() { + tracing::warn!("Failed to persist promote_tray_icon default migration: {error}"); + } + } + if !already_migrated && let Some(parent) = marker.parent() { + // Best-effort marker dir creation; the write below reports failure. + let _created_dir = std::fs::create_dir_all(parent); + if let Err(error) = std::fs::write(&marker, b"1") { + tracing::warn!("Failed to write promote_tray_icon migration marker: {error}"); + } + } + } + + /// Save settings to disk + pub fn save(&self) -> anyhow::Result<()> { + let path = Self::settings_path() + .ok_or_else(|| anyhow::anyhow!("Could not determine settings path"))?; + + self.save_to_path(&path) + } + + pub(super) fn save_to_path(&self, path: &Path) -> anyhow::Result<()> { + // Ensure parent directory exists + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + + let json = serde_json::to_string_pretty(self)?; + crate::secure_file::write_string(path, &json)?; + + Ok(()) + } +} diff --git a/rust/src/settings/provider_accessors.rs b/rust/src/settings/provider_accessors.rs new file mode 100644 index 0000000000..1cf5523ffe --- /dev/null +++ b/rust/src/settings/provider_accessors.rs @@ -0,0 +1,408 @@ +//! Per-provider configuration accessors and legacy field-name aliases. + +use super::*; + +impl Settings { + // ── Per-provider configuration accessors ───────────────────────── + // + // These thin wrappers around `provider_configs` apply provider-specific + // defaults (e.g. cookie/usage source defaults to `"auto"`) so callers + // never have to reach into the raw `Option` fields. The + // `*_str` / boolean / setter pairs intentionally mirror the names of + // the legacy flat fields so call-site migration is mechanical. + + /// Read-only access to a provider's stored config, if any. + pub fn provider_config(&self, id: ProviderId) -> Option<&ProviderConfig> { + self.provider_configs.get(&id) + } + + /// Mutable access to a provider's config, lazily creating an empty + /// entry if none exists. + pub fn provider_config_mut(&mut self, id: ProviderId) -> &mut ProviderConfig { + self.provider_configs.entry(id).or_default() + } + + /// Cookie source for `id`. Kimi, Charm Hyper and Groq follow upstream's + /// automatic default; providers with no specific default retain the + /// legacy manual default. + pub fn cookie_source(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.cookie_source.as_deref()) + .unwrap_or(match id { + ProviderId::Kimi | ProviderId::Hyper | ProviderId::Groq => "auto", + _ => DEFAULT_COOKIE_SOURCE, + }) + } + + pub fn set_cookie_source(&mut self, id: ProviderId, source: impl Into) { + self.provider_config_mut(id).cookie_source = Some(source.into()); + } + + /// Usage source for `id`, or the default `"auto"` if unset. + pub fn usage_source(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.usage_source.as_deref()) + .unwrap_or(DEFAULT_PROVIDER_SOURCE) + } + + pub fn set_usage_source(&mut self, id: ProviderId, source: impl Into) { + self.provider_config_mut(id).usage_source = Some(source.into()); + } + + /// API region for `id`, or the provider-specific default if unset. + pub fn api_region(&self, id: ProviderId) -> &str { + if id == ProviderId::AlibabaTokenPlan && !self.alibaba_token_plan_region.trim().is_empty() { + return self.alibaba_token_plan_region.as_str(); + } + self.provider_configs + .get(&id) + .and_then(|c| c.api_region.as_deref()) + .unwrap_or_else(|| default_api_region(id)) + } + + pub fn set_api_region(&mut self, id: ProviderId, region: impl Into) { + let region = region.into(); + if id == ProviderId::AlibabaTokenPlan { + self.alibaba_token_plan_region = region.clone(); + } + self.provider_config_mut(id).api_region = Some(region); + } + + /// Manual cookie header for `id`, or `""` if unset. + pub fn manual_cookie_header(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.manual_cookie_header.as_deref()) + .unwrap_or("") + } + + /// API token for `id`, or `""` if unset. + pub fn api_token(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.api_token.as_deref()) + .unwrap_or("") + } + + pub fn management_api_token(&self, id: ProviderId) -> Option<&str> { + self.provider_configs + .get(&id) + .and_then(|config| config.management_api_token.as_deref()) + .map(str::trim) + .filter(|value| !value.is_empty()) + } + + pub fn set_management_api_token(&mut self, id: ProviderId, token: Option) { + let token = token + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + self.provider_config_mut(id).management_api_token = token; + } + + /// Workspace ID override for `id`, or `""` if unset. + pub fn workspace_id(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.workspace_id.as_deref()) + .unwrap_or("") + } + + pub fn set_workspace_id(&mut self, id: ProviderId, value: impl Into) { + self.provider_config_mut(id).workspace_id = Some(value.into()); + } + + /// Optional user-entered allowance for Copilot seat AI credits. + /// + /// GitHub reports the absolute `credits_used` counter but does not expose + /// a documented included-credit ceiling, so callers must keep an absent + /// or non-positive value as unknown rather than inventing a denominator. + /// + /// This setter is the single owner of the positive-finite invariant: + /// invalid values are rejected instead of silently dropped, while the + /// getter keeps defensively filtering values persisted by older builds. + pub fn seat_credit_entitlement(&self, id: ProviderId) -> Option { + self.provider_configs + .get(&id) + .and_then(|config| config.seat_credit_entitlement) + .filter(|value| value.is_finite() && *value > 0.0) + } + + pub fn set_seat_credit_entitlement( + &mut self, + id: ProviderId, + value: Option, + ) -> Result<(), String> { + if let Some(value) = value + && (!value.is_finite() || value <= 0.0) + { + return Err( + "Copilot seat AI-credit allowance must be a finite number greater than zero" + .to_string(), + ); + } + self.provider_config_mut(id).seat_credit_entitlement = value; + Ok(()) + } + + /// Wayfinder gateway URL, defaulting to the local loopback gateway. + pub fn gateway_url(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.gateway_url.as_deref()) + .unwrap_or_else(|| { + if id == ProviderId::Wayfinder { + crate::providers::wayfinder::DEFAULT_GATEWAY_URL + } else { + "" + } + }) + } + + pub fn set_gateway_url(&mut self, id: ProviderId, value: impl Into) { + self.provider_config_mut(id).gateway_url = Some(value.into()); + } + + /// IDE base path override for `id`, or `""` if unset. + pub fn ide_base_path(&self, id: ProviderId) -> &str { + self.provider_configs + .get(&id) + .and_then(|c| c.ide_base_path.as_deref()) + .unwrap_or("") + } + + pub fn set_ide_base_path(&mut self, id: ProviderId, value: impl Into) { + self.provider_config_mut(id).ide_base_path = Some(value.into()); + } + + /// Codex `openai_web_extras` toggle, default `true`. + pub fn openai_web_extras(&self, id: ProviderId) -> bool { + self.provider_configs + .get(&id) + .and_then(|c| c.openai_web_extras) + .unwrap_or(DEFAULT_CODEX_OPENAI_WEB_EXTRAS) + } + + /// Codex Spark rows are visible by default. + pub fn spark_usage_visible(&self, id: ProviderId) -> bool { + self.provider_configs + .get(&id) + .and_then(|c| c.spark_usage_visible) + .unwrap_or(DEFAULT_CODEX_SPARK_USAGE_VISIBLE) + } + + pub fn set_spark_usage_visible(&mut self, id: ProviderId, value: bool) { + self.provider_config_mut(id).spark_usage_visible = Some(value); + } + + /// Return the persisted hidden usage-item IDs for `id`. + pub fn hidden_usage_item_ids(&self, id: ProviderId) -> Vec { + self.provider_configs + .get(&id) + .and_then(|config| config.hidden_usage_item_ids.as_ref()) + .map_or_else(Vec::new, |ids| normalize_hidden_usage_item_ids(ids.clone())) + } + + /// Persist an explicit presentation-only usage-item visibility list. + pub fn set_hidden_usage_item_ids(&mut self, id: ProviderId, ids: Vec) { + let hidden = normalize_hidden_usage_item_ids(ids); + self.provider_config_mut(id).hidden_usage_item_ids = Some(hidden); + } + + /// Add or remove `items` from the provider's hidden usage-item list. + /// + /// `visible = false` hides the items; `visible = true` un-hides them. This + /// is the single write path for usage-item visibility; it only touches the + /// presentation list and never the legacy per-provider boolean flags. + pub fn toggle_hidden_items(&mut self, id: ProviderId, items: &[&str], visible: bool) { + let mut hidden = self.hidden_usage_item_ids(id); + if visible { + hidden.retain(|item| !items.contains(&item.as_str())); + } else { + hidden.extend(items.iter().map(|item| (*item).to_string())); + } + self.set_hidden_usage_item_ids(id, hidden); + } + + /// Update the old Claude Daily Routines flag without touching the + /// usage-item list; the flag is presentation-only and kept for callers + /// that still read the boolean directly. + pub fn set_claude_daily_routines_usage_visible(&mut self, value: bool) { + self.claude_daily_routines_usage_visible = value; + } + + /// Per-provider historical-tracking toggle (currently codex-only). + pub fn historical_tracking(&self, id: ProviderId) -> bool { + self.provider_configs + .get(&id) + .map(|c| c.historical_tracking) + .unwrap_or(false) + } + + /// Per-provider "avoid keychain prompts" toggle (currently claude-only). + pub fn avoid_keychain_prompts(&self, id: ProviderId) -> bool { + self.provider_configs + .get(&id) + .map(|c| c.avoid_keychain_prompts) + .unwrap_or(false) + } + + pub fn set_avoid_keychain_prompts(&mut self, id: ProviderId, value: bool) { + self.provider_config_mut(id).avoid_keychain_prompts = value; + } + + /// Whether the desktop shell may reopen a captured CLI session after its + /// provider quota becomes available again. This is intentionally opt-in. + pub fn auto_resume_after_quota_reset(&self, id: ProviderId) -> bool { + self.provider_configs + .get(&id) + .map(|config| config.auto_resume_after_quota_reset) + .unwrap_or(false) + } + + pub fn set_auto_resume_after_quota_reset(&mut self, id: ProviderId, value: bool) { + self.provider_config_mut(id).auto_resume_after_quota_reset = value; + } + + // ── Legacy field-name aliases ──────────────────────────────────── + // + // Keep the names of the old flat per-provider fields available as + // accessor methods so existing call sites only need a `()` (read) or + // `set_` prefix (write). New code should prefer the typed accessors + // above. + + pub fn codex_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Codex) + } + pub fn claude_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Claude) + } + pub fn cursor_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Cursor) + } + pub fn opencode_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::OpenCode) + } + pub fn factory_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Factory) + } + pub fn alibaba_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Alibaba) + } + pub fn kimi_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Kimi) + } + pub fn minimax_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::MiniMax) + } + pub fn augment_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Augment) + } + pub fn amp_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Amp) + } + pub fn ollama_cookie_source(&self) -> &str { + self.cookie_source(ProviderId::Ollama) + } + + pub fn claude_usage_source(&self) -> &str { + self.usage_source(ProviderId::Claude) + } + pub fn codex_usage_source(&self) -> &str { + self.usage_source(ProviderId::Codex) + } + + pub fn alibaba_api_region(&self) -> &str { + self.api_region(ProviderId::Alibaba) + } + pub fn zai_api_region(&self) -> &str { + self.api_region(ProviderId::Zai) + } + pub fn minimax_api_region(&self) -> &str { + self.api_region(ProviderId::MiniMax) + } + + pub fn alibaba_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::Alibaba) + } + pub fn kimi_manual_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::Kimi) + } + pub fn augment_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::Augment) + } + pub fn amp_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::Amp) + } + pub fn ollama_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::Ollama) + } + pub fn minimax_cookie_header(&self) -> &str { + self.manual_cookie_header(ProviderId::MiniMax) + } + + pub fn opencode_workspace_id(&self) -> &str { + self.workspace_id(ProviderId::OpenCode) + } + pub fn minimax_api_token(&self) -> &str { + self.api_token(ProviderId::MiniMax) + } + pub fn jetbrains_ide_base_path(&self) -> &str { + self.ide_base_path(ProviderId::JetBrains) + } + pub fn set_jetbrains_ide_base_path(&mut self, v: impl Into) { + self.set_ide_base_path(ProviderId::JetBrains, v) + } + + pub fn codex_openai_web_extras(&self) -> bool { + self.openai_web_extras(ProviderId::Codex) + } + pub fn codex_spark_usage_visible(&self) -> bool { + self.spark_usage_visible(ProviderId::Codex) + } + pub fn set_codex_spark_usage_visible(&mut self, v: bool) { + self.set_spark_usage_visible(ProviderId::Codex, v) + } + pub fn codex_historical_tracking(&self) -> bool { + self.historical_tracking(ProviderId::Codex) + } + pub fn claude_avoid_keychain_prompts(&self) -> bool { + self.avoid_keychain_prompts(ProviderId::Claude) + } + pub fn set_claude_avoid_keychain_prompts(&mut self, v: bool) { + self.set_avoid_keychain_prompts(ProviderId::Claude, v) + } + + /// Claude-only: whether the external claude-swap (`cswap`) adapter is + /// enabled. Disabled by default. + pub fn claude_swap_enabled(&self) -> bool { + self.provider_configs + .get(&ProviderId::Claude) + .map(|config| config.claude_swap_enabled) + .unwrap_or(false) + } + + pub fn set_claude_swap_enabled(&mut self, value: bool) { + self.provider_config_mut(ProviderId::Claude) + .claude_swap_enabled = value; + } + + /// Claude-only: configured claude-swap executable path, or `""` when unset. + pub fn claude_swap_executable_path(&self) -> &str { + self.provider_configs + .get(&ProviderId::Claude) + .and_then(|config| config.claude_swap_executable_path.as_deref()) + .unwrap_or("") + } + + pub fn set_claude_swap_executable_path(&mut self, value: impl Into) { + let trimmed = value.into().trim().to_string(); + let config = self.provider_config_mut(ProviderId::Claude); + config.claude_swap_executable_path = if trimmed.is_empty() { + None + } else { + Some(trimmed) + } + } +} diff --git a/rust/src/settings/start_at_login.rs b/rust/src/settings/start_at_login.rs new file mode 100644 index 0000000000..20d1c02a8a --- /dev/null +++ b/rust/src/settings/start_at_login.rs @@ -0,0 +1,125 @@ +//! The Windows "start at login" Run-key registration. + +use super::*; + +impl Settings { + pub(super) fn start_at_login_exe_path(current_exe: &std::path::Path) -> std::path::PathBuf { + let file_name = current_exe.file_name().and_then(|name| name.to_str()); + if file_name.is_some_and(|name| { + name.eq_ignore_ascii_case("codexbar-cli.exe") + || name.eq_ignore_ascii_case("codexbar-desktop.exe") + }) && let Some(desktop_exe) = current_exe + .parent() + .map(|dir| dir.join("codexbar.exe")) + .filter(|path| path.exists()) + { + return desktop_exe; + } + + current_exe.to_path_buf() + } + + pub(super) fn start_at_login_command(current_exe: &std::path::Path) -> String { + let exe_path = Self::start_at_login_exe_path(current_exe); + format!("\"{}\"", exe_path.display()) + } + + pub(super) fn start_at_login_command_needs_repair( + existing: &str, + current_exe: &std::path::Path, + ) -> bool { + existing != Self::start_at_login_command(current_exe) + } + + #[cfg(target_os = "windows")] + pub fn apply_start_at_login_registry(enabled: bool) -> anyhow::Result<()> { + use winreg::RegKey; + use winreg::enums::*; + + let hkcu = RegKey::predef(HKEY_CURRENT_USER); + let run_key = hkcu.open_subkey_with_flags( + r"Software\Microsoft\Windows\CurrentVersion\Run", + KEY_READ | KEY_WRITE, + )?; + + if enabled { + let exe_path = std::env::current_exe()?; + let command = Self::start_at_login_command(&exe_path); + run_key.set_value("CodexBar", &command)?; + } else { + // Best-effort removal; a missing value means the desired state already. + let _removed_value = run_key.delete_value("CodexBar"); + } + + Ok(()) + } + + #[cfg(target_os = "windows")] + pub(super) fn sync_start_at_login_registry() -> bool { + use winreg::RegKey; + use winreg::enums::*; + + let hkcu = RegKey::predef(HKEY_CURRENT_USER); + let Ok(run_key) = hkcu.open_subkey_with_flags( + r"Software\Microsoft\Windows\CurrentVersion\Run", + KEY_READ | KEY_WRITE, + ) else { + return false; + }; + + let Ok(existing) = run_key.get_value::("CodexBar") else { + return false; + }; + + match std::env::current_exe() { + Ok(exe_path) if Self::start_at_login_command_needs_repair(&existing, &exe_path) => { + let command = Self::start_at_login_command(&exe_path); + if let Err(error) = run_key.set_value("CodexBar", &command) { + tracing::warn!("Failed to repair CodexBar start-at-login command: {error}"); + } + } + Err(error) => { + tracing::warn!( + "Failed to resolve current executable for start-at-login sync: {error}" + ); + } + _ => {} + } + + true + } + + #[cfg(not(target_os = "windows"))] + pub fn apply_start_at_login_registry(_enabled: bool) -> anyhow::Result<()> { + Ok(()) + } + + /// Set start at login (updates Windows registry) + pub fn set_start_at_login(&mut self, enabled: bool) -> anyhow::Result<()> { + self.start_at_login = enabled; + Self::apply_start_at_login_registry(enabled)?; + Ok(()) + } + + /// Check if start at login is actually enabled in registry + #[cfg(target_os = "windows")] + pub fn is_start_at_login_enabled() -> bool { + use winreg::RegKey; + use winreg::enums::*; + + let hkcu = RegKey::predef(HKEY_CURRENT_USER); + if let Ok(run_key) = hkcu.open_subkey(r"Software\Microsoft\Windows\CurrentVersion\Run") { + run_key.get_value::("CodexBar").is_ok() + } else { + false + } + } + + #[cfg(not(target_os = "windows"))] + pub fn is_start_at_login_enabled() -> bool { + false + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/settings/start_at_login/tests.rs b/rust/src/settings/start_at_login/tests.rs new file mode 100644 index 0000000000..6ca78e4b51 --- /dev/null +++ b/rust/src/settings/start_at_login/tests.rs @@ -0,0 +1,81 @@ +//! Run-key command shape and post-update repair. + +use super::*; + +#[test] +fn test_start_at_login_command_uses_only_the_executable_path() { + let path = std::path::PathBuf::from(r"C:\Program Files\CodexBar\codexbar-desktop-tauri.exe"); + let command = Settings::start_at_login_command(&path); + assert_eq!( + command, + "\"C:\\Program Files\\CodexBar\\codexbar-desktop-tauri.exe\"" + ); + assert!(!command.contains("menubar")); +} + +#[test] +fn test_start_at_login_prefers_desktop_sibling_when_called_from_cli() { + let temp = tempfile::tempdir().expect("temp dir"); + let cli_path = temp.path().join("codexbar-cli.exe"); + let desktop_path = temp.path().join("codexbar.exe"); + std::fs::write(&cli_path, b"cli").expect("write cli"); + std::fs::write(&desktop_path, b"desktop").expect("write desktop"); + + let command = Settings::start_at_login_command(&cli_path); + + assert_eq!(command, format!("\"{}\"", desktop_path.display())); +} + +#[test] +fn test_start_at_login_keeps_current_exe_when_desktop_sibling_missing() { + let temp = tempfile::tempdir().expect("temp dir"); + let cli_path = temp.path().join("codexbar-cli.exe"); + std::fs::write(&cli_path, b"cli").expect("write cli"); + + let command = Settings::start_at_login_command(&cli_path); + + assert_eq!(command, format!("\"{}\"", cli_path.display())); +} + +#[test] +fn test_start_at_login_repairs_stale_cli_command_after_update() { + let temp = tempfile::tempdir().expect("temp dir"); + let cli_path = temp.path().join("codexbar-cli.exe"); + let desktop_path = temp.path().join("codexbar.exe"); + std::fs::write(&cli_path, b"cli").expect("write cli"); + std::fs::write(&desktop_path, b"desktop").expect("write desktop"); + let stale_command = format!("\"{}\"", cli_path.display()); + + assert!(Settings::start_at_login_command_needs_repair( + &stale_command, + &desktop_path + )); +} + +#[test] +fn test_start_at_login_keeps_current_desktop_command_after_update() { + let temp = tempfile::tempdir().expect("temp dir"); + let desktop_path = temp.path().join("codexbar.exe"); + std::fs::write(&desktop_path, b"desktop").expect("write desktop"); + let current_command = format!("\"{}\"", desktop_path.display()); + + assert!(!Settings::start_at_login_command_needs_repair( + ¤t_command, + &desktop_path + )); +} + +#[test] +fn test_start_at_login_repairs_legacy_desktop_command_after_update() { + let temp = tempfile::tempdir().expect("temp dir"); + let desktop_path = temp.path().join("codexbar.exe"); + let legacy_desktop_path = temp.path().join("codexbar-desktop.exe"); + std::fs::write(&desktop_path, b"desktop").expect("write desktop"); + std::fs::write(&legacy_desktop_path, b"legacy desktop").expect("write legacy desktop"); + let stale_command = format!("\"{}\"", legacy_desktop_path.display()); + + assert!(Settings::start_at_login_command_needs_repair( + &stale_command, + &legacy_desktop_path + )); +} diff --git a/rust/src/settings/status.rs b/rust/src/settings/status.rs deleted file mode 100644 index 3b15301b0f..0000000000 --- a/rust/src/settings/status.rs +++ /dev/null @@ -1,46 +0,0 @@ -use super::*; - -/// Provider status for settings UI -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ProviderStatus { - pub id: String, - pub name: String, - pub enabled: bool, -} - -/// Refresh interval options -#[derive(Debug, Clone, Serialize)] -pub struct RefreshIntervalOption { - pub value: u64, - pub label: String, -} - -/// Get available refresh interval options -pub fn get_refresh_interval_options() -> Vec { - vec![ - RefreshIntervalOption { - value: 60, - label: "1 minute".to_string(), - }, - RefreshIntervalOption { - value: 120, - label: "2 minutes".to_string(), - }, - RefreshIntervalOption { - value: 300, - label: "5 minutes".to_string(), - }, - RefreshIntervalOption { - value: 600, - label: "10 minutes".to_string(), - }, - RefreshIntervalOption { - value: 900, - label: "15 minutes".to_string(), - }, - RefreshIntervalOption { - value: 1800, - label: "30 minutes".to_string(), - }, - ] -} diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index 06188918ca..681ef84c3f 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -1,5 +1,20 @@ use super::*; +mod language_theme; +mod provider_configs; + +/// Settings parsed from a document that omits every optional field. +fn defaulted() -> Settings { + serde_json::from_str(r#"{ "enabled_providers": [] }"#).expect("minimal settings parse") +} + +/// Serialize then parse back; the JSON is returned for field-name checks. +fn round_trip(settings: &Settings) -> (String, Settings) { + let json = serde_json::to_string(settings).expect("serialize settings"); + let loaded = serde_json::from_str(&json).expect("deserialize settings"); + (json, loaded) +} + #[test] fn test_settings_default() { let settings = Settings::default(); @@ -60,8 +75,7 @@ fn groq_cookie_source_defaults_to_automatic_session_import() { #[test] fn preferred_currency_defaults_validates_and_round_trips() { - let legacy: Settings = serde_json::from_str(r#"{"enabled_providers": []}"#) - .expect("legacy settings without a preferred currency remain valid"); + let legacy = defaulted(); assert_eq!(legacy.preferred_currency_code, "AUTO"); let selected: Settings = @@ -79,16 +93,14 @@ fn preferred_currency_defaults_validates_and_round_trips() { #[test] fn overview_layout_defaults_to_detailed_and_round_trips() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing overview layout defaults to detailed"); + let defaulted = defaulted(); assert_eq!(defaulted.overview_layout, "detailed"); let compact = Settings { overview_layout: "compact".to_string(), ..Settings::default() }; - let json = serde_json::to_string(&compact).expect("serialize overview layout"); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize overview layout"); + let (_, loaded) = round_trip(&compact); assert_eq!(loaded.overview_layout, "compact"); let unknown: Settings = @@ -99,26 +111,22 @@ fn overview_layout_defaults_to_detailed_and_round_trips() { #[test] fn tray_panel_always_on_top_defaults_off_and_round_trips() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing tray panel topmost field defaults off"); + let defaulted = defaulted(); assert!(!defaulted.tray_panel_always_on_top); let enabled = Settings { tray_panel_always_on_top: true, ..Settings::default() }; - let json = serde_json::to_string(&enabled).expect("serialize tray panel topmost setting"); + let (json, loaded) = round_trip(&enabled); assert!(json.contains(r#""tray_panel_always_on_top":true"#)); - let loaded: Settings = - serde_json::from_str(&json).expect("deserialize tray panel topmost setting"); assert!(loaded.tray_panel_always_on_top); } #[test] fn low_power_mode_migrates_legacy_boolean_and_round_trips_preference() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing low power fields defaults off"); + let defaulted = defaulted(); assert_eq!( defaulted.low_power_mode_preference, LowPowerModePreference::Off @@ -133,9 +141,8 @@ fn low_power_mode_migrates_legacy_boolean_and_round_trips_preference() { low_power_mode_preference: LowPowerModePreference::Automatic, ..Settings::default() }; - let json = serde_json::to_string(&automatic).expect("serialize low power preference"); + let (json, loaded) = round_trip(&automatic); assert!(json.contains(r#""low_power_mode_preference":"automatic""#)); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize low power preference"); assert_eq!( loaded.low_power_mode_preference, LowPowerModePreference::Automatic @@ -144,8 +151,7 @@ fn low_power_mode_migrates_legacy_boolean_and_round_trips_preference() { #[test] fn open_codex_usage_logs_default_off_and_round_trip() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing open_codex_usage_logs_enabled defaults false"); + let defaulted = defaulted(); assert!(!defaulted.open_codex_usage_logs_enabled); let enabled = Settings { @@ -153,35 +159,31 @@ fn open_codex_usage_logs_default_off_and_round_trip() { hide_native_codex_cost_when_open_codex_present: true, ..Settings::default() }; - let json = serde_json::to_string(&enabled).expect("serialize OpenCodex usage opt-in"); + let (json, loaded) = round_trip(&enabled); assert!(json.contains(r#""open_codex_usage_logs_enabled":true"#)); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize OpenCodex usage opt-in"); assert!(loaded.open_codex_usage_logs_enabled); assert!(loaded.hide_native_codex_cost_when_open_codex_present); } #[test] fn tray_pace_color_defaults_off_and_round_trips() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing tray pace color defaults false"); + let defaulted = defaulted(); assert!(!defaulted.menu_bar_color_pace); let enabled = Settings { menu_bar_color_pace: true, ..Settings::default() }; - let json = serde_json::to_string(&enabled).expect("serialize tray pace color"); + let (json, loaded) = round_trip(&enabled); assert!(json.contains(r#""menu_bar_color_pace":true"#)); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize tray pace color"); assert!(loaded.menu_bar_color_pace); } #[test] fn cost_reporting_period_defaults_to_thirty_days_and_round_trips() { - let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("missing cost_reporting_period defaults"); + let defaulted = defaulted(); assert_eq!( defaulted.cost_reporting_period, CostReportingPeriod::Rolling(30) @@ -196,9 +198,8 @@ fn cost_reporting_period_defaults_to_thirty_days_and_round_trips() { cost_reporting_period: period, ..Settings::default() }; - let json = serde_json::to_string(&settings).expect("serialize cost period"); + let (json, loaded) = round_trip(&settings); assert!(json.contains(&format!(r#""cost_reporting_period":"{}""#, period.raw()))); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize cost period"); assert_eq!(loaded.cost_reporting_period, period); } } @@ -224,11 +225,9 @@ fn notification_sound_paths_round_trip_and_default_for_existing_settings() { }, ..Settings::default() }; - let json = serde_json::to_string(&settings).expect("serialize notification sound paths"); + let (json, loaded) = round_trip(&settings); assert!(json.contains("\"criticalUsage\":\"C:\\\\sounds\\\\critical.wav\"")); - let loaded: Settings = - serde_json::from_str(&json).expect("deserialize notification sound paths"); assert_eq!( loaded.notification_sound_paths, settings.notification_sound_paths @@ -238,8 +237,7 @@ fn notification_sound_paths_round_trip_and_default_for_existing_settings() { NotificationSoundTheme::CodexBar ); - let legacy: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) - .expect("deserialize settings without notification sound paths"); + let legacy = defaulted(); assert_eq!( legacy.notification_sound_paths, NotificationSoundPaths::default() @@ -362,14 +360,61 @@ fn main_window_scale_defaults_to_100_percent() { } #[test] -fn main_window_scale_clamp_pins_to_supported_range() { - assert_eq!(clamp_window_scale_percent(0), 100); - assert_eq!(clamp_window_scale_percent(99), 100); - assert_eq!(clamp_window_scale_percent(100), 100); - assert_eq!(clamp_window_scale_percent(125), 125); - assert_eq!(clamp_window_scale_percent(180), 180); - assert_eq!(clamp_window_scale_percent(250), 250); - assert_eq!(clamp_window_scale_percent(251), 250); +fn clamp_helpers_pin_to_supported_ranges() { + // (input, expected) per helper: below range lifts to the floor, in range + // passes through, above range drops to the ceiling. + for (input, expected) in [ + (0, 100), + (99, 100), + (100, 100), + (125, 125), + (180, 180), + (250, 250), + (251, 250), + ] { + assert_eq!( + clamp_window_scale_percent(input), + expected, + "window scale {input}" + ); + } + for (input, expected) in [ + (0, 100), + (99, 100), + (100, 100), + (125, 125), + (180, 180), + (200, 200), + (201, 200), + ] { + assert_eq!( + clamp_tray_scale_percent(input), + expected, + "tray scale {input}" + ); + } + // Opacity floors at 30 so the bar isn't accidentally invisible. + for (input, expected) in [ + (0, 30), + (29, 30), + (45, 45), + (80, 80), + (150, 100), + (255, 100), + ] { + assert_eq!( + clamp_float_bar_opacity(input), + expected, + "float bar opacity {input}" + ); + } + for (input, expected) in [(0, 75), (74, 75), (100, 100), (150, 150), (250, 200)] { + assert_eq!( + clamp_float_bar_scale(input), + expected, + "float bar scale {input}" + ); + } } #[test] @@ -389,17 +434,6 @@ fn tray_scale_defaults_to_100_percent() { assert_eq!(settings.tray_scale_percent, 100); } -#[test] -fn tray_scale_clamp_pins_to_supported_range() { - assert_eq!(clamp_tray_scale_percent(0), 100); - assert_eq!(clamp_tray_scale_percent(99), 100); - assert_eq!(clamp_tray_scale_percent(100), 100); - assert_eq!(clamp_tray_scale_percent(125), 125); - assert_eq!(clamp_tray_scale_percent(180), 180); - assert_eq!(clamp_tray_scale_percent(200), 200); - assert_eq!(clamp_tray_scale_percent(201), 200); -} - #[test] fn raw_settings_clamps_tray_scale_on_load() { let json = r#"{ @@ -411,28 +445,6 @@ fn raw_settings_clamps_tray_scale_on_load() { assert_eq!(loaded.tray_scale_percent, 200); } -#[test] -fn float_bar_opacity_clamp_pins_to_supported_range() { - // Below 30 → 30 so the bar isn't accidentally invisible. - assert_eq!(clamp_float_bar_opacity(0), 30); - assert_eq!(clamp_float_bar_opacity(29), 30); - // Within range → unchanged. - assert_eq!(clamp_float_bar_opacity(45), 45); - assert_eq!(clamp_float_bar_opacity(80), 80); - // Above 100 → 100. - assert_eq!(clamp_float_bar_opacity(150), 100); - assert_eq!(clamp_float_bar_opacity(255), 100); -} - -#[test] -fn float_bar_scale_clamp_pins_to_supported_range() { - assert_eq!(clamp_float_bar_scale(0), 75); - assert_eq!(clamp_float_bar_scale(74), 75); - assert_eq!(clamp_float_bar_scale(100), 100); - assert_eq!(clamp_float_bar_scale(150), 150); - assert_eq!(clamp_float_bar_scale(250), 200); -} - #[test] fn float_bar_orientation_normalization_rejects_unknown_values() { assert_eq!(normalize_float_bar_orientation("horizontal"), "horizontal"); @@ -472,8 +484,7 @@ fn float_bar_settings_round_trip_through_raw() { ..Settings::default() }; - let json = serde_json::to_string(&s).expect("serialize"); - let back: Settings = serde_json::from_str(&json).expect("deserialize"); + let (_, back) = round_trip(&s); assert!(back.float_bar_enabled); assert_eq!(back.float_bar_opacity, 65); assert_eq!(back.float_bar_scale, 140); @@ -548,8 +559,7 @@ fn wayfinder_gateway_round_trips_without_changing_settings_paths() { "https://gateway.example.test/wayfinder/", ); - let json = serde_json::to_string(&settings).expect("serialize settings"); - let loaded: Settings = serde_json::from_str(&json).expect("deserialize settings"); + let (_, loaded) = round_trip(&settings); assert_eq!( loaded.gateway_url(ProviderId::Wayfinder), "https://gateway.example.test/wayfinder/" @@ -618,20 +628,6 @@ fn enabled_provider_ids_follow_custom_provider_order() { ); } -#[test] -fn test_settings_get_all_providers_status() { - let settings = Settings::default(); - let status = settings.get_all_providers_status(); - assert_eq!(status.len(), ProviderId::all().len()); - - let claude_status = status.iter().find(|s| s.id == "claude").unwrap(); - assert_eq!(claude_status.name, "Claude"); - assert!(claude_status.enabled); - - let gemini_status = status.iter().find(|s| s.id == "gemini").unwrap(); - assert!(!gemini_status.enabled); -} - #[test] fn test_api_key_provider_catalog_includes_token_providers() { let providers = get_api_key_providers(); @@ -697,14 +693,6 @@ fn test_t3_chat_is_cookie_configured_not_api_key_configured() { ); } -#[test] -fn test_refresh_interval_options() { - let options = get_refresh_interval_options(); - assert!(!options.is_empty()); - assert!(options.iter().any(|o| o.value == 60)); - assert!(options.iter().any(|o| o.value == 300)); -} - #[test] fn test_manual_cookies_default() { let cookies = ManualCookies::default(); @@ -735,782 +723,38 @@ fn api_key_display_mask_is_utf8_safe() { assert_eq!(display[0].masked_key, "🔑🔒漢字...fgh🔐"); } -#[test] -fn test_start_at_login_command_uses_only_the_executable_path() { - let path = std::path::PathBuf::from(r"C:\Program Files\CodexBar\codexbar-desktop-tauri.exe"); - let command = Settings::start_at_login_command(&path); - assert_eq!( - command, - "\"C:\\Program Files\\CodexBar\\codexbar-desktop-tauri.exe\"" - ); - assert!(!command.contains("menubar")); -} - -#[test] -fn test_start_at_login_prefers_desktop_sibling_when_called_from_cli() { - let temp = tempfile::tempdir().expect("temp dir"); - let cli_path = temp.path().join("codexbar-cli.exe"); - let desktop_path = temp.path().join("codexbar.exe"); - std::fs::write(&cli_path, b"cli").expect("write cli"); - std::fs::write(&desktop_path, b"desktop").expect("write desktop"); - - let command = Settings::start_at_login_command(&cli_path); - - assert_eq!(command, format!("\"{}\"", desktop_path.display())); -} - -#[test] -fn test_start_at_login_keeps_current_exe_when_desktop_sibling_missing() { - let temp = tempfile::tempdir().expect("temp dir"); - let cli_path = temp.path().join("codexbar-cli.exe"); - std::fs::write(&cli_path, b"cli").expect("write cli"); - - let command = Settings::start_at_login_command(&cli_path); - - assert_eq!(command, format!("\"{}\"", cli_path.display())); -} - -#[test] -fn test_start_at_login_repairs_stale_cli_command_after_update() { - let temp = tempfile::tempdir().expect("temp dir"); - let cli_path = temp.path().join("codexbar-cli.exe"); - let desktop_path = temp.path().join("codexbar.exe"); - std::fs::write(&cli_path, b"cli").expect("write cli"); - std::fs::write(&desktop_path, b"desktop").expect("write desktop"); - let stale_command = format!("\"{}\"", cli_path.display()); - - assert!(Settings::start_at_login_command_needs_repair( - &stale_command, - &desktop_path - )); -} - -#[test] -fn test_start_at_login_keeps_current_desktop_command_after_update() { - let temp = tempfile::tempdir().expect("temp dir"); - let desktop_path = temp.path().join("codexbar.exe"); - std::fs::write(&desktop_path, b"desktop").expect("write desktop"); - let current_command = format!("\"{}\"", desktop_path.display()); - - assert!(!Settings::start_at_login_command_needs_repair( - ¤t_command, - &desktop_path - )); -} - -#[test] -fn test_start_at_login_repairs_legacy_desktop_command_after_update() { - let temp = tempfile::tempdir().expect("temp dir"); - let desktop_path = temp.path().join("codexbar.exe"); - let legacy_desktop_path = temp.path().join("codexbar-desktop.exe"); - std::fs::write(&desktop_path, b"desktop").expect("write desktop"); - std::fs::write(&legacy_desktop_path, b"legacy desktop").expect("write legacy desktop"); - let stale_command = format!("\"{}\"", legacy_desktop_path.display()); - - assert!(Settings::start_at_login_command_needs_repair( - &stale_command, - &legacy_desktop_path - )); -} - -#[test] -fn test_language_defaults_to_english() { - let settings = Settings::default(); - assert_eq!(settings.ui_language, Language::English); -} - -#[test] -fn test_language_all_variants_available() { - let languages = Language::all(); - assert_eq!(languages.len(), 10); - assert!(languages.contains(&Language::English)); - assert!(languages.contains(&Language::Chinese)); - assert!(languages.contains(&Language::ChineseTraditional)); - assert!(languages.contains(&Language::Japanese)); - assert!(languages.contains(&Language::Korean)); - assert!(languages.contains(&Language::Spanish)); - assert!(languages.contains(&Language::PortugueseBrazil)); - assert!(languages.contains(&Language::Russian)); - assert!(languages.contains(&Language::Turkish)); - assert!(languages.contains(&Language::Ukrainian)); -} - -#[test] -fn test_language_display_names() { - assert_eq!(Language::English.display_name(), "English"); - assert_eq!(Language::Chinese.display_name(), "中文"); - assert_eq!(Language::ChineseTraditional.display_name(), "繁體中文"); - assert_eq!(Language::Japanese.display_name(), "日本語"); - assert_eq!(Language::Russian.display_name(), "Русский"); - assert_eq!(Language::Turkish.display_name(), "Türkçe"); - assert_eq!(Language::Ukrainian.display_name(), "Українська"); - assert_eq!( - Language::PortugueseBrazil.display_name(), - "Português (Brasil)" - ); -} - -#[test] -fn test_language_resolves_brazilian_portuguese_aliases() { - for alias in [ - "portuguesebrazil", - "pt", - "pt-BR", - "Portuguese", - "Português", - "portugues", - "Português (Brasil)", - ] { - assert_eq!( - Language::resolve(alias), - Some(Language::PortugueseBrazil), - "failed to resolve {alias}" - ); - } -} - -#[test] -fn test_language_resolves_russian_aliases() { - assert_eq!(Language::resolve("russian"), Some(Language::Russian)); - assert_eq!(Language::resolve("ru-RU"), Some(Language::Russian)); - assert_eq!(Language::resolve("Русский"), Some(Language::Russian)); -} - -#[test] -fn test_language_resolves_turkish_aliases() { - assert_eq!(Language::resolve("turkish"), Some(Language::Turkish)); - assert_eq!(Language::resolve("tr-TR"), Some(Language::Turkish)); - assert_eq!(Language::resolve("Türkçe"), Some(Language::Turkish)); - assert_eq!(Language::resolve("turkce"), Some(Language::Turkish)); -} - -#[test] -fn test_language_resolves_ukrainian_aliases() { - assert_eq!(Language::resolve("ukrainian"), Some(Language::Ukrainian)); - assert_eq!(Language::resolve("uk"), Some(Language::Ukrainian)); - assert_eq!(Language::resolve("uk-UA"), Some(Language::Ukrainian)); - assert_eq!(Language::resolve("Українська"), Some(Language::Ukrainian)); -} - -#[test] -fn test_settings_load_missing_language_field_defaults_to_english() { - // Simulate loading legacy settings JSON without ui_language field - let legacy_json = r#"{ - "enabled_providers": ["claude", "codex"], - "refresh_interval_secs": 300, - "start_minimized": false, - "ui_language": "english" - }"#; - - let settings: Result = serde_json::from_str(legacy_json); - assert!(settings.is_ok()); - let settings = settings.unwrap(); - assert_eq!(settings.ui_language, Language::English); -} - -#[test] -fn test_settings_roundtrip_with_language() { - use std::io::Write; - use tempfile::NamedTempFile; - - // Create settings with Chinese language - let settings = Settings { - ui_language: Language::Chinese, - ..Settings::default() - }; - - // Save to a temp file - let mut temp_file = NamedTempFile::new().expect("Failed to create temp file"); - let json = serde_json::to_string_pretty(&settings).expect("Failed to serialize settings"); - temp_file - .write_all(json.as_bytes()) - .expect("Failed to write settings"); - let path = temp_file.path().to_path_buf(); - - // Read back and verify - let content = std::fs::read_to_string(&path).expect("Failed to read settings"); - let loaded: Settings = serde_json::from_str(&content).expect("Failed to deserialize settings"); - - assert_eq!(loaded.ui_language, Language::Chinese); -} - -#[test] -fn test_settings_with_utf8_bom_parses_perprovider_tray_mode() { - let json = "\u{feff}{\n \"enabled_providers\": [\"claude\", \"codex\"],\n \"refresh_interval_secs\": 300,\n \"tray_icon_mode\": \"perprovider\"\n }"; - - let settings: Settings = serde_json::from_str(json.trim_start_matches('\u{feff}')).unwrap(); - - assert_eq!(settings.tray_icon_mode, TrayIconMode::PerProvider); -} - -#[test] -fn stacked_tray_mode_preserves_provider_preferences() { - let json = r#"{ - "tray_icon_mode": "stacked", - "stacked_tray_top_provider": "claude", - "stacked_tray_bottom_provider": "codex" - }"#; - - let settings: Settings = serde_json::from_str(json).unwrap(); - - assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); - assert_eq!( - settings.stacked_tray_top_provider.as_deref(), - Some("claude") - ); - assert_eq!( - settings.stacked_tray_bottom_provider.as_deref(), - Some("codex") - ); - - let saved = serde_json::to_string(&settings).unwrap(); - let reloaded: Settings = serde_json::from_str(&saved).unwrap(); - assert_eq!( - reloaded.stacked_tray_top_provider.as_deref(), - Some("claude") - ); - assert_eq!( - reloaded.stacked_tray_bottom_provider.as_deref(), - Some("codex") - ); -} - -#[test] -fn test_language_serde_serialization() { - // Test that Language serializes to lowercase string - let english = Language::English; - let chinese = Language::Chinese; - let chinese_traditional = Language::ChineseTraditional; - - let english_json = serde_json::to_string(&english).unwrap(); - let chinese_json = serde_json::to_string(&chinese).unwrap(); - let chinese_traditional_json = serde_json::to_string(&chinese_traditional).unwrap(); - - assert_eq!(english_json, "\"english\""); - assert_eq!(chinese_json, "\"chinese\""); - assert_eq!(chinese_traditional_json, "\"chinesetraditional\""); -} - -#[test] -fn test_language_serde_deserialization() { - // Test that lowercase strings deserialize correctly - let english: Language = serde_json::from_str("\"english\"").unwrap(); - let chinese: Language = serde_json::from_str("\"chinese\"").unwrap(); - let chinese_traditional: Language = serde_json::from_str("\"chinesetraditional\"").unwrap(); - - assert_eq!(english, Language::English); - assert_eq!(chinese, Language::Chinese); - assert_eq!(chinese_traditional, Language::ChineseTraditional); -} - -#[test] -fn test_language_resolves_traditional_chinese_aliases() { - assert_eq!( - Language::resolve("chinesetraditional"), - Some(Language::ChineseTraditional) - ); - assert_eq!( - Language::resolve("zh-tw"), - Some(Language::ChineseTraditional) - ); - assert_eq!( - Language::resolve("zh-hant-tw"), - Some(Language::ChineseTraditional) - ); - assert_eq!( - Language::resolve("繁體中文"), - Some(Language::ChineseTraditional) - ); -} - -#[test] -fn test_theme_defaults_to_auto() { - let settings = Settings::default(); - assert_eq!(settings.theme, ThemePreference::Auto); -} - -#[test] -fn test_theme_all_variants_available() { - let themes = ThemePreference::all(); - assert_eq!(themes.len(), 3); - assert!(themes.contains(&ThemePreference::Auto)); - assert!(themes.contains(&ThemePreference::Light)); - assert!(themes.contains(&ThemePreference::Dark)); -} - -#[test] -fn test_theme_serde_roundtrip() { - for variant in [ - ThemePreference::Auto, - ThemePreference::Light, - ThemePreference::Dark, - ] { - let encoded = serde_json::to_string(&variant).unwrap(); - let decoded: ThemePreference = serde_json::from_str(&encoded).unwrap(); - assert_eq!(decoded, variant); - } - assert_eq!( - serde_json::to_string(&ThemePreference::Light).unwrap(), - "\"light\"" - ); - assert_eq!( - serde_json::to_string(&ThemePreference::Dark).unwrap(), - "\"dark\"" - ); - assert_eq!( - serde_json::to_string(&ThemePreference::Auto).unwrap(), - "\"auto\"" - ); -} - -#[test] -fn test_settings_missing_theme_defaults_to_auto() { - // Legacy settings JSON without the theme field should still parse. - let legacy_json = r#"{ - "enabled_providers": ["claude", "codex"], - "refresh_interval_secs": 300, - "ui_language": "english" - }"#; - - let settings: Settings = serde_json::from_str(legacy_json).unwrap(); - assert_eq!(settings.theme, ThemePreference::Auto); -} - -#[test] -fn test_settings_roundtrip_with_theme() { - let settings = Settings { - theme: ThemePreference::Dark, - ..Settings::default() - }; - let json = serde_json::to_string(&settings).unwrap(); - let loaded: Settings = serde_json::from_str(&json).unwrap(); - assert_eq!(loaded.theme, ThemePreference::Dark); -} - -// ── Phase 3: provider_configs migration tests ─────────────────────── - -/// Loading a legacy `settings.json` (with flat per-provider fields) -/// must populate `provider_configs` and surface every value through the -/// per-provider accessors. -#[test] -fn test_legacy_per_provider_fields_migrate_into_provider_configs() { - // NOTE: placeholder values only — no real cookies/tokens. - let legacy_json = r#"{ - "enabled_providers": ["claude", "codex"], - "refresh_interval_secs": 300, - "codex_cookie_source": "manual", - "claude_cookie_source": "browser", - "cursor_cookie_source": "manual", - "alibaba_cookie_source": "manual", - "alibaba_cookie_header": "ali=PLACEHOLDER", - "alibaba_api_region": "cn", - "zai_api_region": "cn", - "minimax_api_region": "cn", - "minimax_api_token": "TOK_PLACEHOLDER", - "claude_usage_source": "ccusage", - "codex_usage_source": "manual", - "codex_openai_web_extras": false, - "codex_historical_tracking": true, - "claude_avoid_keychain_prompts": true, - "opencode_workspace_id": "ws_placeholder", - "jetbrains_ide_base_path": "C:/JB" - }"#; - - let settings: Settings = serde_json::from_str(legacy_json).unwrap(); - - // Cookie sources - assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); - assert_eq!(settings.cookie_source(ProviderId::Claude), "browser"); - assert_eq!(settings.cookie_source(ProviderId::Cursor), "manual"); - assert_eq!(settings.cookie_source(ProviderId::Alibaba), "manual"); - // Untouched providers fall through to the default "manual" to avoid - // background browser-cookie reads unless the user opts into Automatic. - assert_eq!(settings.cookie_source(ProviderId::Amp), "manual"); - - // Manual cookie headers + api regions - assert_eq!( - settings.manual_cookie_header(ProviderId::Alibaba), - "ali=PLACEHOLDER" - ); - assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); - assert_eq!(settings.api_region(ProviderId::Zai), "cn"); - assert_eq!(settings.api_region(ProviderId::MiniMax), "cn"); - - // Usage sources - assert_eq!(settings.usage_source(ProviderId::Claude), "ccusage"); - assert_eq!(settings.usage_source(ProviderId::Codex), "manual"); - - // Codex booleans - assert!(!settings.openai_web_extras(ProviderId::Codex)); - assert!(settings.historical_tracking(ProviderId::Codex)); - - // Claude per-provider boolean - assert!(settings.avoid_keychain_prompts(ProviderId::Claude)); - - // Misc per-provider strings - assert_eq!( - settings.workspace_id(ProviderId::OpenCode), - "ws_placeholder" - ); - assert_eq!(settings.api_token(ProviderId::MiniMax), "TOK_PLACEHOLDER"); - assert_eq!(settings.ide_base_path(ProviderId::JetBrains), "C:/JB"); - - // Legacy field-name aliases agree with typed accessors. - assert_eq!(settings.codex_cookie_source(), "manual"); - assert_eq!(settings.alibaba_api_region(), "cn"); - assert!(settings.codex_historical_tracking()); - assert!(!settings.codex_openai_web_extras()); - assert!(settings.claude_avoid_keychain_prompts()); -} - -/// Round-trip: build a `Settings` programmatically via the new map + -/// accessors, serialize, parse back, and assert equality of every -/// per-provider field. -#[test] -fn test_provider_configs_roundtrip() { - let mut settings = Settings::default(); - settings.set_cookie_source(ProviderId::Codex, "manual"); - settings.set_cookie_source(ProviderId::Claude, "browser"); - settings.set_usage_source(ProviderId::Claude, "ccusage"); - settings.set_api_region(ProviderId::Alibaba, "cn"); - settings.set_api_region(ProviderId::Zai, "cn"); - settings.set_manual_cookie_header(ProviderId::Amp, "amp=PLACEHOLDER"); - settings.set_api_token(ProviderId::MiniMax, "TOK_PLACEHOLDER"); - settings.set_workspace_id(ProviderId::OpenCode, "ws_placeholder"); - settings.set_ide_base_path(ProviderId::JetBrains, "C:/JB"); - settings.set_openai_web_extras(ProviderId::Codex, false); - settings.set_historical_tracking(ProviderId::Codex, true); - settings.set_avoid_keychain_prompts(ProviderId::Claude, true); - settings.set_auto_resume_after_quota_reset(ProviderId::Codex, true); - settings - .set_seat_credit_entitlement(ProviderId::Copilot, Some(300.0)) - .expect("valid seat credit entitlement"); - - let json = serde_json::to_string(&settings).unwrap(); - // The legacy flat fields must NOT appear in serialized output. - assert!(!json.contains("\"codex_cookie_source\""), "json: {json}"); - assert!(!json.contains("\"alibaba_api_region\""), "json: {json}"); - assert!( - !json.contains("\"claude_avoid_keychain_prompts\""), - "json: {json}" - ); - assert!(json.contains("\"provider_configs\""), "json: {json}"); - - let loaded: Settings = serde_json::from_str(&json).unwrap(); - assert_eq!(loaded.cookie_source(ProviderId::Codex), "manual"); - assert_eq!(loaded.cookie_source(ProviderId::Claude), "browser"); - assert_eq!(loaded.usage_source(ProviderId::Claude), "ccusage"); - assert_eq!(loaded.api_region(ProviderId::Alibaba), "cn"); - assert_eq!(loaded.api_region(ProviderId::Zai), "cn"); - assert_eq!( - loaded.manual_cookie_header(ProviderId::Amp), - "amp=PLACEHOLDER" - ); - assert_eq!(loaded.api_token(ProviderId::MiniMax), "TOK_PLACEHOLDER"); - assert_eq!(loaded.workspace_id(ProviderId::OpenCode), "ws_placeholder"); - assert_eq!(loaded.ide_base_path(ProviderId::JetBrains), "C:/JB"); - assert!(!loaded.openai_web_extras(ProviderId::Codex)); - assert!(loaded.historical_tracking(ProviderId::Codex)); - assert!(loaded.avoid_keychain_prompts(ProviderId::Claude)); - assert!(loaded.auto_resume_after_quota_reset(ProviderId::Codex)); - assert_eq!( - loaded.seat_credit_entitlement(ProviderId::Copilot), - Some(300.0) - ); - assert_eq!( - loaded.provider_configs.get(&ProviderId::Codex), - settings.provider_configs.get(&ProviderId::Codex) - ); -} - -/// New-format files (no legacy flat fields, only `provider_configs`) -/// must load identically. -#[test] -fn test_new_format_provider_configs_only() { - let json = r#"{ - "enabled_providers": ["claude"], - "refresh_interval_secs": 300, - "provider_configs": { - "codex": { "cookie_source": "manual", "openai_web_extras": false }, - "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } - } - }"#; - - let settings: Settings = serde_json::from_str(json).unwrap(); - assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); - assert!(!settings.openai_web_extras(ProviderId::Codex)); - assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); - assert_eq!( - settings.manual_cookie_header(ProviderId::Alibaba), - "ali=PLACEHOLDER" - ); - // Untouched providers still get their defaults. - assert_eq!(settings.cookie_source(ProviderId::Claude), "manual"); - assert_eq!(settings.api_region(ProviderId::Zai), "global"); -} - -#[test] -fn retired_provider_config_is_ignored_until_explicit_save() { - let original = r#"{ - "enabled_providers": ["codex", "crof"], - "refresh_interval_secs": 300, - "provider_metrics": { "codex": "weekly", "crof": "session" }, - "float_bar_provider_ids": ["codex", "crof"], - "stacked_tray_top_provider": "crof", - "stacked_tray_bottom_provider": "crof", - "provider_configs": { - "crof": { "api_token": "retired-fixture-key" }, - "codex": { "cookie_source": "manual", "openai_web_extras": false }, - "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } - } - }"#; - let original_bytes = original.as_bytes().to_vec(); - - let settings: Settings = - serde_json::from_str(original).expect("load settings with retired key"); - - assert_eq!(original.as_bytes(), original_bytes); - assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); - assert!(!settings.openai_web_extras(ProviderId::Codex)); - assert_eq!( - settings.enabled_providers, - HashSet::from(["codex".to_string()]) - ); - assert_eq!(settings.provider_metrics.len(), 1); - assert_eq!(settings.float_bar_provider_ids, ["codex"]); - assert_eq!(settings.stacked_tray_top_provider, None); - assert_eq!(settings.stacked_tray_bottom_provider, None); - assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); - assert_eq!( - settings.manual_cookie_header(ProviderId::Alibaba), - "ali=PLACEHOLDER" - ); - - let saved = serde_json::to_string(&settings).expect("serialize sanitized settings"); - let saved_value: serde_json::Value = serde_json::from_str(&saved).unwrap(); - let saved_configs = saved_value["provider_configs"].as_object().unwrap(); - assert!(!saved_configs.contains_key("crof")); - assert!(saved_configs.contains_key("codex")); - assert!(saved_configs.contains_key("alibaba")); - assert!( - !saved.contains("\"crof\""), - "saved settings retained Crof: {saved}" - ); -} - -#[test] -fn provider_aliases_are_canonicalized_at_the_load_boundary() { - let settings: Settings = serde_json::from_str( - r#"{ - "enabled_providers": ["openai", "ClAuDe", "not-a-provider"], - "provider_metrics": { - "openai": "weekly", - "CoDeX": "session", - "not-a-provider": "weekly" - }, - "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"], - "stacked_tray_top_provider": "OPENAI", - "stacked_tray_bottom_provider": "ClAuDe" - }"#, - ) - .expect("load settings containing provider aliases"); - - assert_eq!( - settings.enabled_providers, - HashSet::from(["claude".to_string(), "codex".to_string()]) - ); - assert_eq!( - settings.provider_metrics.get("codex"), - Some(&MetricPreference::Session) - ); - assert_eq!(settings.provider_metrics.len(), 1); - assert_eq!(settings.float_bar_provider_ids, ["codex", "claude"]); - assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); - assert_eq!( - settings.stacked_tray_bottom_provider.as_deref(), - Some("claude") - ); -} - -#[test] -fn stacked_preferences_preserve_known_disabled_providers() { - let settings: Settings = serde_json::from_str( - r#"{ - "enabled_providers": ["claude"], - "stacked_tray_top_provider": "OPENAI", - "stacked_tray_bottom_provider": "not-a-provider" - }"#, - ) - .expect("load stacked preferences independently of enablement"); - - assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); - assert_eq!(settings.stacked_tray_bottom_provider, None); - assert_eq!( - settings.enabled_providers, - HashSet::from(["claude".to_string()]) - ); -} - -/// Default `Settings` should serialize WITHOUT a `provider_configs` -/// field (empty map skipped). -#[test] -fn test_default_settings_skip_empty_provider_configs() { - let settings = Settings::default(); - let json = serde_json::to_string(&settings).unwrap(); - assert!( - !json.contains("\"provider_configs\""), - "empty map should be skipped, json: {json}" - ); -} - -/// Per-provider defaults are applied even when the entry is absent. -#[test] -fn test_per_provider_defaults_applied() { - let settings = Settings::default(); - assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); - assert_eq!(settings.usage_source(ProviderId::Codex), "auto"); - assert_eq!(settings.api_region(ProviderId::Alibaba), "singapore"); - assert_eq!(settings.api_region(ProviderId::Zai), "global"); - assert_eq!(settings.api_region(ProviderId::MiniMax), "global"); - assert!(settings.openai_web_extras(ProviderId::Codex)); - assert!(!settings.historical_tracking(ProviderId::Codex)); - assert!(!settings.avoid_keychain_prompts(ProviderId::Claude)); - assert!(!settings.auto_resume_after_quota_reset(ProviderId::Codex)); - assert!(!settings.auto_resume_after_quota_reset(ProviderId::Claude)); -} - -#[test] -fn codex_spark_usage_visibility_defaults_to_visible_and_roundtrips() { - let mut settings = Settings::default(); - assert!(settings.codex_spark_usage_visible()); - - settings.set_codex_spark_usage_visible(false); - let serialized = serde_json::to_string(&settings).unwrap(); - let loaded: Settings = serde_json::from_str(&serialized).unwrap(); - - assert!(!loaded.codex_spark_usage_visible()); -} - -#[test] -fn migrate_legacy_visibility_flags_materializes_hidden_usage_item_ids() { - let mut settings = Settings::default(); - settings.set_spark_usage_visible(ProviderId::Codex, false); - settings.claude_daily_routines_usage_visible = false; - - settings.migrate_legacy_usage_item_flags(); - - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Codex), - CODEX_SPARK_USAGE_ITEM_IDS - .iter() - .map(|id| (*id).to_string()) - .collect::>() - ); - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Claude), - vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()] - ); -} - -#[test] -fn generic_claude_visibility_writes_only_the_usage_item_list() { - let mut settings = Settings::default(); - - settings.set_hidden_usage_item_ids( - ProviderId::Claude, - vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()], - ); - - assert!(settings.claude_daily_routines_usage_visible); - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Claude), - vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()] - ); - - settings.set_hidden_usage_item_ids(ProviderId::Claude, Vec::new()); - - assert!(settings.claude_daily_routines_usage_visible); - assert!( - settings - .hidden_usage_item_ids(ProviderId::Claude) - .is_empty() - ); -} - -#[test] -fn explicit_hidden_usage_item_ids_roundtrip_and_restore_defaults() { - let mut settings = Settings::default(); - settings.set_hidden_usage_item_ids( - ProviderId::Codex, - vec![ - "metric:secondary".to_string(), - "metric:secondary".to_string(), - "not-a-metric".to_string(), - ], - ); - - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Codex), - vec!["metric:secondary".to_string()] - ); - assert!(settings.codex_spark_usage_visible()); - - let serialized = serde_json::to_string(&settings).unwrap(); - let loaded: Settings = serde_json::from_str(&serialized).unwrap(); - assert_eq!( - loaded.hidden_usage_item_ids(ProviderId::Codex), - vec!["metric:secondary".to_string()] - ); - - settings.set_hidden_usage_item_ids(ProviderId::Codex, Vec::new()); - settings.set_hidden_usage_item_ids(ProviderId::Claude, Vec::new()); - assert!(settings.hidden_usage_item_ids(ProviderId::Codex).is_empty()); - assert!(settings.codex_spark_usage_visible()); -} - -#[test] -fn legacy_visibility_setters_preserve_other_explicit_hidden_items() { - let mut settings = Settings::default(); - settings.set_hidden_usage_item_ids(ProviderId::Claude, vec!["metric:secondary".to_string()]); - - settings.toggle_hidden_items( - ProviderId::Claude, - &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], - false, - ); - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Claude), - vec![ - "metric:extra-claude-routines".to_string(), - "metric:secondary".to_string(), +/// Characterization pin for the parse errors `codexbar config validate` shows: the message and the +/// line/column must not change when the deserializer is restructured. +#[test] +fn settings_parse_errors_are_pinned() { + let cases = [ + // Wrongly typed canonical field. + "{\n \"refresh_interval_secs\": \"soon\"\n}", + // Wrongly typed legacy flat field. + "{\n \"codex_cookie_source\": 5\n}", + // Wrong type inside provider_configs. + "{\n \"provider_configs\": { \"codex\": { \"cookie_source\": 5 } }\n}", + // Canonical field error after legacy and unknown keys. + "{\n \"claude_cookie_source\": \"auto\",\n \"not_a_setting\": 1,\n \"theme\": 7\n}", + // Truncated document. + "{\n \"theme\": \"dark\",", + ]; + let errors: Vec = cases + .iter() + .map(|json| { + serde_json::from_str::(json) + .unwrap_err() + .to_string() + }) + .collect(); + assert_eq!( + errors, + [ + "invalid type: string \"soon\", expected u64 at line 2 column 33", + "invalid type: integer `5`, expected a string at line 2 column 26", + "invalid type: integer `5`, expected a string at line 2 column 53", + "expected value at line 4 column 12", + "EOF while parsing a value at line 2 column 18", ] ); - - settings.toggle_hidden_items( - ProviderId::Claude, - &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], - true, - ); - assert_eq!( - settings.hidden_usage_item_ids(ProviderId::Claude), - vec!["metric:secondary".to_string()] - ); -} - -/// Cookie-denial settings must survive the same path-based persistence used by -/// `Settings::save` and `Settings::load`, including the secure-file wrapper -/// shared by desktop and CLI settings. -#[test] -fn cookie_denial_round_trips_through_settings_persistence() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("settings.json"); - let mut settings = Settings::default(); - settings.set_cookie_source(ProviderId::Codex, "off"); - settings.set_openai_web_extras(ProviderId::Codex, false); - - settings.save_to_path(&path).unwrap(); - let loaded = Settings::load_from_path(Some(&path)); - - assert_eq!(loaded.cookie_source(ProviderId::Codex), "off"); - assert!(!loaded.openai_web_extras(ProviderId::Codex)); } diff --git a/rust/src/settings/tests/language_theme.rs b/rust/src/settings/tests/language_theme.rs new file mode 100644 index 0000000000..e0a2e09b75 --- /dev/null +++ b/rust/src/settings/tests/language_theme.rs @@ -0,0 +1,183 @@ +//! UI language and theme preference defaults, tags and round trips. + +use super::*; + +#[test] +fn test_language_defaults_to_english() { + let settings = Settings::default(); + assert_eq!(settings.ui_language, Language::English); +} + +#[test] +fn test_language_table_pins_display_names_serde_tags_and_aliases() { + // (variant, display name, serde tag, extra inputs `resolve` must accept) + let table: [(Language, &str, &str, &[&str]); 10] = [ + (Language::English, "English", "english", &[]), + (Language::Chinese, "中文", "chinese", &[]), + ( + Language::ChineseTraditional, + "繁體中文", + "chinesetraditional", + &["zh-tw", "zh-hant-tw", "繁體中文"], + ), + (Language::Japanese, "日本語", "japanese", &[]), + (Language::Korean, "한국어", "korean", &[]), + (Language::Spanish, "Español", "spanish", &[]), + ( + Language::PortugueseBrazil, + "Português (Brasil)", + "portuguesebrazil", + &[ + "pt", + "pt-BR", + "Portuguese", + "Português", + "portugues", + "Português (Brasil)", + ], + ), + ( + Language::Russian, + "Русский", + "russian", + &["ru-RU", "Русский"], + ), + ( + Language::Turkish, + "Türkçe", + "turkish", + &["tr-TR", "Türkçe", "turkce"], + ), + ( + Language::Ukrainian, + "Українська", + "ukrainian", + &["uk", "uk-UA", "Українська"], + ), + ]; + let listed: Vec = table.iter().map(|row| row.0).collect(); + assert_eq!(Language::all(), listed.as_slice()); + for (language, display_name, tag, aliases) in table { + assert_eq!(language.display_name(), display_name); + let quoted = format!("\"{tag}\""); + assert_eq!(serde_json::to_string(&language).unwrap(), quoted); + assert_eq!(serde_json::from_str::("ed).unwrap(), language); + for input in std::iter::once(tag).chain(aliases.iter().copied()) { + assert_eq!( + Language::resolve(input), + Some(language), + "failed to resolve {input}" + ); + } + } +} + +#[test] +fn test_settings_load_missing_language_field_defaults_to_english() { + // Legacy settings JSON written before the ui_language field existed. + let legacy_json = r#"{ + "enabled_providers": ["claude", "codex"], + "refresh_interval_secs": 300, + "start_minimized": false + }"#; + + let settings: Settings = serde_json::from_str(legacy_json).expect("legacy settings parse"); + assert_eq!(settings.ui_language, Language::English); +} + +#[test] +fn test_settings_roundtrip_with_language() { + let settings = Settings { + ui_language: Language::Chinese, + ..Settings::default() + }; + let (_, loaded) = round_trip(&settings); + assert_eq!(loaded.ui_language, Language::Chinese); +} + +#[test] +fn test_settings_with_utf8_bom_parses_perprovider_tray_mode() { + let json = "\u{feff}{\n \"enabled_providers\": [\"claude\", \"codex\"],\n \"refresh_interval_secs\": 300,\n \"tray_icon_mode\": \"perprovider\"\n }"; + + let settings: Settings = serde_json::from_str(json.trim_start_matches('\u{feff}')).unwrap(); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::PerProvider); +} + +#[test] +fn stacked_tray_mode_preserves_provider_preferences() { + let json = r#"{ + "tray_icon_mode": "stacked", + "stacked_tray_top_provider": "claude", + "stacked_tray_bottom_provider": "codex" + }"#; + + let settings: Settings = serde_json::from_str(json).unwrap(); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); + assert_eq!( + settings.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); + + let (_, reloaded) = round_trip(&settings); + assert_eq!( + reloaded.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + reloaded.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); +} + +#[test] +fn test_theme_defaults_to_auto() { + let settings = Settings::default(); + assert_eq!(settings.theme, ThemePreference::Auto); +} + +#[test] +fn test_theme_variants_serialize_to_stable_tags() { + let table = [ + (ThemePreference::Auto, "\"auto\""), + (ThemePreference::Light, "\"light\""), + (ThemePreference::Dark, "\"dark\""), + ]; + let listed: Vec = table.iter().map(|row| row.0).collect(); + assert_eq!(ThemePreference::all(), listed.as_slice()); + for (variant, quoted) in table { + assert_eq!(serde_json::to_string(&variant).unwrap(), quoted); + assert_eq!( + serde_json::from_str::(quoted).unwrap(), + variant + ); + } +} + +#[test] +fn test_settings_missing_theme_defaults_to_auto() { + // Legacy settings JSON without the theme field should still parse. + let legacy_json = r#"{ + "enabled_providers": ["claude", "codex"], + "refresh_interval_secs": 300, + "ui_language": "english" + }"#; + + let settings: Settings = serde_json::from_str(legacy_json).unwrap(); + assert_eq!(settings.theme, ThemePreference::Auto); +} + +#[test] +fn test_settings_roundtrip_with_theme() { + let settings = Settings { + theme: ThemePreference::Dark, + ..Settings::default() + }; + let (_, loaded) = round_trip(&settings); + assert_eq!(loaded.theme, ThemePreference::Dark); +} diff --git a/rust/src/settings/tests/provider_configs.rs b/rust/src/settings/tests/provider_configs.rs new file mode 100644 index 0000000000..6d59bd5335 --- /dev/null +++ b/rust/src/settings/tests/provider_configs.rs @@ -0,0 +1,438 @@ +//! `provider_configs` migration, canonicalization and per-provider accessors. + +use super::*; + +// ── Phase 3: provider_configs migration tests ─────────────────────── + +/// Loading a legacy `settings.json` (with flat per-provider fields) +/// must populate `provider_configs` and surface every value through the +/// per-provider accessors. +#[test] +fn test_legacy_per_provider_fields_migrate_into_provider_configs() { + // NOTE: placeholder values only — no real cookies/tokens. + let legacy_json = r#"{ + "enabled_providers": ["claude", "codex"], + "refresh_interval_secs": 300, + "codex_cookie_source": "manual", + "claude_cookie_source": "browser", + "cursor_cookie_source": "manual", + "alibaba_cookie_source": "manual", + "alibaba_cookie_header": "ali=PLACEHOLDER", + "alibaba_api_region": "cn", + "zai_api_region": "cn", + "minimax_api_region": "cn", + "minimax_api_token": "TOK_PLACEHOLDER", + "claude_usage_source": "ccusage", + "codex_usage_source": "manual", + "codex_openai_web_extras": false, + "codex_historical_tracking": true, + "claude_avoid_keychain_prompts": true, + "opencode_workspace_id": "ws_placeholder", + "jetbrains_ide_base_path": "C:/JB" + }"#; + + let settings: Settings = serde_json::from_str(legacy_json).unwrap(); + + // Cookie sources + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert_eq!(settings.cookie_source(ProviderId::Claude), "browser"); + assert_eq!(settings.cookie_source(ProviderId::Cursor), "manual"); + assert_eq!(settings.cookie_source(ProviderId::Alibaba), "manual"); + // Untouched providers fall through to the default "manual" to avoid + // background browser-cookie reads unless the user opts into Automatic. + assert_eq!(settings.cookie_source(ProviderId::Amp), "manual"); + + // Manual cookie headers + api regions + assert_eq!( + settings.manual_cookie_header(ProviderId::Alibaba), + "ali=PLACEHOLDER" + ); + assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); + assert_eq!(settings.api_region(ProviderId::Zai), "cn"); + assert_eq!(settings.api_region(ProviderId::MiniMax), "cn"); + + // Usage sources + assert_eq!(settings.usage_source(ProviderId::Claude), "ccusage"); + assert_eq!(settings.usage_source(ProviderId::Codex), "manual"); + + // Codex booleans + assert!(!settings.openai_web_extras(ProviderId::Codex)); + assert!(settings.historical_tracking(ProviderId::Codex)); + + // Claude per-provider boolean + assert!(settings.avoid_keychain_prompts(ProviderId::Claude)); + + // Misc per-provider strings + assert_eq!( + settings.workspace_id(ProviderId::OpenCode), + "ws_placeholder" + ); + assert_eq!(settings.api_token(ProviderId::MiniMax), "TOK_PLACEHOLDER"); + assert_eq!(settings.ide_base_path(ProviderId::JetBrains), "C:/JB"); + + // Legacy field-name aliases agree with typed accessors. + assert_eq!(settings.codex_cookie_source(), "manual"); + assert_eq!(settings.alibaba_api_region(), "cn"); + assert!(settings.codex_historical_tracking()); + assert!(!settings.codex_openai_web_extras()); + assert!(settings.claude_avoid_keychain_prompts()); +} + +/// Round-trip: build a `Settings` programmatically via the new map + +/// accessors, serialize, parse back, and assert equality of every +/// per-provider field. +#[test] +fn test_provider_configs_roundtrip() { + let mut settings = Settings::default(); + settings.set_cookie_source(ProviderId::Codex, "manual"); + settings.set_cookie_source(ProviderId::Claude, "browser"); + settings.set_usage_source(ProviderId::Claude, "ccusage"); + settings.set_api_region(ProviderId::Alibaba, "cn"); + settings.set_api_region(ProviderId::Zai, "cn"); + settings + .provider_config_mut(ProviderId::Amp) + .manual_cookie_header = Some("amp=PLACEHOLDER".into()); + settings.provider_config_mut(ProviderId::MiniMax).api_token = Some("TOK_PLACEHOLDER".into()); + settings.set_workspace_id(ProviderId::OpenCode, "ws_placeholder"); + settings.set_ide_base_path(ProviderId::JetBrains, "C:/JB"); + settings + .provider_config_mut(ProviderId::Codex) + .openai_web_extras = Some(false); + settings + .provider_config_mut(ProviderId::Codex) + .historical_tracking = true; + settings.set_avoid_keychain_prompts(ProviderId::Claude, true); + settings.set_auto_resume_after_quota_reset(ProviderId::Codex, true); + settings + .set_seat_credit_entitlement(ProviderId::Copilot, Some(300.0)) + .expect("valid seat credit entitlement"); + + let (json, loaded) = round_trip(&settings); + // The legacy flat fields must NOT appear in serialized output. + assert!(!json.contains("\"codex_cookie_source\""), "json: {json}"); + assert!(!json.contains("\"alibaba_api_region\""), "json: {json}"); + assert!( + !json.contains("\"claude_avoid_keychain_prompts\""), + "json: {json}" + ); + assert!(json.contains("\"provider_configs\""), "json: {json}"); + + assert_eq!(loaded.cookie_source(ProviderId::Codex), "manual"); + assert_eq!(loaded.cookie_source(ProviderId::Claude), "browser"); + assert_eq!(loaded.usage_source(ProviderId::Claude), "ccusage"); + assert_eq!(loaded.api_region(ProviderId::Alibaba), "cn"); + assert_eq!(loaded.api_region(ProviderId::Zai), "cn"); + assert_eq!( + loaded.manual_cookie_header(ProviderId::Amp), + "amp=PLACEHOLDER" + ); + assert_eq!(loaded.api_token(ProviderId::MiniMax), "TOK_PLACEHOLDER"); + assert_eq!(loaded.workspace_id(ProviderId::OpenCode), "ws_placeholder"); + assert_eq!(loaded.ide_base_path(ProviderId::JetBrains), "C:/JB"); + assert!(!loaded.openai_web_extras(ProviderId::Codex)); + assert!(loaded.historical_tracking(ProviderId::Codex)); + assert!(loaded.avoid_keychain_prompts(ProviderId::Claude)); + assert!(loaded.auto_resume_after_quota_reset(ProviderId::Codex)); + assert_eq!( + loaded.seat_credit_entitlement(ProviderId::Copilot), + Some(300.0) + ); + assert_eq!( + loaded.provider_configs.get(&ProviderId::Codex), + settings.provider_configs.get(&ProviderId::Codex) + ); +} + +/// New-format files (no legacy flat fields, only `provider_configs`) +/// must load identically. +#[test] +fn test_new_format_provider_configs_only() { + let json = r#"{ + "enabled_providers": ["claude"], + "refresh_interval_secs": 300, + "provider_configs": { + "codex": { "cookie_source": "manual", "openai_web_extras": false }, + "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } + } + }"#; + + let settings: Settings = serde_json::from_str(json).unwrap(); + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert!(!settings.openai_web_extras(ProviderId::Codex)); + assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); + assert_eq!( + settings.manual_cookie_header(ProviderId::Alibaba), + "ali=PLACEHOLDER" + ); + // Untouched providers still get their defaults. + assert_eq!(settings.cookie_source(ProviderId::Claude), "manual"); + assert_eq!(settings.api_region(ProviderId::Zai), "global"); +} + +#[test] +fn retired_provider_config_is_ignored_until_explicit_save() { + let original = r#"{ + "enabled_providers": ["codex", "crof"], + "refresh_interval_secs": 300, + "provider_metrics": { "codex": "weekly", "crof": "session" }, + "float_bar_provider_ids": ["codex", "crof"], + "stacked_tray_top_provider": "crof", + "stacked_tray_bottom_provider": "crof", + "provider_configs": { + "crof": { "api_token": "retired-fixture-key" }, + "codex": { "cookie_source": "manual", "openai_web_extras": false }, + "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } + } + }"#; + let original_bytes = original.as_bytes().to_vec(); + + let settings: Settings = + serde_json::from_str(original).expect("load settings with retired key"); + + assert_eq!(original.as_bytes(), original_bytes); + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert!(!settings.openai_web_extras(ProviderId::Codex)); + assert_eq!( + settings.enabled_providers, + HashSet::from(["codex".to_string()]) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex"]); + assert_eq!(settings.stacked_tray_top_provider, None); + assert_eq!(settings.stacked_tray_bottom_provider, None); + assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); + assert_eq!( + settings.manual_cookie_header(ProviderId::Alibaba), + "ali=PLACEHOLDER" + ); + + let saved = serde_json::to_string(&settings).expect("serialize sanitized settings"); + let saved_value: serde_json::Value = serde_json::from_str(&saved).unwrap(); + let saved_configs = saved_value["provider_configs"].as_object().unwrap(); + assert!(!saved_configs.contains_key("crof")); + assert!(saved_configs.contains_key("codex")); + assert!(saved_configs.contains_key("alibaba")); + assert!( + !saved.contains("\"crof\""), + "saved settings retained Crof: {saved}" + ); +} + +#[test] +fn provider_aliases_are_canonicalized_at_the_load_boundary() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["openai", "ClAuDe", "not-a-provider"], + "provider_metrics": { + "openai": "weekly", + "CoDeX": "session", + "not-a-provider": "weekly" + }, + "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "ClAuDe" + }"#, + ) + .expect("load settings containing provider aliases"); + + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string(), "codex".to_string()]) + ); + assert_eq!( + settings.provider_metrics.get("codex"), + Some(&MetricPreference::Session) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex", "claude"]); + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("claude") + ); +} + +#[test] +fn stacked_preferences_preserve_known_disabled_providers() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["claude"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "not-a-provider" + }"#, + ) + .expect("load stacked preferences independently of enablement"); + + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!(settings.stacked_tray_bottom_provider, None); + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string()]) + ); +} + +/// Default `Settings` should serialize WITHOUT a `provider_configs` +/// field (empty map skipped). +#[test] +fn test_default_settings_skip_empty_provider_configs() { + let settings = Settings::default(); + let json = serde_json::to_string(&settings).unwrap(); + assert!( + !json.contains("\"provider_configs\""), + "empty map should be skipped, json: {json}" + ); +} + +/// Per-provider defaults are applied even when the entry is absent. +#[test] +fn test_per_provider_defaults_applied() { + let settings = Settings::default(); + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert_eq!(settings.usage_source(ProviderId::Codex), "auto"); + assert_eq!(settings.api_region(ProviderId::Alibaba), "singapore"); + assert_eq!(settings.api_region(ProviderId::Zai), "global"); + assert_eq!(settings.api_region(ProviderId::MiniMax), "global"); + assert!(settings.openai_web_extras(ProviderId::Codex)); + assert!(!settings.historical_tracking(ProviderId::Codex)); + assert!(!settings.avoid_keychain_prompts(ProviderId::Claude)); + assert!(!settings.auto_resume_after_quota_reset(ProviderId::Codex)); + assert!(!settings.auto_resume_after_quota_reset(ProviderId::Claude)); +} + +#[test] +fn codex_spark_usage_visibility_defaults_to_visible_and_roundtrips() { + let mut settings = Settings::default(); + assert!(settings.codex_spark_usage_visible()); + + settings.set_codex_spark_usage_visible(false); + let serialized = serde_json::to_string(&settings).unwrap(); + let loaded: Settings = serde_json::from_str(&serialized).unwrap(); + + assert!(!loaded.codex_spark_usage_visible()); +} + +#[test] +fn migrate_legacy_visibility_flags_materializes_hidden_usage_item_ids() { + let mut settings = Settings::default(); + settings.set_spark_usage_visible(ProviderId::Codex, false); + settings.claude_daily_routines_usage_visible = false; + + settings.migrate_legacy_usage_item_flags(); + + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Codex), + CODEX_SPARK_USAGE_ITEM_IDS + .iter() + .map(|id| (*id).to_string()) + .collect::>() + ); + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Claude), + vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()] + ); +} + +#[test] +fn generic_claude_visibility_writes_only_the_usage_item_list() { + let mut settings = Settings::default(); + + settings.set_hidden_usage_item_ids( + ProviderId::Claude, + vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()], + ); + + assert!(settings.claude_daily_routines_usage_visible); + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Claude), + vec![CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID.to_string()] + ); + + settings.set_hidden_usage_item_ids(ProviderId::Claude, Vec::new()); + + assert!(settings.claude_daily_routines_usage_visible); + assert!( + settings + .hidden_usage_item_ids(ProviderId::Claude) + .is_empty() + ); +} + +#[test] +fn explicit_hidden_usage_item_ids_roundtrip_and_restore_defaults() { + let mut settings = Settings::default(); + settings.set_hidden_usage_item_ids( + ProviderId::Codex, + vec![ + "metric:secondary".to_string(), + "metric:secondary".to_string(), + "not-a-metric".to_string(), + ], + ); + + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Codex), + vec!["metric:secondary".to_string()] + ); + assert!(settings.codex_spark_usage_visible()); + + let serialized = serde_json::to_string(&settings).unwrap(); + let loaded: Settings = serde_json::from_str(&serialized).unwrap(); + assert_eq!( + loaded.hidden_usage_item_ids(ProviderId::Codex), + vec!["metric:secondary".to_string()] + ); + + settings.set_hidden_usage_item_ids(ProviderId::Codex, Vec::new()); + settings.set_hidden_usage_item_ids(ProviderId::Claude, Vec::new()); + assert!(settings.hidden_usage_item_ids(ProviderId::Codex).is_empty()); + assert!(settings.codex_spark_usage_visible()); +} + +#[test] +fn legacy_visibility_setters_preserve_other_explicit_hidden_items() { + let mut settings = Settings::default(); + settings.set_hidden_usage_item_ids(ProviderId::Claude, vec!["metric:secondary".to_string()]); + + settings.toggle_hidden_items( + ProviderId::Claude, + &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], + false, + ); + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Claude), + vec![ + "metric:extra-claude-routines".to_string(), + "metric:secondary".to_string(), + ] + ); + + settings.toggle_hidden_items( + ProviderId::Claude, + &[CLAUDE_DAILY_ROUTINES_USAGE_ITEM_ID], + true, + ); + assert_eq!( + settings.hidden_usage_item_ids(ProviderId::Claude), + vec!["metric:secondary".to_string()] + ); +} + +/// Cookie-denial settings must survive the same path-based persistence used by +/// `Settings::save` and `Settings::load`, including the secure-file wrapper +/// shared by desktop and CLI settings. +#[test] +fn cookie_denial_round_trips_through_settings_persistence() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("settings.json"); + let mut settings = Settings::default(); + settings.set_cookie_source(ProviderId::Codex, "off"); + settings + .provider_config_mut(ProviderId::Codex) + .openai_web_extras = Some(false); + + settings.save_to_path(&path).unwrap(); + let loaded = Settings::load_from_path(Some(&path)); + + assert_eq!(loaded.cookie_source(ProviderId::Codex), "off"); + assert!(!loaded.openai_web_extras(ProviderId::Codex)); +} diff --git a/rust/src/shortcuts.rs b/rust/src/shortcuts.rs deleted file mode 100755 index d51d35db29..0000000000 --- a/rust/src/shortcuts.rs +++ /dev/null @@ -1,276 +0,0 @@ -//! Global keyboard shortcuts for CodexBar -//! -//! Provides system-wide hotkeys to open the menu - -#![allow(dead_code, reason = "keyboard shortcut types reserved for future global hotkey integration")] - -use global_hotkey::{ - GlobalHotKeyEvent, GlobalHotKeyManager, - hotkey::{Code, HotKey, Modifiers}, -}; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; - -const KEY_ALIASES: &[(&str, Code)] = &[ - ("a", Code::KeyA), - ("b", Code::KeyB), - ("c", Code::KeyC), - ("d", Code::KeyD), - ("e", Code::KeyE), - ("f", Code::KeyF), - ("g", Code::KeyG), - ("h", Code::KeyH), - ("i", Code::KeyI), - ("j", Code::KeyJ), - ("k", Code::KeyK), - ("l", Code::KeyL), - ("m", Code::KeyM), - ("n", Code::KeyN), - ("o", Code::KeyO), - ("p", Code::KeyP), - ("q", Code::KeyQ), - ("r", Code::KeyR), - ("s", Code::KeyS), - ("t", Code::KeyT), - ("u", Code::KeyU), - ("v", Code::KeyV), - ("w", Code::KeyW), - ("x", Code::KeyX), - ("y", Code::KeyY), - ("z", Code::KeyZ), - ("0", Code::Digit0), - ("1", Code::Digit1), - ("2", Code::Digit2), - ("3", Code::Digit3), - ("4", Code::Digit4), - ("5", Code::Digit5), - ("6", Code::Digit6), - ("7", Code::Digit7), - ("8", Code::Digit8), - ("9", Code::Digit9), - ("f1", Code::F1), - ("f2", Code::F2), - ("f3", Code::F3), - ("f4", Code::F4), - ("f5", Code::F5), - ("f6", Code::F6), - ("f7", Code::F7), - ("f8", Code::F8), - ("f9", Code::F9), - ("f10", Code::F10), - ("f11", Code::F11), - ("f12", Code::F12), - ("space", Code::Space), - ("enter", Code::Enter), - ("return", Code::Enter), - ("escape", Code::Escape), - ("esc", Code::Escape), - ("tab", Code::Tab), -]; - -const KEY_LABELS: &[(Code, &str)] = &[ - (Code::KeyA, "A"), - (Code::KeyB, "B"), - (Code::KeyC, "C"), - (Code::KeyD, "D"), - (Code::KeyE, "E"), - (Code::KeyF, "F"), - (Code::KeyG, "G"), - (Code::KeyH, "H"), - (Code::KeyI, "I"), - (Code::KeyJ, "J"), - (Code::KeyK, "K"), - (Code::KeyL, "L"), - (Code::KeyM, "M"), - (Code::KeyN, "N"), - (Code::KeyO, "O"), - (Code::KeyP, "P"), - (Code::KeyQ, "Q"), - (Code::KeyR, "R"), - (Code::KeyS, "S"), - (Code::KeyT, "T"), - (Code::KeyU, "U"), - (Code::KeyV, "V"), - (Code::KeyW, "W"), - (Code::KeyX, "X"), - (Code::KeyY, "Y"), - (Code::KeyZ, "Z"), - (Code::Digit0, "0"), - (Code::Digit1, "1"), - (Code::Digit2, "2"), - (Code::Digit3, "3"), - (Code::Digit4, "4"), - (Code::Digit5, "5"), - (Code::Digit6, "6"), - (Code::Digit7, "7"), - (Code::Digit8, "8"), - (Code::Digit9, "9"), - (Code::F1, "F1"), - (Code::F2, "F2"), - (Code::F3, "F3"), - (Code::F4, "F4"), - (Code::F5, "F5"), - (Code::F6, "F6"), - (Code::F7, "F7"), - (Code::F8, "F8"), - (Code::F9, "F9"), - (Code::F10, "F10"), - (Code::F11, "F11"), - (Code::F12, "F12"), - (Code::Space, "Space"), - (Code::Enter, "Enter"), - (Code::Escape, "Esc"), - (Code::Tab, "Tab"), -]; - -/// Keyboard shortcut manager -pub struct ShortcutManager { - manager: GlobalHotKeyManager, - open_menu_id: u32, - triggered: Arc, -} - -impl ShortcutManager { - /// Create a new shortcut manager with default shortcuts - pub fn new() -> anyhow::Result { - let manager = GlobalHotKeyManager::new()?; - - // Default shortcut: Ctrl+Shift+U (U for Usage) - let open_menu = HotKey::new(Some(Modifiers::CONTROL | Modifiers::SHIFT), Code::KeyU); - let open_menu_id = open_menu.id(); - - manager.register(open_menu)?; - - Ok(Self { - manager, - open_menu_id, - triggered: Arc::new(AtomicBool::new(false)), - }) - } - - /// Register a custom shortcut for opening the menu - pub fn set_open_menu_shortcut( - &mut self, - modifiers: Modifiers, - key: Code, - ) -> anyhow::Result<()> { - // Unregister old shortcut - let old = HotKey::new(Some(Modifiers::CONTROL | Modifiers::SHIFT), Code::KeyU); - let _ = self.manager.unregister(old); - - // Register new shortcut - let new_hotkey = HotKey::new(Some(modifiers), key); - self.open_menu_id = new_hotkey.id(); - self.manager.register(new_hotkey)?; - - Ok(()) - } - - /// Check if the open menu shortcut was triggered - /// Call this in your event loop - pub fn check_events(&self) -> bool { - if let Ok(event) = GlobalHotKeyEvent::receiver().try_recv() - && event.id == self.open_menu_id - { - return true; - } - false - } - - /// Get the triggered flag (for async usage) - pub fn triggered_flag(&self) -> Arc { - Arc::clone(&self.triggered) - } -} - -impl Drop for ShortcutManager { - fn drop(&mut self) { - // Unregister all shortcuts on drop - let hotkey = HotKey::new(Some(Modifiers::CONTROL | Modifiers::SHIFT), Code::KeyU); - let _ = self.manager.unregister(hotkey); - } -} - -/// Parse a shortcut string like "Ctrl+Shift+U" into modifiers and key -pub fn parse_shortcut(s: &str) -> Option<(Modifiers, Code)> { - let mut modifiers = Modifiers::empty(); - let mut key_code = None; - - for part in s.split('+').map(str::trim).filter(|part| !part.is_empty()) { - if let Some(parsed_modifier) = parse_modifier(part) { - modifiers |= parsed_modifier; - } else if let Some(parsed_key) = parse_key(part) { - key_code = Some(parsed_key); - } - } - - key_code.map(|k| (modifiers, k)) -} - -fn parse_modifier(token: &str) -> Option { - match token.to_ascii_lowercase().as_str() { - "ctrl" | "control" => Some(Modifiers::CONTROL), - "shift" => Some(Modifiers::SHIFT), - "alt" => Some(Modifiers::ALT), - "super" | "win" | "meta" => Some(Modifiers::SUPER), - _ => None, - } -} - -fn parse_key(token: &str) -> Option { - let normalized = token.to_ascii_lowercase(); - KEY_ALIASES - .iter() - .find_map(|(alias, code)| (*alias == normalized).then_some(*code)) -} - -/// Format a shortcut for display -pub fn format_shortcut(modifiers: Modifiers, key: Code) -> String { - let mut parts = Vec::new(); - - if modifiers.contains(Modifiers::CONTROL) { - parts.push("Ctrl"); - } - if modifiers.contains(Modifiers::SHIFT) { - parts.push("Shift"); - } - if modifiers.contains(Modifiers::ALT) { - parts.push("Alt"); - } - if modifiers.contains(Modifiers::SUPER) { - parts.push("Win"); - } - - parts.push(format_key(key)); - parts.join("+") -} - -fn format_key(key: Code) -> &'static str { - KEY_LABELS - .iter() - .find_map(|(code, label)| (*code == key).then_some(*label)) - .unwrap_or("?") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_parse_shortcut() { - let (mods, key) = parse_shortcut("Ctrl+Shift+U").unwrap(); - assert!(mods.contains(Modifiers::CONTROL)); - assert!(mods.contains(Modifiers::SHIFT)); - assert_eq!(key, Code::KeyU); - - let (mods, key) = parse_shortcut("Alt+F1").unwrap(); - assert!(mods.contains(Modifiers::ALT)); - assert_eq!(key, Code::F1); - } - - #[test] - fn test_format_shortcut() { - let s = format_shortcut(Modifiers::CONTROL | Modifiers::SHIFT, Code::KeyU); - assert_eq!(s, "Ctrl+Shift+U"); - } -} diff --git a/rust/src/sound.rs b/rust/src/sound.rs index d2034cc57a..a13871da13 100755 --- a/rust/src/sound.rs +++ b/rust/src/sound.rs @@ -2,11 +2,6 @@ //! //! Handles per-event custom WAV files, built-in CodexBar sounds, and Windows system sounds. -#![allow( - dead_code, - reason = "sound playback types reserved for future alert audio integration" -)] - use crate::settings::{NotificationSoundPaths, NotificationSoundTheme, Settings}; use serde::{Deserialize, Serialize}; use std::path::Path; @@ -14,7 +9,6 @@ use std::path::Path; #[cfg(target_os = "windows")] use std::sync::{OnceLock, mpsc}; -const MINIMUM_WAV_SIZE: usize = 44; const TAG_SIZE: usize = 4; const RIFF_TAG_OFFSET: usize = 0; const RIFF_SIZE_OFFSET: usize = 4; @@ -415,6 +409,9 @@ fn playback_worker(receiver: mpsc::Receiver) { mod tests { use super::*; + /// A canonical PCM WAV header is 44 bytes; generated sounds must carry samples past it. + const MINIMUM_WAV_SIZE: usize = 44; + fn find_chunk<'a>(wav: &'a [u8], tag: &[u8; TAG_SIZE]) -> Option<&'a [u8]> { let mut offset = RIFF_HEADER_SIZE; while offset + CHUNK_HEADER_SIZE <= wav.len() { diff --git a/rust/src/status/indicators.rs b/rust/src/status/indicators.rs deleted file mode 100755 index 8c36833dc1..0000000000 --- a/rust/src/status/indicators.rs +++ /dev/null @@ -1,334 +0,0 @@ -//! Provider Status Indicators -//! -//! Status overlay system for showing provider health in tray icons. -//! Integrates with status page monitoring (Statuspage.io, etc.) - -#![allow(dead_code, reason = "status indicator types reserved for future UI integration")] - -use serde::{Deserialize, Serialize}; - -/// Provider status level -#[derive( - Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default, -)] -pub enum StatusLevel { - /// All systems operational - #[default] - Operational, - /// Minor issues (small yellow dot) - Minor, - /// Major issues (orange dot) - Major, - /// Critical outage (red dot) - Critical, - /// Scheduled maintenance (blue dot) - Maintenance, - /// Status unknown or fetch failed (gray dot) - Unknown, -} - -impl StatusLevel { - /// Get the RGB color for this status level - pub fn color(&self) -> (u8, u8, u8) { - match self { - StatusLevel::Operational => (0, 255, 0), // Green - StatusLevel::Minor => (255, 255, 0), // Yellow - StatusLevel::Major => (255, 165, 0), // Orange - StatusLevel::Critical => (255, 0, 0), // Red - StatusLevel::Maintenance => (100, 149, 237), // Cornflower blue - StatusLevel::Unknown => (128, 128, 128), // Gray - } - } - - /// Get the RGBA color with alpha - pub fn color_rgba(&self) -> (u8, u8, u8, u8) { - let (r, g, b) = self.color(); - (r, g, b, 255) - } - - /// Whether this status should show an overlay - pub fn should_show_overlay(&self) -> bool { - !matches!(self, StatusLevel::Operational) - } - - /// Get a short label for the status - pub fn label(&self) -> &'static str { - match self { - StatusLevel::Operational => "Operational", - StatusLevel::Minor => "Minor Issues", - StatusLevel::Major => "Major Issues", - StatusLevel::Critical => "Critical Outage", - StatusLevel::Maintenance => "Maintenance", - StatusLevel::Unknown => "Unknown", - } - } - - /// Get an emoji for the status - pub fn emoji(&self) -> &'static str { - match self { - StatusLevel::Operational => "✓", - StatusLevel::Minor => "⚠", - StatusLevel::Major => "⚠", - StatusLevel::Critical => "✕", - StatusLevel::Maintenance => "🔧", - StatusLevel::Unknown => "?", - } - } - - /// Get a status prefix for menu items (colored Unicode dots) - /// Returns empty string for Operational status (no dot needed) - pub fn status_prefix(&self) -> &'static str { - match self { - StatusLevel::Operational => "", - StatusLevel::Minor => "\u{1F7E1} ", // Yellow circle - StatusLevel::Major => "\u{1F7E0} ", // Orange circle - StatusLevel::Critical => "\u{1F534} ", // Red circle - StatusLevel::Maintenance => "\u{1F535} ", // Blue circle - StatusLevel::Unknown => "", // No dot for unknown - } - } - - /// Get an ASCII-safe status prefix for menu items (fallback for systems without Unicode support) - pub fn status_prefix_ascii(&self) -> &'static str { - match self { - StatusLevel::Operational => "", - StatusLevel::Minor => "[!] ", - StatusLevel::Major => "[!!] ", - StatusLevel::Critical => "[X] ", - StatusLevel::Maintenance => "[M] ", - StatusLevel::Unknown => "", - } - } - - /// Parse from Statuspage.io status string - pub fn from_statuspage(status: &str) -> Self { - match status.to_lowercase().as_str() { - "operational" | "none" => StatusLevel::Operational, - "degraded_performance" | "degraded" | "minor" => StatusLevel::Minor, - "partial_outage" | "partial" | "major" => StatusLevel::Major, - "major_outage" | "critical" | "outage" => StatusLevel::Critical, - "under_maintenance" | "maintenance" | "scheduled" => StatusLevel::Maintenance, - _ => StatusLevel::Unknown, - } - } -} - -/// Provider status information -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ProviderStatus { - /// Current status level - pub level: StatusLevel, - /// Human-readable description - pub description: Option, - /// URL to status page - pub status_url: Option, - /// When this status was last updated - pub updated_at: chrono::DateTime, - /// Active incidents count - pub incident_count: usize, -} - -impl Default for ProviderStatus { - fn default() -> Self { - Self { - level: StatusLevel::Unknown, - description: None, - status_url: None, - updated_at: chrono::Utc::now(), - incident_count: 0, - } - } -} - -impl ProviderStatus { - pub fn operational() -> Self { - Self { - level: StatusLevel::Operational, - description: Some("All systems operational".to_string()), - ..Default::default() - } - } - - pub fn with_level(mut self, level: StatusLevel) -> Self { - self.level = level; - self - } - - pub fn with_description(mut self, description: impl Into) -> Self { - self.description = Some(description.into()); - self - } - - pub fn with_url(mut self, url: impl Into) -> Self { - self.status_url = Some(url.into()); - self - } -} - -/// Status overlay position on the icon -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] -pub enum OverlayPosition { - /// Top-right corner (default) - #[default] - TopRight, - /// Top-left corner - TopLeft, - /// Bottom-right corner - BottomRight, - /// Bottom-left corner - BottomLeft, -} - -impl OverlayPosition { - /// Get (x, y) offset from icon origin for a given icon size - pub fn offset(&self, icon_size: u32, dot_size: u32) -> (i32, i32) { - let padding = 1; - let icon_size = icon_size as i32; - let dot_size = dot_size as i32; - - match self { - OverlayPosition::TopRight => (icon_size - dot_size - padding, padding), - OverlayPosition::TopLeft => (padding, padding), - OverlayPosition::BottomRight => ( - icon_size - dot_size - padding, - icon_size - dot_size - padding, - ), - OverlayPosition::BottomLeft => (padding, icon_size - dot_size - padding), - } - } -} - -/// Status overlay configuration -#[derive(Debug, Clone)] -pub struct StatusOverlayConfig { - /// Position of the overlay dot - pub position: OverlayPosition, - /// Size of the overlay dot in pixels - pub dot_size: u32, - /// Whether to pulse/animate the dot for critical status - pub animate_critical: bool, -} - -impl Default for StatusOverlayConfig { - fn default() -> Self { - Self { - position: OverlayPosition::TopRight, - dot_size: 4, - animate_critical: true, - } - } -} - -/// Statuspage.io API response structures -#[derive(Debug, Deserialize)] -pub struct StatuspageResponse { - pub status: StatuspageStatus, - #[serde(default)] - pub incidents: Vec, -} - -#[derive(Debug, Deserialize)] -pub struct StatuspageStatus { - pub indicator: String, - pub description: String, -} - -#[derive(Debug, Deserialize)] -pub struct StatuspageIncident { - pub id: String, - pub name: String, - pub status: String, - pub impact: String, -} - -impl StatuspageResponse { - /// Convert to ProviderStatus - pub fn to_provider_status(&self, status_url: Option<&str>) -> ProviderStatus { - let level = StatusLevel::from_statuspage(&self.status.indicator); - ProviderStatus { - level, - description: Some(self.status.description.clone()), - status_url: status_url.map(|s| s.to_string()), - updated_at: chrono::Utc::now(), - incident_count: self.incidents.len(), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_status_level_ordering() { - assert!(StatusLevel::Minor > StatusLevel::Operational); - assert!(StatusLevel::Major > StatusLevel::Minor); - assert!(StatusLevel::Critical > StatusLevel::Major); - } - - #[test] - fn test_status_level_colors() { - let (r, g, b) = StatusLevel::Critical.color(); - assert_eq!(r, 255); - assert_eq!(g, 0); - assert_eq!(b, 0); - } - - #[test] - fn test_from_statuspage() { - assert_eq!( - StatusLevel::from_statuspage("operational"), - StatusLevel::Operational - ); - assert_eq!( - StatusLevel::from_statuspage("degraded_performance"), - StatusLevel::Minor - ); - assert_eq!( - StatusLevel::from_statuspage("major_outage"), - StatusLevel::Critical - ); - assert_eq!( - StatusLevel::from_statuspage("under_maintenance"), - StatusLevel::Maintenance - ); - } - - #[test] - fn test_overlay_position() { - let (x, y) = OverlayPosition::TopRight.offset(16, 4); - assert_eq!(x, 11); // 16 - 4 - 1 - assert_eq!(y, 1); - } - - #[test] - fn test_should_show_overlay() { - assert!(!StatusLevel::Operational.should_show_overlay()); - assert!(StatusLevel::Minor.should_show_overlay()); - assert!(StatusLevel::Critical.should_show_overlay()); - } - - #[test] - fn test_status_prefix() { - // Operational should have no prefix - assert_eq!(StatusLevel::Operational.status_prefix(), ""); - // Non-operational statuses should have colored dot prefixes - assert!(!StatusLevel::Minor.status_prefix().is_empty()); - assert!(!StatusLevel::Major.status_prefix().is_empty()); - assert!(!StatusLevel::Critical.status_prefix().is_empty()); - assert!(!StatusLevel::Maintenance.status_prefix().is_empty()); - // Unknown should have no prefix - assert_eq!(StatusLevel::Unknown.status_prefix(), ""); - } - - #[test] - fn test_status_prefix_ascii() { - // Operational should have no prefix - assert_eq!(StatusLevel::Operational.status_prefix_ascii(), ""); - // Non-operational statuses should have ASCII prefixes - assert_eq!(StatusLevel::Minor.status_prefix_ascii(), "[!] "); - assert_eq!(StatusLevel::Major.status_prefix_ascii(), "[!!] "); - assert_eq!(StatusLevel::Critical.status_prefix_ascii(), "[X] "); - assert_eq!(StatusLevel::Maintenance.status_prefix_ascii(), "[M] "); - } -} diff --git a/rust/src/status/mod.rs b/rust/src/status/mod.rs index 8cf2cbbd7b..cda2674254 100755 --- a/rust/src/status/mod.rs +++ b/rust/src/status/mod.rs @@ -2,7 +2,6 @@ //! //! Fetches operational status from provider status pages -#![allow(dead_code, reason = "status types reserved for future UI integration")] use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -232,21 +231,3 @@ fn last_cached_success(provider: &str) -> Option { .ok() .and_then(|cache| cache.get(provider).cloned()) } - -/// Fetch status for all providers in parallel -pub async fn fetch_all_statuses(providers: &[&str]) -> HashMap { - let futures: Vec<_> = providers - .iter() - .map(|&p| async move { - let status = fetch_provider_status(p).await; - (p.to_string(), status) - }) - .collect(); - - let results = futures::future::join_all(futures).await; - - results - .into_iter() - .filter_map(|(provider, status)| status.map(|s| (provider, s))) - .collect() -} diff --git a/rust/src/tray/icon.rs b/rust/src/tray/icon.rs index 459d9bb1b9..75dc01f24d 100755 --- a/rust/src/tray/icon.rs +++ b/rust/src/tray/icon.rs @@ -14,10 +14,6 @@ pub enum UsageLevel { /// 95-100% used - red Critical, /// Unknown/error state - gray - #[allow( - dead_code, - reason = "tray icon types reserved for future system tray integration" - )] Unknown, } @@ -43,21 +39,6 @@ impl UsageLevel { } } -/// Badge type for status indicators -#[allow( - dead_code, - reason = "tray icon types reserved for future system tray integration" -)] -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum BadgeType { - /// Warning indicator (yellow) - Warning, - /// Error/incident indicator (red) - Incident, - /// No badge - None, -} - /// Loading animation patterns for tray icon #[derive(Debug, Clone, Copy, PartialEq, Default)] pub enum LoadingPattern { @@ -121,18 +102,6 @@ impl LoadingPattern { } } - /// Get secondary bar offset (to make it animate differently) - pub fn secondary_offset(&self) -> f64 { - match self { - LoadingPattern::KnightRider => 0.25, - LoadingPattern::Cylon => 0.15, - LoadingPattern::OutsideIn => 0.5, - LoadingPattern::Race => 0.2, - LoadingPattern::Pulse => 0.3, - LoadingPattern::Unbraid => 0.1, - } - } - /// Get all available patterns pub fn all() -> &'static [LoadingPattern] { &[ @@ -183,12 +152,4 @@ mod tests { let (r, g, b) = UsageLevel::Critical.color(); assert!(r > g && r > b); // Red should be dominant for critical } - - #[test] - fn test_badge_type_equality() { - assert_eq!(BadgeType::None, BadgeType::None); - assert_eq!(BadgeType::Warning, BadgeType::Warning); - assert_eq!(BadgeType::Incident, BadgeType::Incident); - assert_ne!(BadgeType::None, BadgeType::Warning); - } } diff --git a/rust/src/updater.rs b/rust/src/updater.rs index a5b740aba3..0d4460e993 100755 --- a/rust/src/updater.rs +++ b/rust/src/updater.rs @@ -4,7 +4,6 @@ use crate::settings::UpdateChannel; use serde::Deserialize; use std::path::{Path, PathBuf}; -use std::sync::Arc; use tokio::sync::watch; const GITHUB_REPO: &str = "nesszer/Win-CodexBar"; @@ -37,15 +36,7 @@ pub struct UpdateInfo { pub version: String, pub download_url: String, pub expected_sha256: Option, - #[allow( - dead_code, - reason = "update metadata fields are deserialized for version comparison but not all are read" - )] pub release_url: String, - #[allow( - dead_code, - reason = "update metadata fields are deserialized for version comparison but not all are read" - )] pub release_notes: String, pub delivery: UpdateDelivery, } @@ -68,12 +59,6 @@ struct GitHubRelease { assets: Vec, #[serde(default)] draft: bool, - #[serde(default)] - #[allow( - dead_code, - reason = "update metadata fields are deserialized for version comparison but not all are read" - )] - prerelease: bool, } #[derive(Debug, Deserialize)] @@ -88,10 +73,6 @@ struct GitHubAsset { /// /// When `channel` is `UpdateChannel::Beta`, includes pre-release versions. /// When `channel` is `UpdateChannel::Stable`, only considers stable releases. -#[allow( - dead_code, - reason = "update check response fields are deserialized for parsing but not all are read" -)] pub async fn check_for_updates() -> Option { check_for_updates_with_channel(UpdateChannel::Stable).await } @@ -202,31 +183,6 @@ fn parse_version_triplet(v: &str) -> (u32, u32, u32) { ) } -fn installer_version_from_name(name: &str) -> Option<(u32, u32, u32)> { - let lower = name.to_ascii_lowercase(); - let stem = lower - .strip_suffix("-setup.exe") - .or_else(|| lower.strip_suffix(".msi"))?; - - let version_candidate = stem.split_once('-').map(|(_, rest)| rest).unwrap_or(stem); - let version_text: String = version_candidate - .chars() - .skip_while(|ch| !ch.is_ascii_digit()) - .take_while(|ch| ch.is_ascii_digit() || *ch == '.') - .collect(); - - if version_text.is_empty() { - return None; - } - - let version = parse_version_triplet(&version_text); - if version == (0, 0, 0) { - return None; - } - - Some(version) -} - fn parse_sha256_digest(digest: &str) -> Option<&str> { let (algo, hex) = digest.split_once(':')?; if !algo.eq_ignore_ascii_case("sha256") { @@ -249,15 +205,6 @@ fn is_newer_version(remote: &str, current: &str) -> bool { remote_v > current_v } -/// Get the current version -#[allow( - dead_code, - reason = "update info struct reserved for future UI integration" -)] -pub fn current_version() -> &'static str { - CURRENT_VERSION -} - /// Get the download directory for updates fn get_download_dir() -> Option { dirs::cache_dir().map(|p| p.join("CodexBar").join("updates")) @@ -441,40 +388,6 @@ fn sha256_file(file_path: &Path) -> Result { Ok(format!("{:x}", hasher.finalize())) } -/// Start background download of an update -/// -/// Returns a receiver that can be polled for progress updates. -#[allow( - dead_code, - reason = "update info struct reserved for future UI integration" -)] -pub fn start_background_download( - update_info: UpdateInfo, -) -> ( - Arc>, - std::thread::JoinHandle<()>, -) { - let (tx, rx) = watch::channel(UpdateState::Available); - let rx = Arc::new(rx); - - let handle = std::thread::spawn(move || { - let rt = tokio::runtime::Runtime::new().unwrap(); - rt.block_on(async { - match download_update(&update_info, tx.clone()).await { - Ok(_path) => { - // UpdateState::Ready is already sent by download_update - } - Err(e) => { - // Best-effort failure signal; the receiver may already be gone. - let _failure_signal = tx.send(UpdateState::Failed(e)); - } - } - }); - }); - - (rx, handle) -} - /// Apply a downloaded update by spawning the installer and exiting /// /// This function will: @@ -646,62 +559,6 @@ fn windows_powershell_path() -> PathBuf { .unwrap_or_else(|| PathBuf::from("powershell.exe")) } -/// Check if there's a pending update ready to install -#[allow( - dead_code, - reason = "update info struct reserved for future UI integration" -)] -pub fn get_pending_update() -> Option { - let download_dir = get_download_dir()?; - - if !download_dir.exists() { - return None; - } - - find_pending_installer_in_dir(&download_dir) -} - -fn find_pending_installer_in_dir(download_dir: &Path) -> Option { - let current_version = parse_version_triplet(CURRENT_VERSION); - - // Only treat newer installer assets as pending updates, and prefer the highest - // installer version when multiple cached installers are present. - std::fs::read_dir(download_dir) - .ok()? - .filter_map(|entry| entry.ok()) - .filter_map(|entry| { - let path = entry.path(); - let file_name = path.file_name()?.to_str()?; - let installer_version = installer_version_from_name(file_name)?; - if installer_version <= current_version { - return None; - } - - let modified = entry - .metadata() - .ok() - .and_then(|meta| meta.modified().ok()) - .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok()) - .map(|duration| duration.as_secs()) - .unwrap_or(0); - - Some(((installer_version, modified), path)) - }) - .max_by_key(|(sort_key, _)| *sort_key) - .map(|(_, path)| path) -} - -/// Clean up downloaded updates -#[allow( - dead_code, - reason = "update info struct reserved for future UI integration" -)] -pub fn cleanup_downloads() { - if let Some(download_dir) = get_download_dir() { - let _cleaned = std::fs::remove_dir_all(&download_dir); - } -} - #[cfg(test)] mod tests { use super::*; @@ -739,7 +596,6 @@ mod tests { }, ], draft: false, - prerelease: false, }; let update = select_release_target(&release).expect("update target"); @@ -764,7 +620,6 @@ mod tests { digest: None, }], draft: false, - prerelease: false, }; let update = select_release_target(&release).expect("update target"); @@ -776,64 +631,6 @@ mod tests { assert!(!update.supports_auto_apply()); } - #[test] - fn finds_newest_pending_installer_and_ignores_portable_exe() { - let temp = tempfile::tempdir().expect("temp dir"); - let (major, minor, patch) = parse_version_triplet(CURRENT_VERSION); - let portable = temp.path().join("codexbar.exe"); - let older = temp - .path() - .join(format!("CodexBar-{}.{}.{}-Setup.exe", major, minor, patch)); - let newer = temp.path().join(format!( - "CodexBar-{}.{}.{}-Setup.exe", - major, - minor, - patch + 1 - )); - - std::fs::write(&portable, b"portable").expect("write portable"); - std::fs::write(&older, b"older installer").expect("write older installer"); - std::fs::write(&newer, b"newer installer").expect("write newer installer"); - - let pending = find_pending_installer_in_dir(temp.path()).expect("pending installer"); - - assert_eq!(pending, newer); - } - - #[test] - fn ignores_cached_installers_for_current_or_older_versions() { - let temp = tempfile::tempdir().expect("temp dir"); - let (major, minor, patch) = parse_version_triplet(CURRENT_VERSION); - let current = temp - .path() - .join(format!("CodexBar-{}.{}.{}-Setup.exe", major, minor, patch)); - let older = temp.path().join(format!( - "CodexBar-{}.{}.{}-Setup.exe", - major, - minor, - patch.saturating_sub(1) - )); - - std::fs::write(¤t, b"current installer").expect("write current installer"); - std::fs::write(&older, b"older installer").expect("write older installer"); - - assert!(find_pending_installer_in_dir(temp.path()).is_none()); - } - - #[test] - fn parses_prerelease_installer_names_for_beta_updates() { - let (major, minor, patch) = parse_version_triplet(CURRENT_VERSION); - assert_eq!( - installer_version_from_name(&format!( - "CodexBar-{}.{}.{}-beta.1-Setup.exe", - major, - minor, - patch + 1 - )), - Some((major, minor, patch + 1)) - ); - } - #[test] fn verify_installer_hash_accepts_matching_sha256() { let temp = tempfile::tempdir().expect("temp dir"); diff --git a/rust/src/wsl.rs b/rust/src/wsl.rs index c8e678b4a2..24ca5815af 100644 --- a/rust/src/wsl.rs +++ b/rust/src/wsl.rs @@ -105,25 +105,6 @@ fn is_system_user_dir(name: &str) -> bool { ) } -/// Convert a Windows path to its WSL equivalent. -/// -/// `C:\Users\John\AppData\Local` becomes `/mnt/c/Users/John/AppData/Local`. -#[allow( - dead_code, - reason = "WSL detection helper reserved for future cross-platform integration" -)] -pub fn windows_path_to_wsl(windows_path: &str) -> Option { - let path = windows_path.replace('\\', "/"); - - if path.len() >= 2 && path.as_bytes()[1] == b':' { - let drive_letter = (path.as_bytes()[0] as char).to_lowercase().next()?; - let rest = path[2..].trim_start_matches('/'); - return Some(PathBuf::from(format!("/mnt/{}/{}", drive_letter, rest))); - } - - None -} - /// Get the Windows AppData/Local path from within WSL pub fn windows_appdata_local() -> Option { let info = get_wsl_info()?; @@ -162,17 +143,4 @@ mod tests { assert!(!is_system_user_dir("John")); assert!(!is_system_user_dir("alice")); } - - #[test] - fn test_windows_path_to_wsl() { - assert_eq!( - windows_path_to_wsl(r"C:\Users\John\AppData\Local"), - Some(PathBuf::from("/mnt/c/Users/John/AppData/Local")) - ); - assert_eq!( - windows_path_to_wsl("D:\\Games"), - Some(PathBuf::from("/mnt/d/Games")) - ); - assert_eq!(windows_path_to_wsl("/home/user"), None); - } }