From ddca46e99d75fac7470d0712fa8b858cc1491256 Mon Sep 17 00:00:00 2001 From: Jonathan Liebig Date: Fri, 10 Jul 2026 21:41:12 +0200 Subject: [PATCH] Fix imported account refresh authority and recovery Make imported account profiles the sole rotating refresh-token authority, serialize cross-process auth transitions, retire terminal credentials safely, and resolve structural account markers before workspace enforcement. --- codex-rs/cli/src/account_cmd.rs | 6 +- codex-rs/cli/src/doctor.rs | 24 +- codex-rs/cli/src/login.rs | 14 +- codex-rs/cli/src/plugin_cmd.rs | 1 + codex-rs/core/tests/suite/client.rs | 1 + codex-rs/login/src/account.rs | 219 ++++++++++- codex-rs/login/src/account_lease.rs | 18 + codex-rs/login/src/account_tests.rs | 362 +++++++++++++++++- codex-rs/login/src/auth/auth_tests.rs | 65 ++++ codex-rs/login/src/auth/manager.rs | 290 +++++++++++--- .../auth/manager/imported_account_refresh.rs | 336 ++++++++++++++++ codex-rs/login/tests/suite/auth_refresh.rs | 358 ++++++++++++++++- codex-rs/login/tests/suite/logout.rs | 194 +++++++++- codex-rs/models-manager/src/manager_tests.rs | 1 + codex-rs/tui/src/account_usage.rs | 89 ++++- codex-rs/tui/src/account_usage_tests.rs | 11 + codex-rs/tui/src/lib.rs | 17 +- 17 files changed, 1883 insertions(+), 123 deletions(-) create mode 100644 codex-rs/login/src/auth/manager/imported_account_refresh.rs diff --git a/codex-rs/cli/src/account_cmd.rs b/codex-rs/cli/src/account_cmd.rs index 7ee4c9e26f77..301d5032daaa 100644 --- a/codex-rs/cli/src/account_cmd.rs +++ b/codex-rs/cli/src/account_cmd.rs @@ -95,11 +95,11 @@ fn print_accounts(accounts: Vec) { return; } - println!("ID Label Enabled Priority"); + println!("ID Label Enabled Login required Priority"); for account in accounts { println!( - "{} {} {} {}", - account.id, account.label, account.enabled, account.priority + "{} {} {} {} {}", + account.id, account.label, account.enabled, account.login_required, account.priority ); } } diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index fb577a2680bd..5a5c6fc1bdc7 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -370,13 +370,17 @@ async fn build_report( reachability_check, ) = tokio::join!( async { run_sync_check("config", progress.clone(), || config_check(config)) }, - async { run_sync_check("auth", progress.clone(), || auth_check(config)) }, + async { + run_sync_check("auth", progress.clone(), || { + auth_check(config, auth_manager.as_ref()) + }) + }, async { run_sync_check("updates", progress.clone(), || updates_check(config)) }, async { run_sync_check("network", progress.clone(), network_check) }, run_async_check( "websocket", progress.clone(), - websocket_reachability_check(config, Some(auth_manager)), + websocket_reachability_check(config, Some(Arc::clone(&auth_manager))), ), run_async_check("MCP", progress.clone(), mcp_check(config)), async { @@ -1172,7 +1176,7 @@ fn config_toml_details(config: &Config, details: &mut Vec) { } } -fn auth_check(config: &Config) -> DoctorCheck { +fn auth_check(config: &Config, auth_manager: &AuthManager) -> DoctorCheck { let mut details = Vec::new(); let auth_path = config.codex_home.join("auth.json"); details.push(format!( @@ -1218,7 +1222,19 @@ fn auth_check(config: &Config) -> DoctorCheck { "stored agent identity: {}", auth.agent_identity.is_some() )); - let auth_issues = stored_auth_issues(&auth, env_var_present); + let imported_auth_is_usable = auth_manager.active_account_id().is_some() + && auth_manager + .auth_cached() + .and_then(|auth| auth.get_token_data().ok()) + .is_some_and(|tokens| { + !tokens.access_token.trim().is_empty() + && !tokens.refresh_token.trim().is_empty() + }); + let auth_issues = if imported_auth_is_usable { + Vec::new() + } else { + stored_auth_issues(&auth, env_var_present) + }; details.extend( auth_issues .iter() diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 98d56bec12f9..257969c8d25d 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -10,6 +10,7 @@ use codex_config::types::AuthCredentialsStoreMode; use codex_core::config::Config; use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; use codex_login::AuthRouteConfig; use codex_login::CLIENT_ID; use codex_login::CodexAuth; @@ -428,6 +429,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { match CodexAuth::from_auth_storage( &config.codex_home, config.cli_auth_credentials_store_mode, + config.forced_chatgpt_workspace_id.as_deref(), Some(&config.chatgpt_base_url), config.auth_keyring_backend_kind(), auth_route_config.as_ref(), @@ -478,16 +480,10 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; - let auth_route_config = config.auth_route_config(); + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; - match logout_with_revoke( - &config.codex_home, - config.cli_auth_credentials_store_mode, - config.auth_keyring_backend_kind(), - auth_route_config.as_ref(), - ) - .await - { + match auth_manager.logout_with_revoke().await { Ok(true) => { eprintln!("Successfully logged out"); std::process::exit(0); diff --git a/codex-rs/cli/src/plugin_cmd.rs b/codex-rs/cli/src/plugin_cmd.rs index 056c256b394d..97e1984e0729 100644 --- a/codex-rs/cli/src/plugin_cmd.rs +++ b/codex-rs/cli/src/plugin_cmd.rs @@ -608,6 +608,7 @@ pub(crate) async fn load_cli_auth_mode(config: &Config) -> Option { CodexAuth::from_auth_storage( &config.codex_home, config.cli_auth_credentials_store_mode, + config.forced_chatgpt_workspace_id.as_deref(), Some(&config.chatgpt_base_url), config.auth_keyring_backend_kind(), auth_route_config.as_ref(), diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 4d298998afb0..d1f68ca8393a 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1554,6 +1554,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { let auth = CodexAuth::from_auth_storage( codex_home.path(), AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), /*auth_route_config*/ None, diff --git a/codex-rs/login/src/account.rs b/codex-rs/login/src/account.rs index a7ad0c0580df..cde1b7a65ac1 100644 --- a/codex-rs/login/src/account.rs +++ b/codex-rs/login/src/account.rs @@ -19,6 +19,7 @@ use crate::save_auth; const ACCOUNTS_DIR: &str = "accounts"; const INDEX_FILE: &str = "index.json"; +const INDEX_LOCK_FILE: &str = "index.lock"; #[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] #[serde(transparent)] @@ -43,6 +44,8 @@ pub struct AccountProfile { #[serde(default = "default_enabled")] pub enabled: bool, #[serde(default)] + pub login_required: bool, + #[serde(default)] pub priority: u32, #[serde(default, skip_serializing_if = "Option::is_none")] pub usage_limit_resets_at: Option, @@ -71,12 +74,12 @@ pub enum AccountAuthScope { File, } -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct AccountStore { codex_home: PathBuf, } -#[derive(Default, Deserialize, Serialize)] +#[derive(Clone, Default, Deserialize, Serialize)] struct AccountIndex { #[serde(default)] accounts: Vec, @@ -93,21 +96,63 @@ impl AccountStore { root_store_mode: AuthCredentialsStoreMode, root_keyring_backend_kind: AuthKeyringBackendKind, ) -> std::io::Result { - let auth = + let _root_refresh_guard = AccountLease::acquire_auth_refresh(&self.codex_home)?; + let root_auth = load_auth_dot_json(&self.codex_home, root_store_mode, root_keyring_backend_kind)? .ok_or_else(|| { std::io::Error::new(std::io::ErrorKind::NotFound, "not logged in") })?; + let mut auth = root_auth.clone(); + let imported_from_root_marker = is_root_account_marker(&auth); + let source_account_id = account_id_for_auth(&auth)?; + let account_home = self.account_home(&source_account_id); + let _account_refresh_guard = AccountLease::acquire_auth_refresh(&account_home)?; + if imported_from_root_marker { + auth = load_auth_dot_json( + &account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )? + .ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::NotFound, + "current imported account is missing auth.json", + ) + })?; + } if !is_managed_chatgpt_auth(&auth) { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "current auth is not a ChatGPT login", )); } + if auth + .tokens + .as_ref() + .is_none_or(|tokens| tokens.refresh_token.trim().is_empty()) + { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "current ChatGPT login is missing a refresh token", + )); + } let account_id = account_id_for_auth(&auth)?; + if account_id != source_account_id { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "root account marker does not match imported account auth", + )); + } let label = account_label_for_auth(&auth, label, &account_id)?; - let account_home = self.account_home(&account_id); + let _index_guard = self.acquire_index_lock()?; + let previous_account_auth = load_auth_dot_json( + &account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let mut index = self.load_index()?; + let previous_index = index.clone(); save_auth( &account_home, &auth, @@ -115,8 +160,7 @@ impl AccountStore { AuthKeyringBackendKind::default(), )?; - let mut index = self.load_index()?; - let auth = AccountAuthStorage { + let auth_storage = AccountAuthStorage { scope: AccountAuthScope::File, path: auth_path_for_id(&account_id), }; @@ -132,15 +176,56 @@ impl AccountStore { id: account_id, label, enabled: true, + login_required: imported_from_root_marker + && existing.is_some_and(|profile| profile.login_required), priority, usage_limit_resets_at, - auth, + auth: auth_storage, }; index.accounts.retain(|existing| existing.id != profile.id); index.accounts.push(profile.clone()); sort_profiles(&mut index.accounts); - self.save_index(&index)?; + if let Err(err) = self.save_index(&index) { + return match restore_file_auth(&account_home, previous_account_auth.as_ref()) { + Ok(()) => Err(err), + Err(rollback_err) => Err(std::io::Error::other(format!( + "failed to update imported account index: {err}; failed to restore account auth: {rollback_err}" + ))), + }; + } + if let Err(err) = save_root_account_marker( + &self.codex_home, + &auth, + root_store_mode, + root_keyring_backend_kind, + ) { + let mut rollback_errors = Vec::new(); + if let Err(rollback_err) = self.save_index(&previous_index) { + rollback_errors.push(format!("restore account index: {rollback_err}")); + } + if let Err(rollback_err) = + restore_file_auth(&account_home, previous_account_auth.as_ref()) + { + rollback_errors.push(format!("restore account auth: {rollback_err}")); + } + if let Err(rollback_err) = save_auth( + &self.codex_home, + &root_auth, + root_store_mode, + root_keyring_backend_kind, + ) { + rollback_errors.push(format!("restore root auth: {rollback_err}")); + } + return if rollback_errors.is_empty() { + Err(err) + } else { + Err(std::io::Error::other(format!( + "failed to save imported account marker: {err}; rollback failed: {}", + rollback_errors.join("; ") + ))) + }; + } Ok(profile) } @@ -176,6 +261,7 @@ impl AccountStore { account_id: &AccountId, resets_at: i64, ) -> std::io::Result { + let _index_guard = self.acquire_index_lock()?; let mut index = self.load_index()?; let Some(account) = index .accounts @@ -191,12 +277,58 @@ impl AccountStore { Ok(true) } + pub fn record_login_required(&self, account_id: &AccountId) -> std::io::Result { + let _index_guard = self.acquire_index_lock()?; + self.record_login_required_unlocked(account_id) + } + + pub fn record_login_required_if_auth_matches( + &self, + account_id: &AccountId, + expected_auth: &AuthDotJson, + ) -> std::io::Result { + let account_home = self.account_home(account_id); + let _refresh_guard = AccountLease::acquire_auth_refresh(&account_home)?; + let current_auth = load_auth_dot_json( + &account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + if current_auth.as_ref() != Some(expected_auth) { + return Ok(false); + } + let _index_guard = self.acquire_index_lock()?; + self.record_login_required_unlocked(account_id) + } + + fn record_login_required_unlocked(&self, account_id: &AccountId) -> std::io::Result { + let mut index = self.load_index()?; + let Some(account) = index + .accounts + .iter_mut() + .find(|account| &account.id == account_id) + else { + return Ok(false); + }; + if account.login_required { + return Ok(true); + } + + account.login_required = true; + self.save_index(&index)?; + Ok(true) + } + pub fn apply_imported_account_to_root_auth( &self, account_id: &AccountId, root_store_mode: AuthCredentialsStoreMode, root_keyring_backend_kind: AuthKeyringBackendKind, ) -> std::io::Result { + let account_home = self.account_home(account_id); + let _root_refresh_guard = AccountLease::acquire_auth_refresh(&self.codex_home)?; + let _account_refresh_guard = AccountLease::acquire_auth_refresh(&account_home)?; + let _index_guard = self.acquire_index_lock()?; let (profile, account_home) = self .enabled_file_account_profiles()? .into_iter() @@ -219,7 +351,7 @@ impl AccountStore { ) })?; - save_auth( + save_root_account_marker( &self.codex_home, &auth, root_store_mode, @@ -228,7 +360,7 @@ impl AccountStore { Ok(profile) } - pub(crate) fn disable_all(&self) -> std::io::Result { + pub(crate) fn disable_all_unlocked(&self) -> std::io::Result { let mut index = self.load_index()?; let mut changed = false; for account in &mut index.accounts { @@ -280,10 +412,17 @@ impl AccountStore { pub(crate) fn enabled_file_account_profiles( &self, ) -> std::io::Result> { + Ok(self + .file_account_profiles()? + .into_iter() + .filter(|(account, _)| account.enabled && !account.login_required) + .collect()) + } + + pub(crate) fn file_account_profiles(&self) -> std::io::Result> { Ok(self .list()? .into_iter() - .filter(|account| account.enabled) .filter_map(|account| match account.auth.scope { AccountAuthScope::File => { let home = self.account_home(&account.id); @@ -293,6 +432,23 @@ impl AccountStore { .collect()) } + pub(crate) fn file_auth_homes(&self) -> std::io::Result> { + let entries = match std::fs::read_dir(self.accounts_dir()) { + Ok(entries) => entries, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(err) => return Err(err), + }; + let mut auth_homes = Vec::new(); + for entry in entries { + let account_home = entry?.path(); + if account_home.join("auth.json").is_file() { + auth_homes.push(account_home); + } + } + auth_homes.sort(); + Ok(auth_homes) + } + pub(crate) fn account_home(&self, account_id: &AccountId) -> PathBuf { self.accounts_dir().join(account_id.as_str()) } @@ -305,6 +461,10 @@ impl AccountStore { self.accounts_dir().join(INDEX_FILE) } + pub(crate) fn acquire_index_lock(&self) -> std::io::Result { + AccountLease::acquire(&self.accounts_dir().join(INDEX_LOCK_FILE)) + } + fn load_index(&self) -> std::io::Result { let path = self.index_path(); let data = match std::fs::read_to_string(&path) { @@ -457,6 +617,43 @@ fn is_managed_chatgpt_auth(auth: &AuthDotJson) -> bool { } } +pub(crate) fn is_root_account_marker(auth: &AuthDotJson) -> bool { + is_managed_chatgpt_auth(auth) + && auth + .tokens + .as_ref() + .is_some_and(|tokens| tokens.refresh_token.is_empty()) +} + +fn save_root_account_marker( + codex_home: &Path, + auth: &AuthDotJson, + store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> std::io::Result<()> { + let mut marker = auth.clone(); + if let Some(tokens) = marker.tokens.as_mut() { + tokens.refresh_token.clear(); + } + save_auth(codex_home, &marker, store_mode, keyring_backend_kind) +} + +fn restore_file_auth(auth_home: &Path, auth: Option<&AuthDotJson>) -> std::io::Result<()> { + if let Some(auth) = auth { + return save_auth( + auth_home, + auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ); + } + match std::fs::remove_file(auth_home.join("auth.json")) { + Ok(()) => Ok(()), + Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(err) => Err(err), + } +} + fn auth_path_for_id(account_id: &AccountId) -> String { format!("{ACCOUNTS_DIR}/{account_id}/auth.json") } diff --git a/codex-rs/login/src/account_lease.rs b/codex-rs/login/src/account_lease.rs index 7172defbffd2..9619ad91da69 100644 --- a/codex-rs/login/src/account_lease.rs +++ b/codex-rs/login/src/account_lease.rs @@ -9,6 +9,24 @@ pub(crate) struct AccountLease { } impl AccountLease { + pub(crate) fn acquire_auth_refresh(auth_home: &Path) -> io::Result { + Self::acquire(&auth_home.join(".auth-refresh.lock")) + } + + pub(crate) fn acquire(path: &Path) -> io::Result { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path)?; + file.lock_exclusive()?; + Ok(Self { file }) + } + pub(crate) fn try_acquire(path: &Path) -> io::Result> { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; diff --git a/codex-rs/login/src/account_tests.rs b/codex-rs/login/src/account_tests.rs index d367de9b810f..eba4f23bf66a 100644 --- a/codex-rs/login/src/account_tests.rs +++ b/codex-rs/login/src/account_tests.rs @@ -1,5 +1,6 @@ use super::*; use crate::AuthManager; +use crate::CodexAuth; use crate::token_data::IdTokenInfo; use crate::token_data::TokenData; use base64::Engine; @@ -8,10 +9,12 @@ use codex_protocol::auth::AuthMode; use pretty_assertions::assert_eq; use serde::Serialize; use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Barrier; use tempfile::tempdir; #[test] -fn import_current_copies_chatgpt_auth_into_account_home() { +fn import_current_keeps_refresh_token_only_in_account_home() { let codex_home = tempdir().expect("tempdir"); let root_auth = test_auth("account-a", "user-a", "a@example.com"); save_auth( @@ -39,6 +42,17 @@ fn import_current_copies_chatgpt_auth_into_account_home() { ) .expect("load account auth"); assert_eq!(imported_auth, Some(root_auth)); + let root_auth = load_auth_dot_json( + codex_home.path(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("load root auth") + .expect("root account marker"); + assert_eq!( + root_auth.tokens.expect("root tokens").refresh_token, + String::new() + ); } #[test] @@ -69,8 +83,10 @@ fn import_current_reenables_existing_account_profile() { let codex_home = tempdir().expect("tempdir"); let store = AccountStore::new(codex_home.path().to_path_buf()); let first = import_test_account(&store, codex_home.path(), "first", "account-a"); - assert!(store.disable_all().expect("disable accounts")); - save_root_test_auth(codex_home.path(), "account-a"); + { + let _index_guard = store.acquire_index_lock().expect("lock account index"); + assert!(store.disable_all_unlocked().expect("disable accounts")); + } let profile = store .import_current( @@ -102,6 +118,119 @@ fn import_current_reenables_existing_account_profile() { ); } +#[test] +fn import_current_preserves_login_required_for_root_marker() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let imported = import_test_account(&store, codex_home.path(), "first", "account-a"); + assert!( + store + .record_login_required(&imported.id) + .expect("record login required") + ); + + let profile = store + .import_current( + Some("still stale".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("re-import root marker"); + + assert!(profile.login_required); + assert!( + store + .enabled_file_account_profiles() + .expect("eligible accounts") + .is_empty() + ); +} + +#[test] +fn concurrent_login_required_updates_are_preserved() { + let codex_home = tempdir().expect("tempdir"); + let profiles: Vec<_> = (0..16) + .map(|index| { + import_test_account( + &AccountStore::new(codex_home.path().to_path_buf()), + codex_home.path(), + &format!("account {index}"), + &format!("account-{index}"), + ) + }) + .collect(); + let barrier = Arc::new(Barrier::new(profiles.len())); + let handles: Vec<_> = profiles + .iter() + .map(|profile| { + let account_id = profile.id.clone(); + let codex_home = codex_home.path().to_path_buf(); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + AccountStore::new(codex_home).record_login_required(&account_id) + }) + }) + .collect(); + + for handle in handles { + assert!(handle.join().expect("update thread").expect("update")); + } + assert!( + AccountStore::new(codex_home.path().to_path_buf()) + .list() + .expect("accounts") + .iter() + .all(|profile| profile.login_required) + ); +} + +#[test] +fn login_requirement_is_not_recorded_after_account_auth_is_replaced() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let profile = import_test_account(&store, codex_home.path(), "first", "account-a"); + let account_home = store.account_home(&profile.id); + let attempted_auth = load_auth_dot_json( + &account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("load attempted auth") + .expect("attempted auth"); + let replacement_auth = test_auth("account-a", "replacement-user", "replacement@example.com"); + save_auth( + &account_home, + &replacement_auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("replace account auth"); + + assert!( + !store + .record_login_required_if_auth_matches(&profile.id, &attempted_auth) + .expect("compare account auth") + ); + assert!( + store + .list() + .expect("accounts") + .iter() + .find(|account| account.id == profile.id) + .is_some_and(|account| !account.login_required) + ); + assert_eq!( + load_auth_dot_json( + &account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("load replacement auth"), + Some(replacement_auth) + ); +} + #[test] fn candidates_include_usage_limit_blocked_state() { let codex_home = tempdir().expect("tempdir"); @@ -177,10 +306,9 @@ fn apply_imported_account_to_root_auth_switches_root_auth() { .expect("load root auth") .expect("root auth"); assert_eq!(selected, first); - assert_eq!( - root_auth.tokens.and_then(|tokens| tokens.account_id), - Some("account-a".to_string()) - ); + let tokens = root_auth.tokens.expect("root tokens"); + assert_eq!(tokens.account_id, Some("account-a".to_string())); + assert_eq!(tokens.refresh_token, String::new()); } #[tokio::test] @@ -195,6 +323,79 @@ async fn startup_prefers_imported_account_matching_root_chatgpt_auth() { assert_eq!(manager.active_account_id(), Some(second.id)); } +#[tokio::test] +async fn startup_does_not_use_root_account_marker_when_index_is_malformed() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let account = import_test_account(&store, codex_home.path(), "first", "account-a"); + std::fs::write(codex_home.path().join("accounts/index.json"), "{") + .expect("write malformed account index"); + + let manager = test_auth_manager(codex_home.path()).await; + + assert_eq!(manager.active_account_id(), None); + assert_eq!(manager.auth_cached(), None); + assert!(manager.logout().await.expect("logout")); + assert!(!codex_home.path().join("auth.json").exists()); + assert!(!store.account_home(&account.id).join("auth.json").exists()); +} + +#[tokio::test] +async fn live_manager_adopts_imported_account_when_root_becomes_a_marker() { + let codex_home = tempdir().expect("tempdir"); + save_root_test_auth(codex_home.path(), "account-a"); + let manager = test_auth_manager(codex_home.path()).await; + assert_eq!(manager.active_account_id(), None); + + let store = AccountStore::new(codex_home.path().to_path_buf()); + let imported = store + .import_current( + Some("first".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("import account"); + + assert!(manager.reload().await); + assert_eq!(manager.active_account_id(), Some(imported.id)); + assert_eq!( + manager.auth_cached().and_then(|auth| auth.get_account_id()), + Some("account-a".to_string()) + ); +} + +#[tokio::test] +async fn guarded_reload_does_not_commit_a_different_imported_account_source() { + let codex_home = tempdir().expect("tempdir"); + save_root_test_auth(codex_home.path(), "account-a"); + let live_manager = test_auth_manager(codex_home.path()).await; + let store = AccountStore::new(codex_home.path().to_path_buf()); + let account_a = import_test_account(&store, codex_home.path(), "first", "account-a"); + let _account_b = import_test_account(&store, codex_home.path(), "second", "account-b"); + store + .apply_imported_account_to_root_auth( + &account_a.id, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("select first account"); + let lease_owner = test_auth_manager(codex_home.path()).await; + assert_eq!(lease_owner.active_account_id(), Some(account_a.id)); + + live_manager + .refresh_token() + .await + .expect_err("alternate imported account should fail guarded reload"); + + assert_eq!(live_manager.active_account_id(), None); + assert_eq!( + live_manager + .auth_cached() + .and_then(|auth| auth.get_account_id()), + Some("account-a".to_string()) + ); +} + #[tokio::test] async fn startup_avoids_account_leased_by_another_manager() { let codex_home = tempdir().expect("tempdir"); @@ -257,6 +458,150 @@ async fn startup_keeps_root_api_key_auth_over_imported_accounts() { assert_eq!(manager.auth_mode(), Some(AuthMode::ApiKey)); } +#[tokio::test] +async fn startup_keeps_fresh_root_login_when_imported_profile_requires_login() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let imported = import_test_account(&store, codex_home.path(), "first", "account-a"); + store + .record_login_required(&imported.id) + .expect("record login requirement"); + let fresh_auth = test_auth("account-a", "user-account-a", "account-a@example.com"); + save_auth( + codex_home.path(), + &fresh_auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("save fresh root auth"); + + let manager = test_auth_manager(codex_home.path()).await; + + assert_eq!(manager.active_account_id(), None); + assert_eq!( + manager + .auth_cached() + .expect("fresh root auth") + .get_token_data() + .expect("fresh root tokens"), + fresh_auth.tokens.expect("expected root tokens") + ); +} + +#[tokio::test] +async fn startup_keeps_fresh_root_login_when_another_imported_account_is_available() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let _healthy = import_test_account(&store, codex_home.path(), "first", "account-a"); + let reauthenticated = import_test_account(&store, codex_home.path(), "second", "account-b"); + store + .record_login_required(&reauthenticated.id) + .expect("record login requirement"); + let fresh_auth = test_auth("account-b", "fresh-user-b", "fresh-b@example.com"); + save_auth( + codex_home.path(), + &fresh_auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("save fresh root auth"); + + let manager = test_auth_manager(codex_home.path()).await; + + assert_eq!(manager.active_account_id(), None); + assert_eq!( + manager + .auth_cached() + .expect("fresh root auth") + .get_token_data() + .expect("fresh root tokens"), + fresh_auth.tokens.expect("expected root tokens") + ); +} + +#[tokio::test] +async fn startup_honors_forced_workspace_over_fresh_root_reauthentication() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let forced = import_test_account(&store, codex_home.path(), "first", "account-a"); + let reauthenticated = import_test_account(&store, codex_home.path(), "second", "account-b"); + store + .record_login_required(&reauthenticated.id) + .expect("record login requirement"); + save_auth( + codex_home.path(), + &test_auth("account-b", "fresh-user-b", "fresh-b@example.com"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("save fresh root auth"); + + let manager = + test_auth_manager_with_forced_workspace(codex_home.path(), vec!["account-a"]).await; + + assert_eq!(manager.active_account_id(), Some(forced.id)); + assert_eq!( + manager + .auth_cached() + .expect("forced imported auth") + .get_account_id(), + Some("account-a".to_string()) + ); +} + +#[tokio::test] +async fn startup_rejects_disallowed_root_reauthentication_without_forced_account() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let reauthenticated = import_test_account(&store, codex_home.path(), "first", "account-a"); + store + .record_login_required(&reauthenticated.id) + .expect("record login requirement"); + save_auth( + codex_home.path(), + &test_auth("account-a", "fresh-user-a", "fresh-a@example.com"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("save fresh root auth"); + + let manager = + test_auth_manager_with_forced_workspace(codex_home.path(), vec!["account-b"]).await; + + assert_eq!(manager.active_account_id(), None); + assert_eq!(manager.auth_cached(), None); +} + +#[tokio::test] +async fn direct_marker_auth_resolution_honors_forced_workspace() { + let codex_home = tempdir().expect("tempdir"); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let first = import_test_account(&store, codex_home.path(), "first", "account-a"); + let _second = import_test_account(&store, codex_home.path(), "second", "account-b"); + store + .apply_imported_account_to_root_auth( + &first.id, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("select first account"); + let forced_workspaces = vec!["account-b".to_string()]; + + let auth = CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + Some(&forced_workspaces), + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await + .expect("load marker auth") + .expect("matching imported auth"); + + assert_eq!(auth.get_account_id(), Some("account-b".to_string())); +} + #[tokio::test] async fn switch_to_next_imported_account_skips_attempted_local_account_ids() { let codex_home = tempdir().expect("tempdir"); @@ -350,6 +695,9 @@ async fn logout_clears_imported_accounts_that_startup_would_select() { let store = AccountStore::new(codex_home.path().to_path_buf()); let first = import_test_account(&store, codex_home.path(), "first", "account-a"); let second = import_test_account(&store, codex_home.path(), "second", "account-b"); + store + .record_login_required(&second.id) + .expect("record login requirement"); save_root_test_auth(codex_home.path(), "account-a"); let manager = test_auth_manager(codex_home.path()).await; assert_eq!(manager.active_account_id(), Some(first.id.clone())); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 780cd8600984..cd9e6aa1d70d 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -788,6 +788,7 @@ async fn missing_auth_json_returns_none() { let auth = CodexAuth::from_auth_storage( dir.path(), AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), /*auth_route_config*/ None, @@ -1725,6 +1726,70 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { ); } +#[tokio::test] +#[serial(codex_auth_env)] +async fn enforce_login_restrictions_resolves_marker_to_allowed_imported_account() { + let codex_home = tempdir().unwrap(); + let _access_token_guard = remove_access_token_env_var(); + let store = AccountStore::new(codex_home.path().to_path_buf()); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_DISALLOWED.to_string()), + }, + codex_home.path(), + ) + .expect("seed disallowed auth"); + let disallowed = store + .import_current( + Some("disallowed".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("import disallowed account"); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), + }, + codex_home.path(), + ) + .expect("seed allowed auth"); + store + .import_current( + Some("allowed".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("import allowed account"); + store + .apply_imported_account_to_root_auth( + &disallowed.id, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("select disallowed marker"); + let candidates_before = store.candidates().expect("account candidates"); + let config = build_config( + codex_home.path(), + /*forced_login_method*/ None, + Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), + ) + .await; + + super::enforce_login_restrictions(&config) + .await + .expect("allowed imported account should satisfy workspace restriction"); + + assert!(codex_home.path().join("auth.json").exists()); + assert_eq!( + store.candidates().expect("account candidates"), + candidates_before + ); +} + #[tokio::test] #[serial(codex_auth_env)] async fn enforce_login_restrictions_logs_out_for_personal_access_token_workspace_mismatch() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 7108aff2f1fb..18b907bbfbde 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -44,6 +44,7 @@ use crate::account::AccountId; use crate::account::AccountProfile; use crate::account::AccountStore; use crate::account::account_id_for_auth; +use crate::account::is_root_account_marker; use crate::account_lease::AccountLease; use crate::auth::AuthHeaders; pub use crate::auth::agent_identity::AgentIdentityAuth; @@ -73,6 +74,8 @@ use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; +mod imported_account_refresh; + /// Authentication mechanism used by the current user. #[derive(Debug, Clone)] pub enum CodexAuth { @@ -250,13 +253,37 @@ impl From for std::io::Error { } } +#[derive(Error)] +#[error("{source}")] +pub struct RefreshAuthFromStorageError { + #[source] + source: std::io::Error, + attempted_auth: Option, +} + +impl Debug for RefreshAuthFromStorageError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RefreshAuthFromStorageError") + .field("source", &self.source) + .field("attempted_auth_available", &self.attempted_auth.is_some()) + .finish() + } +} + +impl RefreshAuthFromStorageError { + pub fn attempted_auth(&self) -> Option<&AuthDotJson> { + self.attempted_auth.as_ref() + } +} + pub async fn refresh_auth_from_storage( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option<&str>, keyring_backend_kind: AuthKeyringBackendKind, auth_route_config: Option<&AuthRouteConfig>, -) -> std::io::Result> { +) -> Result, RefreshAuthFromStorageError> { let manager = AuthManager::new( codex_home.to_path_buf(), /*enable_codex_api_key_env*/ false, @@ -267,7 +294,14 @@ pub async fn refresh_auth_from_storage( auth_route_config.cloned(), ) .await; - manager.refresh_token().await?; + if let Err(err) = manager.refresh_token().await { + return Err(RefreshAuthFromStorageError { + source: err.into(), + attempted_auth: manager + .auth_cached() + .and_then(|auth| auth.get_current_auth_json()), + }); + } Ok(manager.auth().await) } @@ -377,23 +411,37 @@ impl CodexAuth { pub async fn from_auth_storage( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: Option<&[String]>, chatgpt_base_url: Option<&str>, keyring_backend_kind: AuthKeyringBackendKind, auth_route_config: Option<&AuthRouteConfig>, ) -> std::io::Result> { let agent_identity_authapi_base_url = agent_identity_authapi_base_url(chatgpt_base_url).ok(); - load_auth( + let auth = load_auth( codex_home, /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, - /*forced_chatgpt_workspace_id*/ None, + forced_chatgpt_workspace_id, chatgpt_base_url, keyring_backend_kind, agent_identity_authapi_base_url.as_deref(), auth_route_config, ) + .await?; + if !imported_account_refresh::root_auth_is_account_marker(auth.as_ref()) { + return Ok(auth); + } + Ok(load_initial_imported_account_auth( + codex_home, + auth.as_ref(), + forced_chatgpt_workspace_id, + chatgpt_base_url, + agent_identity_authapi_base_url.as_deref(), + auth_route_config, + ) .await + .map(|(_account_id, _account_home, auth)| auth)) } pub async fn from_agent_identity_jwt( @@ -919,6 +967,7 @@ pub async fn logout_with_revoke( None } }; + let auth_dot_json = auth_dot_json.filter(|auth| !is_root_account_marker(auth)); if let Err(err) = revoke_auth_tokens(auth_dot_json.as_ref(), auth_route_config).await { tracing::warn!("failed to revoke auth tokens during logout: {err}"); } @@ -1091,7 +1140,7 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( config: &AuthConfig, agent_identity_authapi_base_url: Option<&str>, ) -> std::io::Result<()> { - let Some(auth) = load_auth( + let Some(root_auth) = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, @@ -1105,6 +1154,23 @@ async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( else { return Ok(()); }; + let auth = if imported_account_refresh::root_auth_is_account_marker(Some(&root_auth)) { + let Some((_, _, auth)) = load_initial_imported_account_auth( + &config.codex_home, + Some(&root_auth), + config.forced_chatgpt_workspace_id.as_deref(), + config.chatgpt_base_url.as_deref(), + agent_identity_authapi_base_url, + config.auth_route_config.as_ref(), + ) + .await + else { + return Ok(()); + }; + auth + } else { + root_auth + }; if let Some(required_method) = config.forced_login_method { let method_violation = match (required_method, auth.auth_mode()) { @@ -1593,6 +1659,20 @@ enum ReloadOutcome { Skipped, } +struct LoadedAuth { + auth: Option, + imported_source: Option<(AccountId, PathBuf)>, +} + +impl LoadedAuth { + fn from_current_source(auth: Option) -> Self { + Self { + auth, + imported_source: None, + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum UnauthorizedRecoveryMode { Managed, @@ -2006,16 +2086,30 @@ impl AuthManager { .await .ok() .flatten(); - let should_prefer_imported_accounts = match root_auth.as_ref() { - Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_)) | None => true, - Some( - CodexAuth::ApiKey(_) - | CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::PersonalAccessToken(_) - | CodexAuth::BedrockApiKey(_), - ) => false, + let root_auth_workspace_allowed = match root_auth.as_ref() { + Some(auth) if auth.is_chatgpt_auth() => { + chatgpt_auth_workspace_allowed(auth, forced_chatgpt_workspace_id.as_deref()) + } + Some(_) | None => true, }; + let root_reauthenticates_login_required_account = root_auth_workspace_allowed + && imported_account_refresh::root_auth_reauthenticates_login_required_account( + &codex_home, + root_auth.as_ref(), + ); + let should_prefer_imported_accounts = !root_reauthenticates_login_required_account + && match root_auth.as_ref() { + Some(CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_)) | None => true, + Some( + CodexAuth::ApiKey(_) + | CodexAuth::Headers(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::PersonalAccessToken(_) + | CodexAuth::BedrockApiKey(_), + ) => false, + }; + let root_auth_is_account_marker = + imported_account_refresh::root_auth_is_account_marker(root_auth.as_ref()); let (active_account_id, active_auth_home, managed_auth) = if should_prefer_imported_accounts { match load_initial_imported_account_auth( @@ -2031,7 +2125,13 @@ impl AuthManager { Some((account_id, account_home, auth)) => { (Some(account_id), account_home, Some(auth)) } - None => (None, codex_home.clone(), root_auth), + None => ( + None, + codex_home.clone(), + (!root_auth_is_account_marker && root_auth_workspace_allowed) + .then_some(root_auth) + .flatten(), + ), } } else { (None, codex_home.clone(), root_auth) @@ -2226,7 +2326,7 @@ impl AuthManager { && let Err(err) = self.refresh_token().await { tracing::error!("Failed to refresh token: {}", err); - return Some(auth); + return self.auth_cached(); } self.auth_cached() } @@ -2296,9 +2396,21 @@ impl AuthManager { /// Reloads auth from the active source. Returns whether the auth value changed. pub async fn reload(&self) -> bool { + let Ok(_refresh_guard) = self.refresh_lock.acquire().await else { + return false; + }; + self.reload_unlocked().await + } + + async fn reload_unlocked(&self) -> bool { tracing::info!("Reloading auth"); - let new_auth = self.load_auth().await; - self.set_cached_auth(new_auth) + let active_account_id_before_reload = self.active_account_id(); + let loaded_auth = self.load_auth().await; + if let Some((account_id, account_home)) = loaded_auth.imported_source { + self.set_active_imported_account_source(account_id, account_home); + } + let active_account_changed = active_account_id_before_reload != self.active_account_id(); + self.set_cached_auth(loaded_auth.auth) || active_account_changed } async fn reload_if_account_id_matches( @@ -2313,8 +2425,12 @@ impl AuthManager { } }; - let new_auth = self.load_auth().await; - let new_account_id = new_auth.as_ref().and_then(CodexAuth::get_account_id); + let active_account_id_before_reload = self.active_account_id(); + let loaded_auth = self.load_auth().await; + let new_account_id = loaded_auth + .auth + .as_ref() + .and_then(CodexAuth::get_account_id); if new_account_id.as_deref() != Some(expected_account_id) { let found_account_id = new_account_id.as_deref().unwrap_or("unknown"); @@ -2324,11 +2440,17 @@ impl AuthManager { return ReloadOutcome::Skipped; } + if let Some((account_id, account_home)) = loaded_auth.imported_source { + self.set_active_imported_account_source(account_id, account_home); + } tracing::info!("Reloading auth for account {expected_account_id}"); let cached_before_reload = self.auth_cached(); - let auth_changed = - !Self::auths_equal_for_refresh(cached_before_reload.as_ref(), new_auth.as_ref()); - self.set_cached_auth(new_auth); + let auth_changed = active_account_id_before_reload != self.active_account_id() + || !Self::auths_equal_for_refresh( + cached_before_reload.as_ref(), + loaded_auth.auth.as_ref(), + ); + self.set_cached_auth(loaded_auth.auth); if auth_changed { ReloadOutcome::ReloadedChanged } else { @@ -2387,20 +2509,21 @@ impl AuthManager { } } - async fn load_auth(&self) -> Option { + async fn load_auth(&self) -> LoadedAuth { if let Some(external_auth) = self.external_auth() { - return match self.resolve_external_auth(&external_auth).await { + let auth = match self.resolve_external_auth(&external_auth).await { Ok(auth) => Some(auth), Err(err) => { tracing::error!("Failed to resolve external auth: {err}"); None } }; + return LoadedAuth::from_current_source(auth); } let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); let auth_home = self.active_auth_home(); - load_auth( + let auth = load_auth( &auth_home, self.enable_codex_api_key_env, self.active_auth_credentials_store_mode(), @@ -2412,7 +2535,29 @@ impl AuthManager { ) .await .ok() - .flatten() + .flatten(); + if !imported_account_refresh::root_auth_is_account_marker(auth.as_ref()) { + return LoadedAuth::from_current_source(auth); + } + if self.active_account_id().is_some() { + return LoadedAuth::from_current_source(/*auth*/ None); + } + let imported_source = load_initial_imported_account_auth( + &self.codex_home, + auth.as_ref(), + forced_chatgpt_workspace_id.as_deref(), + self.chatgpt_base_url.as_deref(), + self.agent_identity_authapi_base_url.as_deref(), + self.auth_route_config.as_ref(), + ) + .await; + match imported_source { + Some((account_id, account_home, auth)) => LoadedAuth { + auth: Some(auth), + imported_source: Some((account_id, account_home)), + }, + None => LoadedAuth::from_current_source(/*auth*/ None), + } } fn set_cached_auth(&self, new_auth: Option) -> bool { @@ -2476,6 +2621,11 @@ impl AuthManager { } pub async fn activate_imported_account(&self, account_id: &AccountId) -> std::io::Result<()> { + let _refresh_guard = self + .refresh_lock + .acquire() + .await + .map_err(|_| std::io::Error::other("auth refresh lock is closed"))?; if self.active_account_id().as_ref() == Some(account_id) { return Ok(()); } @@ -2522,11 +2672,22 @@ impl AuthManager { pub async fn switch_to_next_imported_account( &self, attempted_account_ids: &HashSet, + ) -> bool { + let Ok(_refresh_guard) = self.refresh_lock.acquire().await else { + return false; + }; + self.switch_to_next_imported_account_unlocked(attempted_account_ids) + .await + } + + async fn switch_to_next_imported_account_unlocked( + &self, + attempted_account_ids: &HashSet, ) -> bool { let store = AccountStore::new(self.codex_home.clone()); let accounts = store.enabled_file_account_profiles().unwrap_or_default(); let active_account_id = self.active_account_id(); - if active_account_id.is_some() && accounts.len() < 2 { + if accounts.is_empty() { return false; } @@ -2583,6 +2744,11 @@ impl AuthManager { account_home: PathBuf, auth: CodexAuth, ) { + self.set_active_imported_account_source(account_id, account_home); + self.set_cached_auth(Some(auth)); + } + + fn set_active_imported_account_source(&self, account_id: AccountId, account_home: PathBuf) { let account_lease = AccountStore::new(self.codex_home.clone()) .try_acquire_lease(&account_id) .ok() @@ -2596,7 +2762,6 @@ impl AuthManager { if let Ok(mut active_auth_home) = self.active_auth_home.write() { *active_auth_home = account_home; } - self.set_cached_auth(Some(auth)); } pub async fn set_external_auth( @@ -2730,13 +2895,21 @@ impl AuthManager { REFRESH_TOKEN_UNKNOWN_MESSAGE.to_string(), )) })?; - let auth_before_reload = self.auth_cached(); - if auth_before_reload + if self + .auth_cached() .as_ref() .is_some_and(|auth| auth.is_api_key_auth() || auth.is_personal_access_token_auth()) { return Ok(()); } + let _file_guard = self.acquire_refresh_file_lock().await?; + if matches!( + self.reconcile_imported_account_refresh_readiness().await?, + imported_account_refresh::ImportedAccountRefreshReadiness::Recovered + ) { + return Ok(()); + } + let auth_before_reload = self.auth_cached(); let expected_account_id = auth_before_reload .as_ref() .and_then(CodexAuth::get_account_id); @@ -2749,7 +2922,12 @@ impl AuthManager { tracing::info!("Skipping token refresh because auth changed after guarded reload."); Ok(()) } - ReloadOutcome::ReloadedNoChange => self.refresh_token_from_authority_impl().await, + ReloadOutcome::ReloadedNoChange => { + let attempted_account_id = self.active_account_id(); + let result = self.refresh_token_from_authority_impl().await; + self.recover_terminal_imported_refresh(result, attempted_account_id) + .await + } ReloadOutcome::Skipped => { Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( RefreshTokenFailedReason::Other, @@ -2763,6 +2941,9 @@ impl AuthManager { /// it and update the shared cache. If the token refresh fails, returns the /// error to the caller. pub async fn refresh_token_from_authority(&self) -> Result<(), RefreshTokenError> { + if !self.has_external_auth() { + return self.refresh_token().await; + } let _refresh_guard = self.refresh_lock.acquire().await.map_err(|_| { RefreshTokenError::Permanent(RefreshTokenFailedError::new( RefreshTokenFailedReason::Other, @@ -2817,46 +2998,55 @@ impl AuthManager { /// reloads the in‑memory auth cache so callers immediately observe the /// unauthenticated state. pub async fn logout(&self) -> std::io::Result { - let removed = self.logout_all_managed_auth()?; + let _refresh_guard = self + .refresh_lock + .acquire() + .await + .map_err(|_| std::io::Error::other("auth refresh lock is closed"))?; + let auth_locks = self.acquire_managed_auth_refresh_locks().await?; + let removed = self.logout_all_managed_auth(&auth_locks)?; // Always reload to clear any cached auth (even if file absent). self.clear_external_auth(); self.clear_active_imported_account(); - self.reload().await; + self.reload_unlocked().await; Ok(removed) } pub async fn logout_with_revoke(&self) -> std::io::Result { - let auth_dot_json = self - .auth_cached() - .and_then(|auth| auth.get_current_auth_json()); - if let Err(err) = - revoke_auth_tokens(auth_dot_json.as_ref(), self.auth_route_config.as_ref()).await - { - tracing::warn!("failed to revoke auth tokens during logout: {err}"); - } - let result = self.logout_all_managed_auth()?; + let _refresh_guard = self + .refresh_lock + .acquire() + .await + .map_err(|_| std::io::Error::other("auth refresh lock is closed"))?; + let mut auth_locks = self.acquire_managed_auth_refresh_locks().await?; + auth_locks.release_index_lock(); + self.revoke_managed_auth(&auth_locks).await; + auth_locks.reacquire_index_lock().await?; + let result = self.logout_all_managed_auth(&auth_locks)?; // Always reload to clear any cached auth (even if file absent). self.clear_external_auth(); self.clear_active_imported_account(); - self.reload().await; + self.reload_unlocked().await; Ok(result) } - fn logout_all_managed_auth(&self) -> std::io::Result { + fn logout_all_managed_auth( + &self, + auth_locks: &imported_account_refresh::ManagedAuthRefreshLocks, + ) -> std::io::Result { let mut removed = logout_all_stores( &self.codex_home, self.auth_credentials_store_mode, self.keyring_backend_kind, )?; - let account_store = AccountStore::new(self.codex_home.clone()); - for (_account, account_home) in account_store.enabled_file_account_profiles()? { + for account_home in auth_locks.account_homes() { removed |= logout_all_stores( - &account_home, + account_home, AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), )?; } - removed |= account_store.disable_all()?; + removed |= auth_locks.disable_all()?; Ok(removed) } @@ -2990,7 +3180,7 @@ impl AuthManager { refresh_response.refresh_token, ) .map_err(RefreshTokenError::from)?; - self.reload().await; + self.reload_unlocked().await; Ok(()) } diff --git a/codex-rs/login/src/auth/manager/imported_account_refresh.rs b/codex-rs/login/src/auth/manager/imported_account_refresh.rs new file mode 100644 index 000000000000..83c796aa5a35 --- /dev/null +++ b/codex-rs/login/src/auth/manager/imported_account_refresh.rs @@ -0,0 +1,336 @@ +use std::collections::HashSet; +use std::path::Path; +use std::path::PathBuf; + +use codex_config::types::AuthCredentialsStoreMode; +use codex_protocol::auth::RefreshTokenFailedReason; + +use super::AuthManager; +use super::CodexAuth; +use super::RefreshTokenError; +use super::RefreshTokenFailedError; +use super::ReloadOutcome; +use super::load_auth_dot_json; +use super::revoke_auth_tokens; +use crate::account::AccountId; +use crate::account::AccountStore; +use crate::account::account_id_for_auth; +use crate::account::is_root_account_marker; +use crate::account_lease::AccountLease; +use crate::auth::storage::AuthDotJson; +use crate::auth::storage::AuthKeyringBackendKind; + +const IMPORTED_ACCOUNT_LOGIN_REQUIRED_MESSAGE: &str = + "This account needs you to sign in again. Run `codex account add` to continue."; + +pub(super) enum ImportedAccountRefreshReadiness { + Ready, + Recovered, +} + +pub(super) fn root_auth_is_account_marker(root_auth: Option<&CodexAuth>) -> bool { + root_auth + .and_then(CodexAuth::get_current_auth_json) + .is_some_and(|auth| is_root_account_marker(&auth)) +} + +pub(super) fn root_auth_reauthenticates_login_required_account( + codex_home: &Path, + root_auth: Option<&CodexAuth>, +) -> bool { + root_auth + .and_then(CodexAuth::get_current_auth_json) + .filter(|auth| !is_root_account_marker(auth)) + .and_then(|auth| account_id_for_auth(&auth).ok()) + .is_some_and(|root_account_id| { + AccountStore::new(codex_home.to_path_buf()) + .list() + .is_ok_and(|accounts| { + accounts + .iter() + .any(|account| account.id == root_account_id && account.login_required) + }) + }) +} + +pub(super) struct ManagedAuthRefreshLocks { + account_store: AccountStore, + account_homes: Vec, + index_readable: bool, + index_guard: Option, + _refresh_guards: Vec, +} + +impl ManagedAuthRefreshLocks { + pub(super) fn account_homes(&self) -> &[PathBuf] { + &self.account_homes + } + + pub(super) fn disable_all(&self) -> std::io::Result { + if self.index_guard.is_none() { + return Err(std::io::Error::other("account index lock is not held")); + } + if self.index_readable { + self.account_store.disable_all_unlocked() + } else { + Ok(false) + } + } + + pub(super) fn release_index_lock(&mut self) { + self.index_guard = None; + } + + pub(super) async fn reacquire_index_lock(&mut self) -> std::io::Result<()> { + let account_store = self.account_store.clone(); + self.index_guard = Some( + tokio::task::spawn_blocking(move || account_store.acquire_index_lock()) + .await + .map_err(std::io::Error::other)??, + ); + Ok(()) + } +} + +impl AuthManager { + pub(super) async fn acquire_refresh_file_lock( + &self, + ) -> Result, RefreshTokenError> { + if self.has_external_auth() { + return Ok(None); + } + let auth_home = self.active_auth_home(); + tokio::task::spawn_blocking(move || acquire_refresh_file_lock(&auth_home)) + .await + .map_err(|err| RefreshTokenError::Transient(std::io::Error::other(err)))? + .map(Some) + .map_err(RefreshTokenError::Transient) + } + + pub(super) async fn acquire_managed_auth_refresh_locks( + &self, + ) -> std::io::Result { + let codex_home = self.codex_home.clone(); + tokio::task::spawn_blocking(move || acquire_managed_auth_refresh_locks(&codex_home)) + .await + .map_err(std::io::Error::other)? + } + + pub(super) async fn revoke_managed_auth(&self, locks: &ManagedAuthRefreshLocks) { + let mut auth_snapshots = Vec::new(); + if let Some(auth) = self + .auth_cached() + .and_then(|auth| auth.get_current_auth_json()) + .filter(|auth| !is_root_account_marker(auth)) + { + auth_snapshots.push(auth); + } + if let Some(auth) = load_auth_snapshot( + &self.codex_home, + AuthCredentialsStoreMode::Ephemeral, + AuthKeyringBackendKind::default(), + ) + .filter(|auth| !is_root_account_marker(auth)) + { + auth_snapshots.push(auth); + } + if self.auth_credentials_store_mode != AuthCredentialsStoreMode::Ephemeral + && let Some(auth) = load_auth_snapshot( + &self.codex_home, + self.auth_credentials_store_mode, + self.keyring_backend_kind, + ) + .filter(|auth| !is_root_account_marker(auth)) + { + auth_snapshots.push(auth); + } + for account_home in locks.account_homes() { + if let Some(auth) = load_auth_snapshot( + account_home, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) { + auth_snapshots.push(auth); + } + } + + let mut revoked_tokens = HashSet::new(); + for auth in auth_snapshots { + let Some(token) = revocation_token(&auth) else { + continue; + }; + if !revoked_tokens.insert(token.to_string()) { + continue; + } + if let Err(err) = revoke_auth_tokens(Some(&auth), self.auth_route_config.as_ref()).await + { + tracing::warn!("failed to revoke auth tokens during logout: {err}"); + } + } + } + + pub(super) async fn recover_terminal_imported_refresh( + &self, + result: Result<(), RefreshTokenError>, + attempted_account_id: Option, + ) -> Result<(), RefreshTokenError> { + let terminal = matches!( + result + .as_ref() + .err() + .and_then(RefreshTokenError::failed_reason), + Some( + RefreshTokenFailedReason::Expired + | RefreshTokenFailedReason::Exhausted + | RefreshTokenFailedReason::Revoked + ) + ); + let Some(attempted_account_id) = terminal.then_some(attempted_account_id).flatten() else { + return result; + }; + if self.active_account_id().as_ref() == Some(&attempted_account_id) { + let expected_account_id = self + .auth_cached() + .as_ref() + .and_then(CodexAuth::get_account_id); + if matches!( + self.reload_if_account_id_matches(expected_account_id.as_deref()) + .await, + ReloadOutcome::ReloadedChanged + ) { + return Ok(()); + } + } + + AccountStore::new(self.codex_home.clone()) + .record_login_required(&attempted_account_id) + .map_err(RefreshTokenError::Transient)?; + if self.active_account_id().as_ref() == Some(&attempted_account_id) { + self.move_off_imported_account_requiring_login(attempted_account_id) + .await + } else { + Ok(()) + } + } + + pub(super) async fn reconcile_imported_account_refresh_readiness( + &self, + ) -> Result { + let Some(active_account_id) = self.active_account_id() else { + return Ok(ImportedAccountRefreshReadiness::Ready); + }; + let login_required = AccountStore::new(self.codex_home.clone()) + .list() + .map_err(RefreshTokenError::Transient)? + .into_iter() + .find(|account| account.id == active_account_id) + .is_some_and(|account| account.login_required); + if !login_required { + return Ok(ImportedAccountRefreshReadiness::Ready); + } + + self.move_off_imported_account_requiring_login(active_account_id) + .await?; + Ok(ImportedAccountRefreshReadiness::Recovered) + } + + async fn move_off_imported_account_requiring_login( + &self, + active_account_id: AccountId, + ) -> Result<(), RefreshTokenError> { + let attempted_account_ids = HashSet::from([active_account_id.to_string()]); + if self + .switch_to_next_imported_account_unlocked(&attempted_account_ids) + .await + { + tracing::info!(%active_account_id, "switched away from imported account that requires login"); + Ok(()) + } else { + self.clear_active_imported_account(); + self.set_cached_auth(/*new_auth*/ None); + Err(RefreshTokenError::Permanent(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Other, + IMPORTED_ACCOUNT_LOGIN_REQUIRED_MESSAGE.to_string(), + ))) + } + } +} + +fn acquire_refresh_file_lock(auth_home: &Path) -> std::io::Result { + AccountLease::acquire_auth_refresh(auth_home) +} + +fn acquire_managed_auth_refresh_locks( + codex_home: &Path, +) -> std::io::Result { + loop { + let account_store = AccountStore::new(codex_home.to_path_buf()); + let (account_homes, index_readable) = file_account_homes(&account_store)?; + let mut auth_homes = vec![codex_home.to_path_buf()]; + auth_homes.extend(account_homes.iter().cloned()); + auth_homes.sort(); + auth_homes.dedup(); + let refresh_guards = auth_homes + .iter() + .map(|auth_home| acquire_refresh_file_lock(auth_home)) + .collect::>>()?; + let index_guard = account_store.acquire_index_lock()?; + let (current_account_homes, current_index_readable) = file_account_homes(&account_store)?; + if current_account_homes == account_homes && current_index_readable == index_readable { + return Ok(ManagedAuthRefreshLocks { + account_store, + account_homes, + index_readable, + index_guard: Some(index_guard), + _refresh_guards: refresh_guards, + }); + } + } +} + +fn file_account_homes(account_store: &AccountStore) -> std::io::Result<(Vec, bool)> { + let (mut account_homes, index_readable) = match account_store.file_account_profiles() { + Ok(profiles) => ( + profiles + .into_iter() + .map(|(_account, account_home)| account_home) + .collect(), + true, + ), + Err(err) => { + tracing::warn!(%err, "account index is unreadable during logout; scanning account auth files"); + (account_store.file_auth_homes()?, false) + } + }; + account_homes.sort(); + account_homes.dedup(); + Ok((account_homes, index_readable)) +} + +fn load_auth_snapshot( + auth_home: &Path, + store_mode: AuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Option { + match load_auth_dot_json(auth_home, store_mode, keyring_backend_kind) { + Ok(auth) => auth, + Err(err) => { + tracing::warn!( + auth_home = %auth_home.display(), + "failed to load stored auth during logout: {err}" + ); + None + } + } +} + +fn revocation_token(auth: &AuthDotJson) -> Option<&str> { + let tokens = auth.tokens.as_ref()?; + if !tokens.refresh_token.is_empty() { + Some(tokens.refresh_token.as_str()) + } else if !tokens.access_token.is_empty() { + Some(tokens.access_token.as_str()) + } else { + None + } +} diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index e47b2978e8d6..9bd491b1ccd2 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -4,10 +4,13 @@ use base64::Engine; use chrono::Duration; use chrono::Utc; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AccountProfile; +use codex_login::AccountStore; use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::CLIENT_ID_OVERRIDE_ENV_VAR; +use codex_login::CodexAuth; use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use codex_login::RefreshTokenError; use codex_login::load_auth_dot_json; @@ -21,6 +24,7 @@ use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; use std::ffi::OsString; +use std::path::Path; use std::sync::Arc; use tempfile::TempDir; use wiremock::Mock; @@ -839,6 +843,300 @@ async fn refresh_token_does_not_retry_after_bad_request_reused_failure() -> Resu Ok(()) } +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn reused_refresh_token_fails_over_to_another_imported_account() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "refresh_token_reused" + } + }))) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let _env_guard = EnvGuard::set( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", server.uri()), + ); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let stale_auth = account_auth("stale-account", "stale-access", "stale-refresh"); + let stale_profile = import_account(&store, codex_home.path(), "stale", &stale_auth)?; + + let healthy_auth = account_auth("healthy-account", "healthy-access", "healthy-refresh"); + let healthy_profile = import_account(&store, codex_home.path(), "healthy", &healthy_auth)?; + + // Match the startup picker behavior that made the stale imported account active. + save_auth( + codex_home.path(), + &stale_auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let auth_manager = shared_auth_manager(codex_home.path()).await; + assert_eq!( + auth_manager.active_account_id(), + Some(stale_profile.id.clone()) + ); + + auth_manager + .refresh_token() + .await + .context("reused imported account should fail over")?; + + assert_eq!( + auth_manager.active_account_id(), + Some(healthy_profile.id.clone()) + ); + let profiles = store.list()?; + assert!( + profiles + .iter() + .find(|profile| profile.id == stale_profile.id) + .is_some_and(|profile| profile.login_required) + ); + assert!( + profiles + .iter() + .find(|profile| profile.id == healthy_profile.id) + .is_some_and(|profile| !profile.login_required) + ); + server.verify().await; + Ok(()) +} + +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn reused_refresh_token_without_fallback_requires_login_instead_of_retrying() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "refresh_token_reused" + } + }))) + .expect(2) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let _env_guard = EnvGuard::set( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", server.uri()), + ); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let stale_auth = account_auth("stale-account", "stale-access", "stale-refresh"); + let stale_profile = import_account(&store, codex_home.path(), "stale", &stale_auth)?; + + let auth_manager = shared_auth_manager(codex_home.path()).await; + assert_eq!( + auth_manager.active_account_id(), + Some(stale_profile.id.clone()) + ); + + let error = auth_manager + .refresh_token() + .await + .expect_err("terminal imported refresh should require login"); + + assert_eq!(auth_manager.active_account_id(), None); + assert_eq!(auth_manager.auth_cached(), None); + assert_eq!(error.failed_reason(), Some(RefreshTokenFailedReason::Other)); + assert_eq!( + error.to_string(), + "This account needs you to sign in again. Run `codex account add` to continue." + ); + assert!( + store + .list()? + .iter() + .find(|profile| profile.id == stale_profile.id) + .is_some_and(|profile| profile.login_required) + ); + assert!(!auth_manager.reload().await); + assert_eq!(auth_manager.auth_cached(), None); + assert_eq!( + CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await?, + None + ); + + let mut proactively_stale_auth = stale_auth.clone(); + proactively_stale_auth + .tokens + .as_mut() + .expect("stale tokens") + .access_token = access_token_with_expiration(Utc::now() + Duration::minutes(4)); + proactively_stale_auth.last_refresh = Some(Utc::now()); + save_auth( + codex_home.path(), + &proactively_stale_auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let reimported = store.import_current( + Some("stale".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + assert!(!reimported.login_required); + assert_eq!( + CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await? + .and_then(|auth| auth.get_account_id()), + Some("stale-account".to_string()) + ); + let proactive_manager = shared_auth_manager(codex_home.path()).await; + assert_eq!(proactive_manager.auth().await, None); + assert_eq!(proactive_manager.active_account_id(), None); + server.verify().await; + Ok(()) +} + +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn concurrent_auth_managers_refresh_a_rotating_token_once() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(std::time::Duration::from_millis(100)) + .set_body_json(json!({ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token" + })), + ) + .expect(1) + .mount(&server) + .await; + + let ctx = RefreshTokenTestContext::new(&server).await?; + let initial_auth = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN)), + last_refresh: Some(Utc::now() - Duration::days(1)), + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + }; + ctx.write_auth(&initial_auth).await?; + let second_manager = shared_auth_manager(ctx.codex_home.path()).await; + + let (first, second) = tokio::join!( + ctx.auth_manager.refresh_token(), + second_manager.refresh_token() + ); + first.context("first manager should refresh")?; + second.context("second manager should adopt the refreshed auth")?; + + let expected_tokens = TokenData { + access_token: "new-access-token".to_string(), + refresh_token: "new-refresh-token".to_string(), + ..build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN) + }; + assert_eq!( + ctx.auth_manager + .auth_cached() + .context("first manager should cache auth")? + .get_token_data()?, + expected_tokens + ); + assert_eq!( + second_manager + .auth_cached() + .context("second manager should cache auth")? + .get_token_data()?, + expected_tokens + ); + server.verify().await; + Ok(()) +} + +#[tokio::test] +async fn refresh_token_does_not_create_a_lock_for_api_key_auth() -> Result<()> { + let temp = TempDir::new()?; + let missing_home = temp.path().join("missing"); + let auth_manager = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("sk-test"), + missing_home.clone(), + ); + + auth_manager.refresh_token().await?; + + assert!(!missing_home.exists()); + Ok(()) +} + +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn concurrent_managers_attempt_a_terminal_imported_refresh_once() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": { + "code": "refresh_token_reused" + } + }))) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let _env_guard = EnvGuard::set( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", server.uri()), + ); + let store = AccountStore::new(codex_home.path().to_path_buf()); + let stale_auth = account_auth("stale-account", "stale-access", "stale-refresh"); + let stale_profile = import_account(&store, codex_home.path(), "stale", &stale_auth)?; + let first = shared_auth_manager(codex_home.path()).await; + let second = shared_auth_manager(codex_home.path()).await; + + let (first_result, second_result) = tokio::join!(first.refresh_token(), second.refresh_token()); + + assert!(first_result.is_err()); + assert!(second_result.is_err()); + assert!( + store + .list()? + .iter() + .find(|profile| profile.id == stale_profile.id) + .is_some_and(|profile| profile.login_required) + ); + server.verify().await; + Ok(()) +} + #[serial_test::serial(auth_env)] #[tokio::test] async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result<()> { @@ -1213,16 +1511,7 @@ impl RefreshTokenTestContext { let endpoint = format!("{}/oauth/token", server.uri()); let env_guard = EnvGuard::set(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, endpoint); - let auth_manager = AuthManager::shared( - codex_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*forced_chatgpt_workspace_id*/ None, - /*chatgpt_base_url*/ None, - AuthKeyringBackendKind::default(), - /*auth_route_config*/ None, - ) - .await; + let auth_manager = shared_auth_manager(codex_home.path()).await; Ok(Self { codex_home, @@ -1253,6 +1542,19 @@ impl RefreshTokenTestContext { } } +async fn shared_auth_manager(codex_home: &Path) -> Arc { + AuthManager::shared( + codex_home.to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await +} + struct EnvGuard { key: &'static str, original: Option, @@ -1325,3 +1627,39 @@ fn build_tokens(access_token: &str, refresh_token: &str) -> TokenData { account_id: Some("account-id".to_string()), } } + +fn account_auth(account_id: &str, access_token: &str, refresh_token: &str) -> AuthDotJson { + AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + account_id: Some(account_id.to_string()), + ..build_tokens(access_token, refresh_token) + }), + last_refresh: Some(Utc::now() - Duration::days(1)), + agent_identity: None, + personal_access_token: None, + bedrock_api_key: None, + } +} + +fn import_account( + store: &AccountStore, + codex_home: &std::path::Path, + label: &str, + auth: &AuthDotJson, +) -> Result { + save_auth( + codex_home, + auth, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + store + .import_current( + Some(label.to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .context("import account") +} diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index bf2873779b32..8c86f3b3ae39 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -2,6 +2,7 @@ use anyhow::Context; use anyhow::Result; use base64::Engine; use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AccountStore; use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; @@ -130,6 +131,58 @@ async fn logout_with_revoke_uses_stored_auth_when_access_token_env_is_set() -> R Ok(()) } +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn logout_with_revoke_does_not_revoke_imported_account_marker() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let _env_guard = EnvGuard::set( + REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/revoke", server.uri()), + ); + let codex_home = TempDir::new()?; + save_auth( + codex_home.path(), + &chatgpt_auth(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let profile = AccountStore::new(codex_home.path().to_path_buf()).import_current( + Some("imported".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + assert!( + logout_with_revoke( + codex_home.path(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await? + ); + + assert!(!codex_home.path().join("auth.json").exists()); + assert!( + codex_home + .path() + .join("accounts") + .join(profile.id.as_str()) + .join("auth.json") + .exists() + ); + server.verify().await; + Ok(()) +} + #[serial_test::serial(auth_env)] #[tokio::test] async fn logout_with_revoke_removes_auth_when_revoke_fails() -> Result<()> { @@ -176,7 +229,7 @@ async fn logout_with_revoke_removes_auth_when_revoke_fails() -> Result<()> { #[serial_test::serial(auth_env)] #[tokio::test] -async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> { +async fn auth_manager_logout_with_revoke_uses_cached_and_stored_auth() -> Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; @@ -185,7 +238,7 @@ async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> { .respond_with(ResponseTemplate::new(200).set_body_json(json!({ "message": "success" }))) - .expect(1) + .expect(2) .mount(&server) .await; let _env_guard = EnvGuard::set( @@ -227,16 +280,123 @@ async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> { .received_requests() .await .context("failed to fetch revoke requests")?; - assert_eq!(requests.len(), 1); + let mut bodies = requests + .iter() + .map(wiremock::Request::body_json::) + .collect::, _>>()?; + bodies.sort_by(|left, right| left["token"].as_str().cmp(&right["token"].as_str())); assert_eq!( - requests[0] - .body_json::() - .context("revoke request should be JSON")?, - json!({ - "token": REFRESH_TOKEN, - "token_type_hint": "refresh_token", - "client_id": CLIENT_ID, - }) + bodies, + vec![ + json!({ + "token": "newer-disk-refresh-token", + "token_type_hint": "refresh_token", + "client_id": CLIENT_ID, + }), + json!({ + "token": REFRESH_TOKEN, + "token_type_hint": "refresh_token", + "client_id": CLIENT_ID, + }), + ] + ); + server.verify().await; + Ok(()) +} + +#[serial_test::serial(auth_env)] +#[tokio::test] +async fn auth_manager_logout_with_revoke_revokes_all_imported_accounts() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&server) + .await; + let _env_guard = EnvGuard::set( + REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/revoke", server.uri()), + ); + + let codex_home = TempDir::new()?; + let store = AccountStore::new(codex_home.path().to_path_buf()); + save_auth( + codex_home.path(), + &chatgpt_auth_for_account("account-a", "access-a", "refresh-a"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let account_a = store.import_current( + Some("account a".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + save_auth( + codex_home.path(), + &chatgpt_auth_for_account("account-b", "access-b", "refresh-b"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let account_b = store.import_current( + Some("account b".to_string()), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + let manager = AuthManager::new( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + + assert!(manager.logout_with_revoke().await?); + + let requests = server + .received_requests() + .await + .context("failed to fetch revoke requests")?; + let mut bodies = requests + .iter() + .map(wiremock::Request::body_json::) + .collect::, _>>()?; + bodies.sort_by(|left, right| left["token"].as_str().cmp(&right["token"].as_str())); + assert_eq!( + bodies, + vec![ + json!({ + "token": "refresh-a", + "token_type_hint": "refresh_token", + "client_id": CLIENT_ID, + }), + json!({ + "token": "refresh-b", + "token_type_hint": "refresh_token", + "client_id": CLIENT_ID, + }), + ] + ); + assert!( + !codex_home + .path() + .join("accounts") + .join(account_a.id.as_str()) + .join("auth.json") + .exists() + ); + assert!( + !codex_home + .path() + .join("accounts") + .join(account_b.id.as_str()) + .join("auth.json") + .exists() ); server.verify().await; Ok(()) @@ -247,6 +407,14 @@ fn chatgpt_auth() -> AuthDotJson { } fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { + chatgpt_auth_for_account("account-id", ACCESS_TOKEN, refresh_token) +} + +fn chatgpt_auth_for_account( + account_id: &str, + access_token: &str, + refresh_token: &str, +) -> AuthDotJson { AuthDotJson { auth_mode: Some(AuthMode::Chatgpt), openai_api_key: None, @@ -255,9 +423,9 @@ fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { raw_jwt: minimal_jwt(), ..Default::default() }, - access_token: ACCESS_TOKEN.to_string(), + access_token: access_token.to_string(), refresh_token: refresh_token.to_string(), - account_id: Some("account-id".to_string()), + account_id: Some(account_id.to_string()), }), last_refresh: None, agent_identity: None, diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index e74c7cac4cdf..ac6fda1df8a4 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -240,6 +240,7 @@ c2ln", CodexAuth::from_auth_storage( codex_home, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), /*auth_route_config*/ None, diff --git a/codex-rs/tui/src/account_usage.rs b/codex-rs/tui/src/account_usage.rs index 11ce8b184489..a2ea88c4a478 100644 --- a/codex-rs/tui/src/account_usage.rs +++ b/codex-rs/tui/src/account_usage.rs @@ -8,9 +8,12 @@ use codex_backend_client::RequestError; use codex_login::AccountId; use codex_login::AccountStore; use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::CodexAuth; use codex_login::refresh_auth_from_storage; +use codex_protocol::auth::RefreshTokenFailedError; +use codex_protocol::auth::RefreshTokenFailedReason; use codex_protocol::protocol::RateLimitSnapshot; use std::collections::HashMap; use std::path::PathBuf; @@ -42,11 +45,31 @@ impl AccountUsage { } } +#[derive(Default)] +pub(crate) struct AccountUsageLoad { + pub(crate) usage: HashMap, + pub(crate) login_required: HashMap, +} + +struct AccountUsageFetchError { + error: anyhow::Error, + attempted_auth: Option, +} + +impl AccountUsageFetchError { + fn new(error: impl Into) -> Self { + Self { + error: error.into(), + attempted_auth: None, + } + } +} + pub(crate) async fn load( config: &Config, accounts: &[(AccountId, PathBuf)], store: &AccountStore, -) -> HashMap { +) -> AccountUsageLoad { let mut tasks = JoinSet::new(); for (account_id, account_home) in accounts { let account_id = account_id.clone(); @@ -59,13 +82,13 @@ pub(crate) async fn load( fetch(account_home, chatgpt_base_url, auth_route_config), ) .await - .map_err(|_| anyhow!("rate-limit request timed out")) + .map_err(|_| AccountUsageFetchError::new(anyhow!("rate-limit request timed out"))) .and_then(std::convert::identity); (account_id, result) }); } - let mut usage = HashMap::new(); + let mut loaded = AccountUsageLoad::default(); while let Some(result) = tasks.join_next().await { match result { Ok((account_id, Ok(account_usage))) => { @@ -78,46 +101,69 @@ pub(crate) async fn load( "failed to persist imported account usage limit reset" ); } - usage.insert(account_id, account_usage); + loaded.usage.insert(account_id, account_usage); } Ok((account_id, Err(err))) => { - tracing::warn!(%account_id, %err, "failed to load imported account usage"); + if login_required(&err.error) + && let Some(attempted_auth) = err.attempted_auth + { + loaded + .login_required + .insert(account_id.clone(), attempted_auth); + } + tracing::warn!(%account_id, %err.error, "failed to load imported account usage"); } Err(err) => tracing::warn!(%err, "imported account usage task failed"), } } - usage + loaded } async fn fetch( account_home: PathBuf, chatgpt_base_url: String, auth_route_config: Option, -) -> Result { +) -> std::result::Result { let auth = CodexAuth::from_auth_storage( &account_home, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, Some(&chatgpt_base_url), AuthKeyringBackendKind::default(), auth_route_config.as_ref(), ) - .await? - .context("imported account is not authenticated")?; + .await + .map_err(AccountUsageFetchError::new)? + .context("imported account is not authenticated") + .map_err(AccountUsageFetchError::new)?; match fetch_with_auth(&auth, &chatgpt_base_url).await { Err(err) if is_unauthorized(&err) => { - let auth = refresh_auth_from_storage( + let auth = match refresh_auth_from_storage( &account_home, AuthCredentialsStoreMode::File, Some(&chatgpt_base_url), AuthKeyringBackendKind::default(), auth_route_config.as_ref(), ) - .await? - .context("imported account is not authenticated")?; - fetch_with_auth(&auth, &chatgpt_base_url).await + .await + { + Ok(auth) => auth, + Err(err) => { + let attempted_auth = err.attempted_auth().cloned(); + return Err(AccountUsageFetchError { + error: err.into(), + attempted_auth, + }); + } + } + .context("imported account is not authenticated") + .map_err(AccountUsageFetchError::new)?; + fetch_with_auth(&auth, &chatgpt_base_url) + .await + .map_err(AccountUsageFetchError::new) } - result => result, + result => result.map_err(AccountUsageFetchError::new), } } @@ -136,6 +182,21 @@ fn is_unauthorized(err: &anyhow::Error) -> bool { }) } +fn login_required(err: &anyhow::Error) -> bool { + err.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(|error| { + matches!( + error.reason, + RefreshTokenFailedReason::Expired + | RefreshTokenFailedReason::Exhausted + | RefreshTokenFailedReason::Revoked + ) + }) + }) +} + fn account_usage(response: &RateLimitsWithResetCredits) -> AccountUsage { response .rate_limits diff --git a/codex-rs/tui/src/account_usage_tests.rs b/codex-rs/tui/src/account_usage_tests.rs index 13ae09a93433..65b49230ced4 100644 --- a/codex-rs/tui/src/account_usage_tests.rs +++ b/codex-rs/tui/src/account_usage_tests.rs @@ -1,4 +1,5 @@ use super::*; +use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::protocol::RateLimitWindow; #[test] @@ -87,3 +88,13 @@ fn rounded_zero_remaining_does_not_mark_window_exhausted() { } ); } + +#[test] +fn reused_refresh_token_marks_imported_account_as_login_required() { + let error = anyhow!(RefreshTokenFailedError::new( + RefreshTokenFailedReason::Exhausted, + "refresh token already used", + )); + + assert!(login_required(&error)); +} diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 70a5b7a00559..6ef6ce8f7f5c 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -72,6 +72,7 @@ pub use session_archive_commands::SessionArchiveAction; pub use session_archive_commands::SessionArchiveCommandOptions; pub use session_archive_commands::run_session_archive_command; use std::collections::HashMap; +use std::collections::HashSet; use std::fs::OpenOptions; use std::path::Path; use std::path::PathBuf; @@ -634,7 +635,7 @@ async fn maybe_run_startup_account_picker( .collect::>(); let selectable_homes: HashMap = selectable_accounts.iter().cloned().collect(); - let candidates: Vec = store + let mut candidates: Vec = store .candidates()? .into_iter() .filter(|candidate| candidate.enabled && selectable_homes.contains_key(&candidate.id)) @@ -646,6 +647,18 @@ async fn maybe_run_startup_account_picker( } let usage = account_usage::load(config, &selectable_accounts, &store).await; + let mut login_required = HashSet::new(); + for (account_id, attempted_auth) in &usage.login_required { + if store.record_login_required_if_auth_matches(account_id, attempted_auth)? { + login_required.insert(account_id.clone()); + } + } + candidates.retain(|candidate| !login_required.contains(&candidate.id)); + if candidates.is_empty() { + return Ok(StartupAccountSelection::Continue { + selected_account_id: None, + }); + } let current_account_id = store .current_root_account_id( config.cli_auth_credentials_store_mode, @@ -658,7 +671,7 @@ async fn maybe_run_startup_account_picker( .map(|candidate| { account_picker_candidate( candidate, - usage.get(&candidate.id), + usage.usage.get(&candidate.id), store.account_in_use(&candidate.id).unwrap_or(false), current_account_id.as_ref() == Some(&candidate.id), )