From e3cfa83ee7fd6197047209087f7bd76940c9078c Mon Sep 17 00:00:00 2001 From: Chris Busillo Date: Sun, 21 Jun 2026 18:35:27 -0400 Subject: [PATCH] Make stored account activation transactional --- codex-rs/login/src/auth/manager.rs | 11 ++ codex-rs/login/src/auth_accounts.rs | 153 ++++++++++++++++------ codex-rs/login/src/auth_accounts_tests.rs | 125 ++++++++++++++++++ codex-rs/login/src/lib.rs | 3 + codex-rs/tui/src/app/session_lifecycle.rs | 151 ++++++++++++++++++++- 5 files changed, 400 insertions(+), 43 deletions(-) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 4eebf5f3e48c..525edda4e545 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -583,6 +583,17 @@ pub fn logout( Ok(removed) } +/// Delete persisted auth without mirroring the removal into the stored-account +/// catalog. Rollback paths use this to restore a prior "no active auth" state +/// without treating the operation as a user-initiated logout. +pub fn delete_auth( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result { + let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); + storage.delete() +} + pub async fn logout_with_revoke( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, diff --git a/codex-rs/login/src/auth_accounts.rs b/codex-rs/login/src/auth_accounts.rs index 10507d981b39..4592dbbd7fa0 100644 --- a/codex-rs/login/src/auth_accounts.rs +++ b/codex-rs/login/src/auth_accounts.rs @@ -374,6 +374,50 @@ pub fn find_account(codex_home: &Path, account_id: &str) -> io::Result io::Result<(StoredAccount, AuthDotJson)> { + let account = find_account(codex_home, account_id)? + .ok_or_else(|| io::Error::other(format!("account with id {account_id} was not found")))?; + + let auth = + match account.mode { + AuthMode::ApiKey => AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some(account.openai_api_key.clone().ok_or_else(|| { + io::Error::other("stored API key account is missing the key value") + })?), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + }, + AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => AuthDotJson { + auth_mode: Some(account.mode), + openai_api_key: None, + tokens: Some(account.tokens.clone().ok_or_else(|| { + io::Error::other("stored ChatGPT account is missing token data") + })?), + last_refresh: account.last_refresh, + agent_identity: None, + personal_access_token: None, + }, + AuthMode::AgentIdentity => { + return Err(io::Error::other( + "stored agent identity account activation is not supported", + )); + } + AuthMode::PersonalAccessToken => { + return Err(io::Error::other( + "stored personal access token account activation is not supported", + )); + } + }; + + Ok((account, auth)) +} + pub fn update_account_last_refresh( codex_home: &Path, account_id: &str, @@ -419,49 +463,78 @@ pub fn activate_account( account_id: &str, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> io::Result { - let account = find_account(codex_home, account_id)? - .ok_or_else(|| io::Error::other(format!("account with id {account_id} was not found")))?; + let (_account, auth) = auth_for_account(codex_home, account_id)?; + commit_active_account(codex_home, account_id, &auth, auth_credentials_store_mode) +} - let auth = - match account.mode { - AuthMode::ApiKey => AuthDotJson { - auth_mode: Some(AuthMode::ApiKey), - openai_api_key: Some(account.openai_api_key.clone().ok_or_else(|| { - io::Error::other("stored API key account is missing the key value") - })?), - tokens: None, - last_refresh: None, - agent_identity: None, - personal_access_token: None, - }, - AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens => AuthDotJson { - auth_mode: Some(account.mode), - openai_api_key: None, - tokens: Some(account.tokens.clone().ok_or_else(|| { - io::Error::other("stored ChatGPT account is missing token data") - })?), - last_refresh: account.last_refresh, - agent_identity: None, - personal_access_token: None, - }, - AuthMode::AgentIdentity => { - return Err(io::Error::other( - "stored agent identity account activation is not supported", - )); +fn restore_previous_auth( + codex_home: &Path, + previous_auth: Option, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> io::Result<()> { + if let Some(previous_auth) = previous_auth { + save_auth(codex_home, &previous_auth, auth_credentials_store_mode) + } else { + crate::delete_auth(codex_home, auth_credentials_store_mode).map(|_| ()) + } +} + +fn restore_previous_activation( + codex_home: &Path, + previous_auth: Option, + previous_active_account_id: Option, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> io::Result<()> { + let auth_result = restore_previous_auth(codex_home, previous_auth, auth_credentials_store_mode); + let active_account_result = set_active_account_id(codex_home, previous_active_account_id); + + auth_result?; + active_account_result?; + Ok(()) +} + +pub fn commit_active_account( + codex_home: &Path, + account_id: &str, + auth: &AuthDotJson, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> io::Result { + let previous_auth = crate::load_auth_dot_json(codex_home, auth_credentials_store_mode)?; + let previous_active_account_id = get_active_account_id(codex_home)?; + + save_auth(codex_home, auth, auth_credentials_store_mode)?; + match set_active_account_id(codex_home, Some(account_id.to_string())) { + Ok(Some(activated)) => Ok(activated), + Ok(None) => { + let rollback_result = restore_previous_activation( + codex_home, + previous_auth, + previous_active_account_id, + auth_credentials_store_mode, + ); + if let Err(rollback_err) = rollback_result { + tracing::warn!( + "failed to roll back missing stored account activation: {rollback_err}" + ); } - AuthMode::PersonalAccessToken => { - return Err(io::Error::other( - "stored personal access token account activation is not supported", - )); + Err(io::Error::other(format!( + "account with id {account_id} disappeared before activation" + ))) + } + Err(err) => { + if let Err(rollback_err) = restore_previous_activation( + codex_home, + previous_auth, + previous_active_account_id, + auth_credentials_store_mode, + ) { + tracing::warn!( + "failed to roll back stored account activation error: {rollback_err}" + ); } - }; - - save_auth(codex_home, &auth, auth_credentials_store_mode)?; - set_active_account_id(codex_home, Some(account.id))?.ok_or_else(|| { - io::Error::other(format!( - "account with id {account_id} disappeared before activation" - )) - }) + Err(err) + } + } } pub fn remove_account(codex_home: &Path, account_id: &str) -> io::Result> { diff --git a/codex-rs/login/src/auth_accounts_tests.rs b/codex-rs/login/src/auth_accounts_tests.rs index 09abcd949e79..621cb002986e 100644 --- a/codex-rs/login/src/auth_accounts_tests.rs +++ b/codex-rs/login/src/auth_accounts_tests.rs @@ -298,6 +298,131 @@ fn activate_api_key_account_writes_auth_and_marks_active() { ); } +#[test] +fn auth_for_account_returns_auth_without_persisting_activation() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_api_key_account( + temp.path(), + "sk-test".to_string(), + Some("Work".to_string()), + /*make_active*/ false, + ) + .expect("upsert api key"); + + let (account, auth) = auth_for_account(temp.path(), &stored.id).expect("account auth"); + + assert_eq!(stored, account); + assert_eq!( + crate::AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-test".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + }, + auth + ); + assert_eq!(None, get_active_account_id(temp.path()).expect("active id")); + assert_eq!( + None, + crate::load_auth_dot_json(temp.path(), AuthCredentialsStoreMode::File) + .expect("read auth json") + ); +} + +#[test] +fn commit_active_account_rolls_back_auth_and_active_id_when_account_is_missing() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_api_key_account( + temp.path(), + "sk-previous".to_string(), + Some("Previous".to_string()), + /*make_active*/ true, + ) + .expect("upsert previous api key"); + let previous_auth = crate::AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-previous".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + }; + crate::save_auth(temp.path(), &previous_auth, AuthCredentialsStoreMode::File) + .expect("save previous auth"); + let new_auth = crate::AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-new".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + }; + + let err = commit_active_account( + temp.path(), + "missing", + &new_auth, + AuthCredentialsStoreMode::File, + ) + .expect_err("missing account should fail"); + + assert_eq!(io::ErrorKind::Other, err.kind()); + assert_eq!( + Some(stored.id.as_str()), + get_active_account_id(temp.path()) + .expect("active id") + .as_deref() + ); + assert_eq!( + previous_auth, + crate::load_auth_dot_json(temp.path(), AuthCredentialsStoreMode::File) + .expect("read auth json") + .expect("auth json should exist") + ); +} + +#[test] +fn commit_active_account_rollback_without_previous_auth_preserves_stored_accounts() { + let temp = TempDir::new().expect("tempdir"); + let stored = upsert_api_key_account( + temp.path(), + "sk-new".to_string(), + Some("New".to_string()), + /*make_active*/ false, + ) + .expect("upsert api key"); + let auth = crate::AuthDotJson { + auth_mode: Some(AuthMode::ApiKey), + openai_api_key: Some("sk-new".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + personal_access_token: None, + }; + + let err = commit_active_account( + temp.path(), + "missing", + &auth, + AuthCredentialsStoreMode::File, + ) + .expect_err("missing account should fail"); + + assert_eq!(io::ErrorKind::Other, err.kind()); + assert_eq!(None, get_active_account_id(temp.path()).expect("active id")); + assert_eq!( + None, + crate::load_auth_dot_json(temp.path(), AuthCredentialsStoreMode::File) + .expect("read auth json") + ); + assert_eq!( + vec![stored], + list_accounts(temp.path()).expect("list accounts") + ); +} + #[test] fn activate_chatgpt_account_writes_auth_and_marks_active() { let temp = TempDir::new().expect("tempdir"); diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 165d5215df1a..3d680f7d9914 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -39,6 +39,7 @@ pub use auth::REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR; pub use auth::RefreshTokenError; pub use auth::UnauthorizedRecovery; pub use auth::default_client; +pub use auth::delete_auth; pub use auth::enforce_login_restrictions; pub use auth::load_auth_dot_json; pub use auth::login_with_access_token; @@ -56,6 +57,8 @@ pub use auth_account_import::SkippedAuthAccountImport; pub use auth_account_import::import_auth_accounts_from_auth_homes; pub use auth_accounts::StoredAccount; pub use auth_accounts::activate_account; +pub use auth_accounts::auth_for_account; +pub use auth_accounts::commit_active_account; pub use auth_accounts::find_account; pub use auth_accounts::get_active_account_id; pub use auth_accounts::list_accounts; diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index 2a3d5dfa5feb..e76f6acfa116 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -12,6 +12,69 @@ use codex_app_server_protocol::LoginAccountParams; use codex_app_server_protocol::LoginAccountResponse; use codex_cloud_config::cloud_config_bundle_loader_for_storage; use codex_config::CloudConfigBundleLoader; +use codex_config::types::AuthCredentialsStoreMode; + +struct AuthSwitchRollback { + codex_home: PathBuf, + previous_auth: Option, + previous_active_account_id: Option, + auth_credentials_store_mode: AuthCredentialsStoreMode, + armed: bool, +} + +impl AuthSwitchRollback { + fn new( + codex_home: PathBuf, + previous_auth: Option, + previous_active_account_id: Option, + auth_credentials_store_mode: AuthCredentialsStoreMode, + ) -> Self { + Self { + codex_home, + previous_auth, + previous_active_account_id, + auth_credentials_store_mode, + armed: true, + } + } + + fn disarm(mut self) { + self.armed = false; + } + + fn restore_now(&mut self) -> std::io::Result<()> { + self.armed = false; + self.restore_auth() + } + + fn restore_auth(&self) -> std::io::Result<()> { + if let Some(previous_auth) = &self.previous_auth { + codex_login::save_auth( + &self.codex_home, + previous_auth, + self.auth_credentials_store_mode, + )?; + } else { + codex_login::delete_auth(&self.codex_home, self.auth_credentials_store_mode) + .map(|_| ())?; + } + codex_login::set_active_account_id( + &self.codex_home, + self.previous_active_account_id.clone(), + )?; + Ok(()) + } +} + +impl Drop for AuthSwitchRollback { + fn drop(&mut self) { + if self.armed + && let Err(err) = self.restore_auth() + { + tracing::warn!("failed to roll back stored account auth switch: {err}"); + } + } +} impl App { async fn config_for_auth_profile_switch( @@ -941,17 +1004,58 @@ impl App { return; } - if let Err(err) = codex_login::activate_account( + let (_account, selected_auth) = + match codex_login::auth_for_account(&self.config.codex_home, &selection.account_id) { + Ok(auth) => auth, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to load stored account {}: {err}", + selection.label + )); + return; + } + }; + let previous_auth = match codex_login::load_auth_dot_json( &self.config.codex_home, - &selection.account_id, + self.config.cli_auth_credentials_store_mode, + ) { + Ok(previous_auth) => previous_auth, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to read current auth before activating stored account {}: {err}", + selection.label + )); + return; + } + }; + let previous_active_account_id = + match codex_login::get_active_account_id(&self.config.codex_home) { + Ok(previous_active_account_id) => previous_active_account_id, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to read active account before activating stored account {}: {err}", + selection.label + )); + return; + } + }; + if let Err(err) = codex_login::save_auth( + &self.config.codex_home, + &selected_auth, self.config.cli_auth_credentials_store_mode, ) { self.chat_widget.add_error_message(format!( - "Failed to activate stored account {}: {err}", + "Failed to prepare stored account {}: {err}", selection.label )); return; } + let mut rollback = AuthSwitchRollback::new( + self.config.codex_home.to_path_buf(), + previous_auth, + previous_active_account_id, + self.config.cli_auth_credentials_store_mode, + ); let default_auth_home = self.config.codex_home.clone(); let replacement_cloud_config_bundle = cloud_config_bundle_loader_for_storage( @@ -1025,6 +1129,47 @@ impl App { } }; + match codex_login::set_active_account_id( + &self.config.codex_home, + Some(selection.account_id.clone()), + ) { + Ok(Some(_activated)) => rollback.disarm(), + Ok(None) => { + if let Err(rollback_err) = rollback.restore_now() { + tracing::warn!( + "failed to restore auth after missing stored account activation: {rollback_err}" + ); + } + self.chat_widget.add_error_message(format!( + "Failed to activate stored account {}: account disappeared before activation", + selection.label + )); + if let Err(shutdown_err) = replacement_session.shutdown().await { + tracing::warn!( + "failed to shut down replacement app server after account activation failure: {shutdown_err}" + ); + } + return; + } + Err(err) => { + if let Err(rollback_err) = rollback.restore_now() { + tracing::warn!( + "failed to restore auth after stored account activation error: {rollback_err}" + ); + } + self.chat_widget.add_error_message(format!( + "Failed to activate stored account {}: {err}", + selection.label + )); + if let Err(shutdown_err) = replacement_session.shutdown().await { + tracing::warn!( + "failed to shut down replacement app server after account activation failure: {shutdown_err}" + ); + } + return; + } + } + self.shutdown_current_thread(app_server).await; let tracked_thread_ids: Vec = self.thread_event_channels.keys().copied().collect();