From 2d682e7a89d0709dc36b212a8de6486a6929549d Mon Sep 17 00:00:00 2001 From: shiny-code-bot Date: Wed, 17 Jun 2026 16:20:24 -0400 Subject: [PATCH] Harden auth profile switching --- codex-rs/login/src/auth_profiles.rs | 61 +++++++++++- codex-rs/tui/src/app/app_server_events.rs | 29 +++++- codex-rs/tui/src/app/session_lifecycle.rs | 116 ++++++++++++++++++---- codex-rs/tui/src/app_server_session.rs | 8 +- 4 files changed, 188 insertions(+), 26 deletions(-) diff --git a/codex-rs/login/src/auth_profiles.rs b/codex-rs/login/src/auth_profiles.rs index 1e093a73be6a..f062622b95b3 100644 --- a/codex-rs/login/src/auth_profiles.rs +++ b/codex-rs/login/src/auth_profiles.rs @@ -261,8 +261,15 @@ pub fn record_auth_profile_login( upsert_auth_profile(codex_home, profile_name, |metadata| { metadata.last_login_at = Some(now); metadata.last_used_at = Some(now); - metadata.account_id = account_id; - metadata.email = email; + if let Some(account_id) = account_id { + if metadata.account_id.as_deref() != Some(account_id.as_str()) && email.is_none() { + metadata.email = None; + } + metadata.account_id = Some(account_id); + } + if let Some(email) = email { + metadata.email = Some(email); + } }) } @@ -353,6 +360,56 @@ mod tests { ); } + #[test] + fn record_login_preserves_existing_identity_when_new_values_are_missing() { + let temp = TempDir::new().expect("tempdir"); + + record_auth_profile_login( + temp.path(), + "work", + Some("account-1".to_string()), + Some("work@example.com".to_string()), + ) + .expect("record initial login"); + + let updated = record_auth_profile_login( + temp.path(), + "work", + /*account_id*/ None, + /*email*/ None, + ) + .expect("record login without metadata"); + + assert_eq!(updated.metadata.account_id.as_deref(), Some("account-1")); + assert_eq!(updated.metadata.email.as_deref(), Some("work@example.com")); + assert!(updated.metadata.last_login_at.is_some()); + assert!(updated.metadata.last_used_at.is_some()); + } + + #[test] + fn record_login_clears_email_when_account_changes_without_new_email() { + let temp = TempDir::new().expect("tempdir"); + + record_auth_profile_login( + temp.path(), + "work", + Some("account-1".to_string()), + Some("work@example.com".to_string()), + ) + .expect("record initial login"); + + let updated = record_auth_profile_login( + temp.path(), + "work", + Some("account-2".to_string()), + /*email*/ None, + ) + .expect("record changed account login"); + + assert_eq!(updated.metadata.account_id.as_deref(), Some("account-2")); + assert!(updated.metadata.email.is_none()); + } + #[cfg(unix)] #[test] fn saved_profiles_metadata_is_private() { diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index 22d00be984d1..275731acefb2 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -257,12 +257,37 @@ impl App { let Some(pending) = self.pending_auth_profile_login.take() else { return; }; + let auth_home = + match codex_login::profile_home(&self.config.codex_home, &pending.profile_name) { + Ok(auth_home) => auth_home, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Logged in, but failed to resolve auth profile {}: {err}", + pending.profile_label + )); + return; + } + }; + let (account_id, email) = match codex_login::load_auth_dot_json( + &auth_home, + self.config.cli_auth_credentials_store_mode, + ) { + Ok(Some(auth)) => (auth.account_id(), auth.account_email()), + Ok(None) => (None, None), + Err(err) => { + self.chat_widget.add_error_message(format!( + "Logged in, but failed to read auth profile {}: {err}", + pending.profile_label + )); + return; + } + }; if let Err(err) = codex_login::record_auth_profile_login( &self.config.codex_home, &pending.profile_name, - /*account_id*/ None, - /*email*/ None, + account_id, + email, ) { self.chat_widget.add_error_message(format!( "Logged in, but failed to update auth profile {}: {err}", diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index 8d7166f2c23d..bb6015f4252a 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -9,8 +9,36 @@ use crate::app_event::AuthProfileSelection; use codex_app_server_protocol::ClientRequest; 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; impl App { + async fn config_for_auth_profile_switch( + &self, + auth_home: AbsolutePathBuf, + cloud_config_bundle: CloudConfigBundleLoader, + ) -> std::io::Result { + ConfigBuilder::default() + .cli_overrides(self.cli_kv_overrides.clone()) + .harness_overrides(ConfigOverrides { + cwd: Some(self.config.cwd.to_path_buf()), + approval_policy: Some(*self.config.permissions.approval_policy.get()), + codex_self_exe: self.config.codex_self_exe.clone(), + codex_linux_sandbox_exe: self.config.codex_linux_sandbox_exe.clone(), + main_execve_wrapper_exe: self.config.main_execve_wrapper_exe.clone(), + show_raw_agent_reasoning: Some(self.config.show_raw_agent_reasoning), + workspace_roots: Some(self.config.workspace_roots.clone()), + ..Default::default() + }) + .loader_overrides(self.loader_overrides.clone()) + .strict_config(self.strict_config) + .cloud_config_bundle(cloud_config_bundle) + .auth_home(auth_home.to_path_buf()) + .build() + .await + .map_err(std::io::Error::other) + } + pub(super) async fn open_agent_picker(&mut self, app_server: &mut AppServerSession) { let mut thread_ids = self.agent_navigation.tracked_thread_ids(); for thread_id in self.thread_event_channels.keys().copied() { @@ -769,8 +797,7 @@ impl App { return; } - let mut config = self.fresh_session_config(); - config.auth_home = match AbsolutePathBuf::from_absolute_path(auth_home) { + let auth_home = match AbsolutePathBuf::from_absolute_path(auth_home) { Ok(auth_home) => auth_home, Err(err) => { self.chat_widget.add_error_message(format!( @@ -779,6 +806,29 @@ impl App { return; } }; + let replacement_cloud_config_bundle = cloud_config_bundle_loader_for_storage( + self.config.codex_home.to_path_buf(), + auth_home.to_path_buf(), + /*enable_codex_api_key_env*/ false, + self.config.cli_auth_credentials_store_mode, + self.config.chatgpt_base_url.clone(), + ) + .await; + let config = match self + .config_for_auth_profile_switch( + auth_home.clone(), + replacement_cloud_config_bundle.clone(), + ) + .await + { + Ok(config) => config, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to load config for auth profile {profile_label}: {err}" + )); + return; + } + }; let replacement_client = match crate::start_embedded_app_server( self.arg0_paths.clone(), @@ -786,7 +836,7 @@ impl App { self.cli_kv_overrides.clone(), self.loader_overrides.clone(), self.strict_config, - self.cloud_config_bundle.clone(), + replacement_cloud_config_bundle.clone(), self.feedback.clone(), /*log_db*/ None, self.state_db.clone(), @@ -803,6 +853,27 @@ impl App { } }; + let mut replacement_session = + AppServerSession::new(replacement_client, app_server.thread_params_mode()) + .with_remote_cwd_override(app_server.remote_cwd_override().map(PathBuf::from)); + let started = match replacement_session + .start_thread_with_session_start_source(&config, Some(ThreadStartSource::Clear)) + .await + { + Ok(started) => started, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to start auth profile session for {profile_label}: {err}" + )); + if let Err(shutdown_err) = replacement_session.shutdown().await { + tracing::warn!( + "failed to shut down replacement app server after thread start failure: {shutdown_err}" + ); + } + return; + } + }; + self.shutdown_current_thread(app_server).await; let tracked_thread_ids: Vec = self.thread_event_channels.keys().copied().collect(); @@ -812,26 +883,31 @@ impl App { } } - if let Err(err) = app_server.replace_client(replacement_client).await { - self.chat_widget.add_error_message(format!( - "Failed to switch app server to auth profile {profile_label}: {err}" - )); - return; + let old_client = app_server.swap_client(replacement_session.into_client()); + self.config = config.clone(); + self.cloud_config_bundle = replacement_cloud_config_bundle; + if let Err(err) = old_client.shutdown().await { + tracing::warn!( + "failed to shut down previous app server after auth profile switch: {err}" + ); } - let switched = self - .start_fresh_session_with_config( - tui, - app_server, - config, - Some(ThreadStartSource::Clear), - /*initial_user_message*/ None, - /*cleanup_current_thread*/ false, - "To continue the previous session, run ", - "Failed to attach to auth profile session", - "Failed to start auth profile session", + let switched = match self + .replace_chat_widget_with_app_server_thread( + tui, app_server, started, /*initial_user_message*/ None, ) - .await; + .await + { + Ok(()) => { + tui.frame_requester().schedule_frame(); + true + } + Err(err) => { + self.chat_widget + .add_error_message(format!("Failed to attach to auth profile session: {err}")); + false + } + }; if switched { self.chat_widget.add_info_message( format!("Using auth profile {profile_label} for this session."), diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index a2cf4e9cf5e8..ae1e85319cc0 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -1124,13 +1124,17 @@ impl AppServerSession { self.client.shutdown().await } - pub(crate) async fn replace_client(&mut self, client: AppServerClient) -> std::io::Result<()> { + pub(crate) fn into_client(self) -> AppServerClient { + self.client + } + + pub(crate) fn swap_client(&mut self, client: AppServerClient) -> AppServerClient { let old_client = std::mem::replace(&mut self.client, client); self.next_request_id = 1; self.thread_settings_update_supported = true; self.default_model = None; self.available_models.clear(); - old_client.shutdown().await + old_client } pub(crate) fn request_handle(&self) -> AppServerRequestHandle {