diff --git a/docs/providers/codex.md b/docs/providers/codex.md index 086dd35..ac2ee72 100644 --- a/docs/providers/codex.md +++ b/docs/providers/codex.md @@ -29,3 +29,9 @@ Codex usage recorded by pi can also be included. OpenQuota does not upload these - **Subscription usage unavailable** — replace an API-key-only login with a ChatGPT login. - **Session expired or revoked** — sign in again with `codex`. - **No local history** — check the active Codex data directory and the value of `CODEX_HOME`. + +## Multi-Account Support + +OpenQuota supports multiple Codex accounts through the Hermes credentials pool (`~/.hermes/auth.json`). When using the Hermes file format, OpenQuota detects and monitors all available accounts. + +By default, the `Weekly` metric is pinned to the system tray for all active Codex accounts. The tray icon automatically calculates and displays the average health (remaining limit) across all pinned metrics across all your configured accounts. diff --git a/scripts/verify/verify-provider-registry-contract.js b/scripts/verify/verify-provider-registry-contract.js index 9c47d51..97314b3 100644 --- a/scripts/verify/verify-provider-registry-contract.js +++ b/scripts/verify/verify-provider-registry-contract.js @@ -77,6 +77,7 @@ if (!runtimeBlock) { } const runtimeOrder = [ 'ClaudeProvider', + 'CodexProvider', ...[...runtimeBlock.matchAll(/Arc::new\((\w+Provider)::new\b/g)].map(([, provider]) => provider), ]; const expectedRuntimeOrder = [ diff --git a/src-tauri/src/commands/usage.rs b/src-tauri/src/commands/usage.rs index 97ed308..b5c24d0 100644 --- a/src-tauri/src/commands/usage.rs +++ b/src-tauri/src/commands/usage.rs @@ -24,7 +24,7 @@ pub async fn claim_codex_reset_credit( if !settings .enabled_provider_ids() .iter() - .any(|id| id == "codex") + .any(|id| crate::providers::provider_family(id) == "codex") { return Err("Codex is not enabled.".to_owned()); } @@ -36,7 +36,14 @@ pub async fn claim_codex_reset_credit( if outcome != ResetClaimOutcome::Failed { let observed_account_revision = AtomicU64::new(settings.account_revision()); - service.refresh("codex", true).await; + let codex_ids: Vec<_> = settings + .enabled_provider_ids() + .into_iter() + .filter(|id| crate::providers::provider_family(id) == "codex") + .collect(); + for codex_id in &codex_ids { + service.refresh(codex_id, true).await; + } let state = service.state(); emit_settings_if_account_changed(&app, &settings, &observed_account_revision); let _ = app.emit("usage-state", &state); diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 61ec336..dd5b1eb 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -47,11 +47,11 @@ use crate::{ pacing::NotificationEvaluator, pricing::PricingStore, providers::{ - antigravity::AntigravityProvider, claude, codex::reset_claim::CodexResetClaimService, - codex::CodexProvider, copilot::CopilotProvider, cursor::CursorProvider, - detect_local_credentials, devin::DevinProvider, grok::GrokProvider, kimi::KimiProvider, - minimax::MiniMaxProvider, opencode::OpenCodeProvider, openrouter::OpenRouterProvider, - zai::ZaiProvider, ProviderRegistry, UsageProvider, + antigravity::AntigravityProvider, claude, codex, + codex::reset_claim::CodexResetClaimService, copilot::CopilotProvider, + cursor::CursorProvider, detect_local_credentials, devin::DevinProvider, grok::GrokProvider, + kimi::KimiProvider, minimax::MiniMaxProvider, opencode::OpenCodeProvider, + openrouter::OpenRouterProvider, zai::ZaiProvider, ProviderRegistry, UsageProvider, }, storage::Storage, window::{ @@ -71,13 +71,27 @@ fn install_tray(app: &mut App) -> Result<(), Box> { }; #[cfg(not(target_os = "macos"))] let menu = { - let open = MenuItem::with_id(app, "open", "Open OpenQuota", true, None::<&str>)?; + let summary = MenuItem::with_id(app, "summary", "Average Limit: ...", false, None::<&str>)?; + let separator_top = PredefinedMenuItem::separator(app)?; + let open = MenuItem::with_id(app, "open", "Open Dashboard", true, None::<&str>)?; let customize = MenuItem::with_id(app, "customize", "Customize…", true, None::<&str>)?; let settings_item = MenuItem::with_id(app, "settings", "Settings…", true, None::<&str>)?; let separator = PredefinedMenuItem::separator(app)?; let quit = MenuItem::with_id(app, "quit", "Quit OpenQuota", true, None::<&str>)?; - Menu::with_items(app, &[&open, &customize, &settings_item, &separator, &quit])? + Menu::with_items( + app, + &[ + &summary, + &separator_top, + &open, + &customize, + &settings_item, + &separator, + &quit, + ], + )? }; + app.manage(menu.clone()); let icon = app .default_window_icon() @@ -367,9 +381,13 @@ pub fn run() { app_debug!("cache", "application database opened"); let pricing = Arc::new(PricingStore::new(app_data_dir.join("pricing"))?); let mut providers = claude::runtimes(storage.clone(), pricing.clone())?; + providers.extend( + codex::runtimes(storage.clone(), pricing.clone()).unwrap_or_else(|error| { + crate::app_error!("startup", "failed to launch codex runtimes ({error})"); + Vec::new() + }), + ); providers.extend(vec![ - Arc::new(CodexProvider::new(storage.clone(), pricing.clone())?) - as Arc, Arc::new(CursorProvider::new(pricing.clone())?) as Arc, Arc::new(AntigravityProvider::new( app_data_dir.join("antigravity").join("auth.json"), diff --git a/src-tauri/src/providers/codex/accounts.rs b/src-tauri/src/providers/codex/accounts.rs new file mode 100644 index 0000000..b013b6f --- /dev/null +++ b/src-tauri/src/providers/codex/accounts.rs @@ -0,0 +1,439 @@ +#![allow(dead_code)] +use std::{ + collections::{BTreeMap, HashSet}, + path::{Path, PathBuf}, +}; + +use serde::{Deserialize, Serialize}; + +use super::auth::{self, CodexAuthState}; +use crate::{ + hashing::sha256_hex, + storage::{Storage, StorageError}, +}; + +#[derive(Debug, Clone)] +pub(super) struct CodexAccountDiscovery { + pub default_account: Option, + pub accounts: Vec, +} + +#[derive(Debug, Clone)] +pub(super) struct CodexAccount { + pub id: String, + pub display_name: String, + pub label: Option, + pub identity: String, + pub auth_source: CodexAuthSource, + pub session_roots: Vec, +} + +/// Describes where a Codex account's credentials are stored. +#[derive(Debug, Clone)] +pub(crate) enum CodexAuthSource { + /// Default credential locations (standard paths + macOS Keychain). + /// Used when `CODEX_HOME` is not set or has a single value. + Standard, + /// A specific directory containing `auth.json`. + /// Used when `CODEX_HOME` has multiple comma-separated values. + Home(PathBuf), + /// A specific directory containing Hermes `auth.json`, and the credential ID. + Hermes(PathBuf, String), +} + +#[derive(Debug, Clone)] +struct DiscoveredCodexAccount { + identity: String, + auth_source: CodexAuthSource, + session_roots: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +struct StoredCodexAccountPayload { + #[serde(default, skip_serializing_if = "Option::is_none")] + label: Option, +} + +struct StoredAccountRecord { + provider_id: String, + label: Option, +} + +struct DiscoveredCodexAccounts { + default_account: Option, + accounts: Vec, +} + +pub(super) fn discover(storage: &Storage) -> Result { + let raw = discover_accounts(); + if raw.default_account.is_none() && raw.accounts.is_empty() { + return Ok(CodexAccountDiscovery { + default_account: None, + accounts: Vec::new(), + }); + } + reconcile_accounts(storage, raw) +} + +pub(super) fn identity_for_source(source: &CodexAuthSource) -> Option { + match source { + CodexAuthSource::Standard => { + CodexAuthState::observed_account_identity().map(|id| identity_stamp(&id)) + } + CodexAuthSource::Home(path) => { + let auth_path = path.join("auth.json"); + auth::load_from_path(&auth_path) + .ok()? + .account_identity() + .map(|id| identity_stamp(&id)) + } + CodexAuthSource::Hermes(path, hermes_id) => auth::load_hermes_from_path(path, hermes_id) + .ok()? + .account_identity() + .map(|id| identity_stamp(&id)), + } +} + +fn discover_accounts() -> DiscoveredCodexAccounts { + let home = home_directory(); + let configured_home = crate::provider_environment::value("CODEX_HOME"); + + let multi_paths: Option> = configured_home + .as_deref() + .filter(|v| !v.trim().is_empty()) + .map(|configured| { + configured + .split(',') + .map(str::trim) + .filter(|v| !v.is_empty()) + .map(|v| expand_home(v, &home)) + .collect() + }); + + if let Some(paths) = multi_paths { + let result = discover_from_multiple_homes(&paths); + if result.default_account.is_some() || !result.accounts.is_empty() { + crate::app_info!( + "config", + "codex account discovery completed ({} account(s) from {} CODEX_HOME path(s))", + 1 + result.accounts.len(), + paths.len() + ); + return result; + } + } + + discover_from_standard() +} + +fn discover_from_standard() -> DiscoveredCodexAccounts { + let identity = CodexAuthState::observed_account_identity(); + let default_account = identity.map(|raw_identity| DiscoveredCodexAccount { + identity: identity_stamp(&raw_identity), + auth_source: CodexAuthSource::Standard, + session_roots: Vec::new(), + }); + DiscoveredCodexAccounts { + default_account, + accounts: Vec::new(), + } +} + +fn discover_from_multiple_homes(homes: &[PathBuf]) -> DiscoveredCodexAccounts { + let mut seen_identities: BTreeMap = BTreeMap::new(); + let mut default_account = None; + let mut extra_accounts = Vec::new(); + + for (index, home_path) in homes.iter().enumerate() { + let auth_path = home_path.join("auth.json"); + let discovered = auth::discover_identities_from_path(&auth_path); + + for (raw_identity, auth_source) in discovered { + let identity = identity_stamp(&raw_identity); + if seen_identities.contains_key(&identity) { + continue; + } + seen_identities.insert(identity.clone(), index); + + let account = DiscoveredCodexAccount { + identity, + auth_source, + session_roots: vec![home_path.clone()], + }; + + if default_account.is_none() { + default_account = Some(account); + } else { + extra_accounts.push(account); + } + } + } + + DiscoveredCodexAccounts { + default_account, + accounts: extra_accounts, + } +} + +fn reconcile_accounts( + storage: &Storage, + discovery: DiscoveredCodexAccounts, +) -> Result { + let stored = storage.load_provider_account_records("codex")?; + let records: BTreeMap = stored + .into_iter() + .map(|(identity, provider_id, payload)| { + let label = serde_json::from_str::(&payload) + .ok() + .and_then(|p| p.label); + (identity, StoredAccountRecord { provider_id, label }) + }) + .collect(); + + let mut occupied: HashSet = records.values().map(|r| r.provider_id.clone()).collect(); + + let has_bare_scoped_account = discovery.accounts.iter().any(|account| { + records + .get(&account.identity) + .is_some_and(|r| r.provider_id == "codex") + }); + + let default_account = discovery + .default_account + .map(|account| { + reconcile_account( + storage, + account, + &records, + &mut occupied, + !has_bare_scoped_account, + ) + }) + .transpose()?; + + let mut accounts: Vec = discovery + .accounts + .into_iter() + .map(|account| reconcile_account(storage, account, &records, &mut occupied, false)) + .collect::, _>>()?; + accounts.sort_by(|a, b| a.id.cmp(&b.id)); + + Ok(CodexAccountDiscovery { + default_account, + accounts, + }) +} + +fn reconcile_account( + storage: &Storage, + account: DiscoveredCodexAccount, + records: &BTreeMap, + occupied: &mut HashSet, + may_claim_default_id: bool, +) -> Result { + let record = records.get(&account.identity); + let label = record.and_then(|r| r.label.clone()); + let id = record.map(|r| r.provider_id.clone()).unwrap_or_else(|| { + if may_claim_default_id && !occupied.contains("codex") { + "codex".to_owned() + } else { + allocate_account_id(&account.identity, occupied) + } + }); + occupied.insert(id.clone()); + + let reconciled = CodexAccount { + display_name: account_display_name_for_id(label.as_deref(), &id), + id, + label, + identity: account.identity, + auth_source: account.auth_source, + session_roots: account.session_roots, + }; + + let payload = serde_json::to_string(&StoredCodexAccountPayload { + label: reconciled.label.clone(), + })?; + storage.save_provider_account_record( + "codex", + &reconciled.identity, + &reconciled.id, + &payload, + )?; + + Ok(reconciled) +} + +fn allocate_account_id(identity_stamp: &str, occupied: &HashSet) -> String { + for salt in 0_u64.. { + let stamp = if salt == 0 { + identity_stamp.to_owned() + } else { + sha256_hex(format!("{identity_stamp}:{salt}").as_bytes()) + }; + let candidate = format!("codex@{}", &stamp[..8]); + if !occupied.contains(&candidate) { + return candidate; + } + } + unreachable!("an account ID is always available") +} + +fn account_display_name_for_id(label: Option<&str>, id: &str) -> String { + if id == "codex" { + "Codex".to_owned() + } else if let Some(label) = label.map(str::trim).filter(|v| !v.is_empty()) { + format!("Codex — {label}") + } else { + id.to_owned() + } +} + +fn identity_stamp(identity: &str) -> String { + sha256_hex(identity.to_ascii_lowercase().as_bytes()) +} + +fn expand_home(value: &str, home: &Path) -> PathBuf { + if value == "~" { + return home.to_path_buf(); + } + value + .strip_prefix("~/") + .or_else(|| value.strip_prefix("~\\")) + .map(|rest| home.join(rest)) + .unwrap_or_else(|| PathBuf::from(value)) +} + +fn home_directory() -> PathBuf { + std::env::var_os("HOME") + .or_else(|| std::env::var_os("USERPROFILE")) + .map(PathBuf::from) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + + use tempfile::tempdir; + + use crate::storage::Storage; + + use super::{ + account_display_name_for_id, allocate_account_id, identity_stamp, reconcile_accounts, + CodexAuthSource, DiscoveredCodexAccount, DiscoveredCodexAccounts, + }; + + #[test] + fn default_account_claims_bare_id() { + let directory = tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("openquota.db")).unwrap(); + let discovery = DiscoveredCodexAccounts { + default_account: Some(DiscoveredCodexAccount { + identity: identity_stamp("account-a"), + auth_source: CodexAuthSource::Standard, + session_roots: Vec::new(), + }), + accounts: Vec::new(), + }; + + let result = reconcile_accounts(&storage, discovery).unwrap(); + + assert_eq!(result.default_account.as_ref().unwrap().id, "codex"); + assert!(result.accounts.is_empty()); + } + + #[test] + fn extra_accounts_receive_hashed_ids() { + let directory = tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("openquota.db")).unwrap(); + let discovery = DiscoveredCodexAccounts { + default_account: Some(DiscoveredCodexAccount { + identity: identity_stamp("account-a"), + auth_source: CodexAuthSource::Standard, + session_roots: Vec::new(), + }), + accounts: vec![DiscoveredCodexAccount { + identity: identity_stamp("account-b"), + auth_source: CodexAuthSource::Home("/tmp/codex-work".into()), + session_roots: vec!["/tmp/codex-work".into()], + }], + }; + + let result = reconcile_accounts(&storage, discovery).unwrap(); + + assert_eq!(result.default_account.as_ref().unwrap().id, "codex"); + assert_eq!(result.accounts.len(), 1); + assert!(result.accounts[0].id.starts_with("codex@")); + assert_eq!(result.accounts[0].id.len(), "codex@".len() + 8); + } + + #[test] + fn stable_ids_across_restarts() { + let directory = tempdir().unwrap(); + let storage = Storage::open(&directory.path().join("openquota.db")).unwrap(); + let identity_b = identity_stamp("account-b"); + + // First discovery + let discovery = DiscoveredCodexAccounts { + default_account: Some(DiscoveredCodexAccount { + identity: identity_stamp("account-a"), + auth_source: CodexAuthSource::Standard, + session_roots: Vec::new(), + }), + accounts: vec![DiscoveredCodexAccount { + identity: identity_b.clone(), + auth_source: CodexAuthSource::Home("/tmp/codex-work".into()), + session_roots: vec!["/tmp/codex-work".into()], + }], + }; + let first = reconcile_accounts(&storage, discovery).unwrap(); + let first_extra_id = first.accounts[0].id.clone(); + + // Second discovery with same accounts + let discovery2 = DiscoveredCodexAccounts { + default_account: Some(DiscoveredCodexAccount { + identity: identity_stamp("account-a"), + auth_source: CodexAuthSource::Standard, + session_roots: Vec::new(), + }), + accounts: vec![DiscoveredCodexAccount { + identity: identity_b, + auth_source: CodexAuthSource::Home("/tmp/codex-work".into()), + session_roots: vec!["/tmp/codex-work".into()], + }], + }; + let second = reconcile_accounts(&storage, discovery2).unwrap(); + + assert_eq!( + first.default_account.unwrap().id, + second.default_account.unwrap().id + ); + assert_eq!(first_extra_id, second.accounts[0].id); + } + + #[test] + fn allocate_avoids_collisions() { + let mut occupied = std::collections::HashSet::new(); + let stamp = identity_stamp("test"); + let first = allocate_account_id(&stamp, &occupied); + occupied.insert(first.clone()); + let second = allocate_account_id(&stamp, &occupied); + + assert_ne!(first, second); + assert!(first.starts_with("codex@")); + assert!(second.starts_with("codex@")); + } + + #[test] + fn display_name_for_default_and_extra_accounts() { + assert_eq!(account_display_name_for_id(None, "codex"), "Codex"); + assert_eq!( + account_display_name_for_id(Some("Work"), "codex@12345678"), + "Codex — Work" + ); + assert_eq!( + account_display_name_for_id(None, "codex@12345678"), + "codex@12345678" + ); + } +} diff --git a/src-tauri/src/providers/codex/auth.rs b/src-tauri/src/providers/codex/auth.rs index 6b510a2..c62406e 100644 --- a/src-tauri/src/providers/codex/auth.rs +++ b/src-tauri/src/providers/codex/auth.rs @@ -1,3 +1,4 @@ +#![allow(dead_code)] use std::{ fs, io::Write, @@ -27,6 +28,7 @@ pub struct CodexAuthState { #[derive(Debug, Clone)] enum AuthSource { File(PathBuf), + Hermes(PathBuf, String), #[cfg(target_os = "macos")] Keychain, } @@ -96,7 +98,28 @@ impl CodexAuthState { .find_map(|state| state.account_identity()) } - pub(super) fn account_identity(&self) -> Option { + pub fn load_candidates_scoped( + source: &super::accounts::CodexAuthSource, + ) -> Result, CodexError> { + match source { + super::accounts::CodexAuthSource::Standard => Self::load_candidates(), + super::accounts::CodexAuthSource::Home(path) => { + let auth_path = path.join("auth.json"); + let state = load_from_path(&auth_path)?; + Ok(vec![state]) + } + super::accounts::CodexAuthSource::Hermes(path, hermes_id) => { + let state = load_hermes_from_path(path, hermes_id)?; + Ok(vec![state]) + } + } + } + + pub fn has_local_credentials_scoped(source: &super::accounts::CodexAuthSource) -> bool { + Self::load_candidates_scoped(source).is_ok() + } + + pub fn account_identity(&self) -> Option { self.account_id .as_deref() .and_then(nonempty_lowercase) @@ -113,11 +136,21 @@ impl CodexAuthState { .and_then(nonempty_lowercase) }) }) + .or_else(|| { + jwt_payload(&self.access_token).and_then(|payload| { + payload + .pointer("/https:~1~1api.openai.com~1auth/chatgpt_account_id") + .or_else(|| payload.get("chatgpt_account_id")) + .and_then(Value::as_str) + .and_then(nonempty_lowercase) + }) + }) } pub fn reload(&self) -> Result { match &self.source { AuthSource::File(path) => load_from_path(path), + AuthSource::Hermes(path, hermes_id) => load_hermes_from_path(path, hermes_id), #[cfg(target_os = "macos")] AuthSource::Keychain => load_from_keychain(), } @@ -155,28 +188,59 @@ impl CodexAuthState { id_token: Option, now: DateTime, ) -> Result<(), CodexError> { - set_string(&mut self.document, "/tokens/access_token", &access_token)?; - if let Some(value) = refresh_token.as_deref() { - set_string(&mut self.document, "/tokens/refresh_token", value)?; - self.refresh_token = Some(value.to_owned()); - } - if let Some(value) = id_token.as_deref() { - set_string(&mut self.document, "/tokens/id_token", value)?; - } let refreshed_at = now.to_rfc3339(); - set_string(&mut self.document, "/last_refresh", &refreshed_at)?; - self.access_token = access_token; - self.last_refresh = Some(refreshed_at); match &self.source { - AuthSource::File(path) => save_file_document(path, &self.document), + AuthSource::File(path) => { + set_string(&mut self.document, "/tokens/access_token", &access_token)?; + if let Some(value) = refresh_token.as_deref() { + set_string(&mut self.document, "/tokens/refresh_token", value)?; + self.refresh_token = Some(value.to_owned()); + } + if let Some(value) = id_token.as_deref() { + set_string(&mut self.document, "/tokens/id_token", value)?; + } + set_string(&mut self.document, "/last_refresh", &refreshed_at)?; + self.access_token = access_token; + self.last_refresh = Some(refreshed_at); + save_file_document(path, &self.document) + } + AuthSource::Hermes(path, hermes_id) => { + update_hermes_credential( + &mut self.document, + hermes_id, + &access_token, + refresh_token.as_deref(), + id_token.as_deref(), + &refreshed_at, + )?; + self.access_token = access_token; + if let Some(value) = refresh_token { + self.refresh_token = Some(value); + } + self.last_refresh = Some(refreshed_at); + save_file_document(path, &self.document) + } #[cfg(target_os = "macos")] - AuthSource::Keychain => save_keychain_document(&self.document), + AuthSource::Keychain => { + set_string(&mut self.document, "/tokens/access_token", &access_token)?; + if let Some(value) = refresh_token.as_deref() { + set_string(&mut self.document, "/tokens/refresh_token", value)?; + self.refresh_token = Some(value.to_owned()); + } + if let Some(value) = id_token.as_deref() { + set_string(&mut self.document, "/tokens/id_token", value)?; + } + set_string(&mut self.document, "/last_refresh", &refreshed_at)?; + self.access_token = access_token; + self.last_refresh = Some(refreshed_at); + save_keychain_document(&self.document) + } } } } -fn load_from_path(path: &Path) -> Result { +pub(super) fn load_from_path(path: &Path) -> Result { let text = fs::read_to_string(path).map_err(|_| CodexError::InvalidAuth)?; let document = parse_auth_document(&text).ok_or(CodexError::InvalidAuth)?; let access_token = string_at(&document, "/tokens/access_token") @@ -350,6 +414,142 @@ fn set_string(document: &mut Value, pointer: &str, value: &str) -> Result<(), Co Ok(()) } +pub(super) fn discover_identities_from_path( + path: &Path, +) -> Vec<(String, super::accounts::CodexAuthSource)> { + let mut identities = Vec::new(); + let Ok(text) = fs::read_to_string(path) else { + return identities; + }; + let Ok(document) = serde_json::from_str::(&text) else { + return identities; + }; + + if document.get("version").is_some() && document.get("credential_pool").is_some() { + if let Some(pool) = document + .pointer("/credential_pool/openai-codex") + .and_then(Value::as_array) + { + for cred in pool { + if let Some(hermes_id) = cred.get("id").and_then(Value::as_str) { + if let Some(access_token) = cred + .get("access_token") + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + { + let state = CodexAuthState { + source: AuthSource::Hermes(path.to_path_buf(), hermes_id.to_owned()), + document: document.clone(), + access_token: access_token.to_owned(), + refresh_token: cred + .get("refresh_token") + .and_then(Value::as_str) + .map(str::to_owned), + account_id: None, + last_refresh: cred + .get("last_refresh") + .and_then(Value::as_str) + .map(str::to_owned), + }; + if let Some(identity) = state.account_identity() { + identities.push(( + identity, + super::accounts::CodexAuthSource::Hermes( + path.to_path_buf(), + hermes_id.to_owned(), + ), + )); + } + } + } + } + } + } else { + if let Ok(state) = load_from_path(path) { + if let Some(identity) = state.account_identity() { + identities.push(( + identity, + super::accounts::CodexAuthSource::Home(path.to_path_buf()), + )); + } + } + } + identities +} + +pub(super) fn load_hermes_from_path( + path: &Path, + hermes_id: &str, +) -> Result { + let text = fs::read_to_string(path).map_err(|_| CodexError::InvalidAuth)?; + let document: Value = serde_json::from_str(&text).map_err(|_| CodexError::InvalidAuth)?; + + let pool = document + .pointer("/credential_pool/openai-codex") + .and_then(Value::as_array) + .ok_or(CodexError::InvalidAuth)?; + let cred = pool + .iter() + .find(|c| c.get("id").and_then(Value::as_str) == Some(hermes_id)) + .ok_or(CodexError::NotLoggedIn)?; + + let access_token = cred + .get("access_token") + .and_then(Value::as_str) + .filter(|s| !s.is_empty()) + .ok_or(CodexError::NotLoggedIn)? + .to_owned(); + + Ok(CodexAuthState { + source: AuthSource::Hermes(path.to_path_buf(), hermes_id.to_owned()), + refresh_token: cred + .get("refresh_token") + .and_then(Value::as_str) + .map(str::to_owned), + account_id: None, + last_refresh: cred + .get("last_refresh") + .and_then(Value::as_str) + .map(str::to_owned), + document, + access_token, + }) +} + +fn update_hermes_credential( + document: &mut Value, + hermes_id: &str, + access_token: &str, + refresh_token: Option<&str>, + _id_token: Option<&str>, + refreshed_at: &str, +) -> Result<(), CodexError> { + let pool = document + .pointer_mut("/credential_pool/openai-codex") + .and_then(Value::as_array_mut) + .ok_or(CodexError::AuthWrite)?; + let cred = pool + .iter_mut() + .find(|c| c.get("id").and_then(Value::as_str) == Some(hermes_id)) + .ok_or(CodexError::AuthWrite)?; + + if let Some(obj) = cred.as_object_mut() { + obj.insert( + "access_token".to_string(), + Value::String(access_token.to_string()), + ); + if let Some(rt) = refresh_token { + obj.insert("refresh_token".to_string(), Value::String(rt.to_string())); + } + obj.insert( + "last_refresh".to_string(), + Value::String(refreshed_at.to_string()), + ); + Ok(()) + } else { + Err(CodexError::AuthWrite) + } +} #[cfg(test)] mod tests { use std::{fs, path::Path}; diff --git a/src-tauri/src/providers/codex/client.rs b/src-tauri/src/providers/codex/client.rs index 6fcb06c..f4ea501 100644 --- a/src-tauri/src/providers/codex/client.rs +++ b/src-tauri/src/providers/codex/client.rs @@ -27,6 +27,7 @@ pub struct TokenRefresh { pub id_token: Option, } +#[derive(Clone)] pub struct CodexClient { client: Client, refresh_url: String, diff --git a/src-tauri/src/providers/codex/local_usage.rs b/src-tauri/src/providers/codex/local_usage.rs index 1107c65..8019ec6 100644 --- a/src-tauri/src/providers/codex/local_usage.rs +++ b/src-tauri/src/providers/codex/local_usage.rs @@ -1,3 +1,4 @@ +#![allow(dead_code)] use std::{ collections::HashSet, ffi::OsStr, @@ -37,33 +38,33 @@ pub struct TokenEvent { const LOG_CACHE_SCHEMA_VERSION: u8 = 3; -pub fn scan_local_usage( +pub fn scan_local_usage_scoped( storage: &Storage, now: DateTime, pricing: &ModelPricing, + provider_id: &str, + session_roots: &[PathBuf], ) -> Result { - let home = home_directory(); - let configured_home = crate::provider_environment::value("CODEX_HOME"); - let homes = codex_homes(configured_home.as_deref().map(OsStr::new), &home); let since_date = now .with_timezone(&Local) .date_naive() .checked_sub_days(Days::new(30)) .unwrap_or(NaiveDate::MIN); - let events = scan_codex_events(storage, &homes, since_date)?; + let events = scan_codex_events(storage, provider_id, session_roots, since_date)?; let mut accumulator = DailyUsageAccumulator::default(); aggregate_into(events, now, pricing, &mut accumulator); - let includes_pi = match pi_usage::scan_into(storage, now, pricing, "codex", &mut accumulator) { - Ok(includes_pi) => includes_pi, - Err(_) => { - crate::app_warn!( - "plugin:pi", - "pi usage history could not be folded into Codex" - ); - false - } - }; + let includes_pi = + match pi_usage::scan_into(storage, now, pricing, provider_id, &mut accumulator) { + Ok(includes_pi) => includes_pi, + Err(_) => { + crate::app_warn!( + "plugin:pi", + "pi usage history could not be folded into Codex" + ); + false + } + }; let source_note = if includes_pi { "From your Codex logs and pi (estimated)" } else { @@ -74,6 +75,7 @@ pub fn scan_local_usage( fn scan_codex_events( storage: &Storage, + provider_id: &str, homes: &[PathBuf], since_date: NaiveDate, ) -> Result, CodexError> { @@ -85,7 +87,7 @@ fn scan_codex_events( seen_paths.insert(path.clone()); let Some(parsed) = load_or_parse_log( storage, - "codex", + provider_id, &path, LOG_CACHE_SCHEMA_VERSION, parse_jsonl, @@ -103,7 +105,7 @@ fn scan_codex_events( .filter(|event| event.timestamp.with_timezone(&Local).date_naive() >= since_date), ); } - storage.prune_log_events("codex", &seen_paths)?; + storage.prune_log_events(provider_id, &seen_paths)?; Ok(events) } @@ -851,7 +853,8 @@ mod tests { let storage = Storage::open(&directory.path().join("openquota.db")).unwrap(); let since = NaiveDate::from_ymd_opt(2026, 7, 1).unwrap(); - let initial = scan_codex_events(&storage, std::slice::from_ref(&home), since).unwrap(); + let initial = + scan_codex_events(&storage, "codex", std::slice::from_ref(&home), since).unwrap(); assert_eq!(initial.len(), 1); assert!(!initial[0].is_fast); @@ -866,7 +869,7 @@ mod tests { ) .unwrap(); - let refreshed = scan_codex_events(&storage, &[home], since).unwrap(); + let refreshed = scan_codex_events(&storage, "codex", &[home], since).unwrap(); assert_eq!(refreshed.len(), 2); assert_eq!(refreshed.iter().map(|event| event.total).sum::(), 195); assert!(refreshed.iter().all(|event| !event.is_fast)); @@ -900,6 +903,7 @@ mod tests { let events = scan_codex_events( &storage, + "codex", &[home], NaiveDate::from_ymd_opt(2026, 3, 1).unwrap(), ) diff --git a/src-tauri/src/providers/codex/mod.rs b/src-tauri/src/providers/codex/mod.rs index c8bf0a1..efe2eab 100644 --- a/src-tauri/src/providers/codex/mod.rs +++ b/src-tauri/src/providers/codex/mod.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] +pub mod accounts; pub mod auth; pub mod client; pub mod local_usage; @@ -20,17 +22,19 @@ use crate::{ storage::Storage, }; -use self::{ - auth::CodexAuthState, client::CodexClient, local_usage::scan_local_usage, mapper::map_usage, -}; +use self::{auth::CodexAuthState, client::CodexClient, mapper::map_usage}; use crate::providers::log_usage::scan_or_cached_usage; pub(crate) fn definition() -> ProviderDefinition { - ProviderDefinition { + definition_for("codex", "Codex", true) +} + +fn definition_for(id: &str, display_name: &str, fallback_enabled: bool) -> ProviderDefinition { + let mut definition = ProviderDefinition { id: "codex".into(), - display_name: "Codex".into(), + display_name: display_name.into(), short_name: "Cx".into(), - fallback_enabled: true, + fallback_enabled, local_usage_source_note: Some("From your Codex logs (estimated)".into()), links: vec![ ProviderLink::new("Status", "https://status.openai.com/"), @@ -44,7 +48,7 @@ pub(crate) fn definition() -> ProviderDefinition { false, true, MetricSection::AlwaysVisible, - true, + false, "S", ), MetricDefinition::quota( @@ -120,7 +124,17 @@ pub(crate) fn definition() -> ProviderDefinition { "M", ), ], + }; + + if id != "codex" { + definition.id = id.into(); + for metric in &mut definition.metrics { + if let Some(suffix) = metric.id.strip_prefix("codex.") { + metric.id = format!("{id}.{suffix}"); + } + } } + definition } #[derive(Debug, Error)] @@ -163,26 +177,65 @@ impl From for CodexError { } } +pub(crate) fn runtimes( + storage: Arc, + pricing: Arc, +) -> Result>, CodexError> { + let discovery = accounts::discover(&storage)?; + let client = CodexClient::new()?; + let mut runtimes = Vec::new(); + + if let Some(account) = discovery.default_account { + runtimes.push(Arc::new(CodexProvider::new_scoped( + account, + storage.clone(), + pricing.clone(), + client.clone(), + )) as Arc); + } + + for account in discovery.accounts { + runtimes.push(Arc::new(CodexProvider::new_scoped( + account, + storage.clone(), + pricing.clone(), + client.clone(), + )) as Arc); + } + + Ok(runtimes) +} + pub struct CodexProvider { + definition: ProviderDefinition, + auth_source: accounts::CodexAuthSource, account_identity: Option, + session_roots: Vec, storage: Arc, pricing: Arc, client: CodexClient, } impl CodexProvider { - pub fn new(storage: Arc, pricing: Arc) -> Result { - let account_identity = CodexAuthState::observed_account_identity() - .map(|identity| account_identity_key(&identity)); - if let Some(identity) = account_identity.as_deref() { - crate::providers::remember_default_account(&storage, "codex", identity)?; - } - Ok(Self { - account_identity, + fn new_scoped( + account: accounts::CodexAccount, + storage: Arc, + pricing: Arc, + client: CodexClient, + ) -> Self { + Self { + definition: definition_for(&account.id, &account.display_name, account.id == "codex"), + auth_source: account.auth_source, + account_identity: Some(account.identity), + session_roots: account.session_roots, storage, pricing, - client: CodexClient::new()?, - }) + client, + } + } + + fn provider_id(&self) -> &str { + &self.definition.id } pub fn refresh(&self) -> Result { @@ -191,7 +244,7 @@ impl CodexProvider { fn refresh_with_identity(&self) -> Result<(ProviderSnapshot, Option), CodexError> { let now = Utc::now(); - let candidates = CodexAuthState::load_candidates()?; + let candidates = CodexAuthState::load_candidates_scoped(&self.auth_source)?; crate::app_debug!( "auth:codex", "credential candidates loaded ({})", @@ -279,17 +332,25 @@ impl CodexProvider { let pricing = self.pricing.current(); let usage = scan_or_cached_usage( &self.storage, - "codex", + self.provider_id(), account_identity .map(crate::providers::CacheIdentity::Resolved) .unwrap_or(crate::providers::CacheIdentity::Unresolved), - "Codex", - || scan_local_usage(&self.storage, now, &pricing), + &self.definition.display_name, + || { + local_usage::scan_local_usage_scoped( + &self.storage, + now, + &pricing, + self.provider_id(), + &self.session_roots, + ) + }, &mut warnings, ); Self::ensure_candidate_source_current(auth, account_identity)?; Ok(ProviderSnapshot { - provider_id: "codex".into(), + provider_id: self.provider_id().into(), plan: mapped.plan, quotas: mapped.quotas, value_metrics: mapped.value_metrics, @@ -372,11 +433,11 @@ fn provider_error(error: CodexError) -> crate::providers::ProviderError { impl crate::providers::UsageProvider for CodexProvider { fn definition(&self) -> ProviderDefinition { - definition() + self.definition.clone() } fn has_local_credentials(&self) -> bool { - CodexAuthState::has_local_credentials() + CodexAuthState::has_local_credentials_scoped(&self.auth_source) } fn cache_identity(&self) -> crate::providers::CacheIdentity<'_> { @@ -405,10 +466,10 @@ impl crate::providers::UsageProvider for CodexProvider { Ok(crate::providers::ProviderRefresh { snapshot, cache_identity: identity.clone(), - account: identity.map(|identity| crate::providers::AccountRefresh { + account: identity.map(|id| crate::providers::AccountRefresh { family: "codex", - provider_id: "codex", - identity, + provider_id: Box::leak(self.definition.id.clone().into_boxed_str()), + identity: id, }), }) } @@ -451,12 +512,18 @@ mod account_tests { let storage = Arc::new(Storage::open(&directory.path().join("openquota.db")).unwrap()); let pricing = Arc::new(PricingStore::new(directory.path().join("pricing")).unwrap()); let provider = CodexProvider { + definition: super::definition(), + auth_source: super::accounts::CodexAuthSource::Standard, + session_roots: Vec::new(), account_identity: Some("account-a".into()), storage: storage.clone(), pricing: pricing.clone(), client: CodexClient::new().unwrap(), }; let unresolved = CodexProvider { + definition: super::definition(), + auth_source: super::accounts::CodexAuthSource::Standard, + session_roots: Vec::new(), account_identity: None, storage, pricing, diff --git a/src-tauri/src/providers/mod.rs b/src-tauri/src/providers/mod.rs index 57cf18f..2c3905f 100644 --- a/src-tauri/src/providers/mod.rs +++ b/src-tauri/src/providers/mod.rs @@ -1,3 +1,4 @@ +#![allow(dead_code)] pub mod antigravity; pub mod api_key; pub mod claude; @@ -32,10 +33,21 @@ pub fn provider_family(provider_id: &str) -> &str { .unwrap_or(provider_id) } +pub fn is_account_provider_id(provider_id: &str, family: &str) -> bool { + provider_id + .strip_prefix(family) + .and_then(|rest| rest.strip_prefix('@')) + .is_some_and(|suffix| { + suffix.len() == 8 && suffix.bytes().all(|byte| byte.is_ascii_hexdigit()) + }) +} + pub fn is_claude_account_provider_id(provider_id: &str) -> bool { - provider_id.strip_prefix("claude@").is_some_and(|suffix| { - suffix.len() == 8 && suffix.bytes().all(|byte| byte.is_ascii_hexdigit()) - }) + is_account_provider_id(provider_id, "claude") +} + +pub fn is_codex_account_provider_id(provider_id: &str) -> bool { + is_account_provider_id(provider_id, "codex") } pub fn remember_default_account( diff --git a/src-tauri/src/settings.rs b/src-tauri/src/settings.rs index f7d2464..f11cdab 100644 --- a/src-tauri/src/settings.rs +++ b/src-tauri/src/settings.rs @@ -785,7 +785,9 @@ fn normalize_with_persisted_accounts( let mut provider = default_provider(definition, is_detected); provider.enabled = !was_known && is_detected; settings.known_provider_ids.push(definition.id.clone()); - if crate::providers::is_claude_account_provider_id(&provider.id) { + if crate::providers::is_claude_account_provider_id(&provider.id) + || crate::providers::is_codex_account_provider_id(&provider.id) + { let family = crate::providers::provider_family(&provider.id); let index = normalized .iter() diff --git a/src-tauri/src/tray_presentation.rs b/src-tauri/src/tray_presentation.rs index 52b8bde..a8ecbeb 100644 --- a/src-tauri/src/tray_presentation.rs +++ b/src-tauri/src/tray_presentation.rs @@ -1,7 +1,5 @@ use tauri::{image::Image, AppHandle}; -#[cfg(not(target_os = "macos"))] -use crate::tray_icon; use crate::{ models::{ AppSettings, MetricDefinition, MetricSource, MetricValue, MetricValueKind, @@ -10,6 +8,8 @@ use crate::{ providers::ProviderRegistry, service::UsageViewState, }; +#[cfg(not(target_os = "macos"))] +use {crate::tray_icon, tauri::Manager}; const TRAY_ID: &str = "openquota-tray"; @@ -30,7 +30,7 @@ struct TrayGroup { #[derive(Debug, Clone, Copy, PartialEq)] struct TrayGauge { display_fraction: f64, - #[cfg(any(not(target_os = "macos"), test))] + remaining_fraction: f64, } @@ -61,15 +61,22 @@ pub fn update( let tooltip = if groups.is_empty() { "OpenQuota".to_owned() } else { - format!( - "OpenQuota\n{}", - groups - .iter() - .flat_map(|group| group.metrics.iter()) - .map(|metric| metric.detail.as_str()) - .collect::>() - .join(" · ") - ) + let details = groups + .iter() + .flat_map(|group| group.metrics.iter()) + .map(|metric| metric.detail.as_str()) + .collect::>() + .join(" · "); + + if let Some(gauge) = primary_gauge(&groups) { + format!( + "OpenQuota\nAverage Limit: {:.0}% left\n{}", + gauge.remaining_fraction * 100.0, + details + ) + } else { + format!("OpenQuota\n{}", details) + } }; #[cfg(not(target_os = "linux"))] if tray.set_tooltip(Some(tooltip)).is_err() { @@ -80,7 +87,25 @@ pub fn update( #[cfg(not(target_os = "macos"))] { - let icon = primary_gauge(&groups) + let primary = primary_gauge(&groups); + + if let Some(menu) = app.try_state::>() { + if let Some(summary_item) = menu.get("summary") { + if let Some(summary_menu_item) = summary_item.as_menuitem() { + let text = if let Some(gauge) = primary { + format!( + "Average Limit: {:.0}% left", + gauge.remaining_fraction * 100.0 + ) + } else { + "OpenQuota".to_owned() + }; + let _ = summary_menu_item.set_text(text); + } + } + } + + let icon = primary .map(|gauge| tray_icon::render_gauge(gauge.display_fraction, gauge.remaining_fraction)) .unwrap_or_else(mark_icon); if tray.set_icon(Some(icon)).is_err() { @@ -178,12 +203,24 @@ fn bar_fractions(groups: &[TrayGroup]) -> Vec { .collect() } -#[cfg(any(not(target_os = "macos"), test))] fn primary_gauge(groups: &[TrayGroup]) -> Option { - groups + let primary_gauges: Vec = groups .iter() - .flat_map(|group| group.metrics.iter()) - .find_map(|metric| metric.gauge) + .filter_map(|group| group.metrics.iter().find_map(|metric| metric.gauge)) + .collect(); + + if primary_gauges.is_empty() { + return None; + } + + let sum_display: f64 = primary_gauges.iter().map(|g| g.display_fraction).sum(); + let sum_remaining: f64 = primary_gauges.iter().map(|g| g.remaining_fraction).sum(); + let count = primary_gauges.len() as f64; + + Some(TrayGauge { + display_fraction: sum_display / count, + remaining_fraction: sum_remaining / count, + }) } #[cfg(any(target_os = "macos", test))] @@ -275,7 +312,7 @@ fn tray_metric( UsageDisplay::Used => used_fraction, UsageDisplay::Left => 1.0 - used_fraction, }, - #[cfg(any(not(target_os = "macos"), test))] + remaining_fraction: 1.0 - used_fraction, }), }; @@ -296,7 +333,7 @@ fn tray_metric( detail: format!("{} {percent:.0}% {word}", quota.label), gauge: Some(TrayGauge { display_fraction, - #[cfg(any(not(target_os = "macos"), test))] + remaining_fraction: 1.0 - used_fraction, }), } @@ -626,11 +663,13 @@ mod tests { last_full_refresh_at: None, }; let catalog = ProviderRegistry::from_definitions(vec![codex::definition()]).unwrap(); - let groups = resolved_groups( - &state, - &default_settings(&catalog, &HashSet::from(["codex".to_owned()])), - &catalog, - ); + let mut test_settings = default_settings(&catalog, &HashSet::from(["codex".to_owned()])); + test_settings.providers[0].metrics.iter_mut().for_each(|m| { + if m.id == "codex.session" { + m.pinned = true; + } + }); + let groups = resolved_groups(&state, &test_settings, &catalog); assert_eq!(groups.len(), 1); assert_eq!(groups[0].provider_id, "codex"); assert_eq!(groups[0].metrics.len(), 2); @@ -651,12 +690,12 @@ mod tests { }) ); - let mut dashboard_hidden = default_settings(&catalog, &HashSet::from(["codex".to_owned()])); + let mut dashboard_hidden = test_settings.clone(); dashboard_hidden.providers[0].metrics[0].enabled = false; let hidden_groups = resolved_groups(&state, &dashboard_hidden, &catalog); assert_eq!(hidden_groups[0].metrics[0].value, "75%"); - let mut used_settings = default_settings(&catalog, &HashSet::from(["codex".to_owned()])); + let mut used_settings = test_settings.clone(); used_settings.usage_display = crate::models::UsageDisplay::Used; let used_groups = resolved_groups(&state, &used_settings, &catalog); assert_eq!(