From c0922e59309e4cdf26659c9bed4cc007985a5151 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 24 Mar 2026 12:48:36 -0600 Subject: [PATCH] Fix TUI app-server login browser launch and shutdown cleanup --- .../app-server/src/codex_message_processor.rs | 7 ++ codex-rs/app-server/src/in_process.rs | 1 + codex-rs/app-server/src/message_processor.rs | 4 + codex-rs/tui/src/lib.rs | 51 ++++++-- .../tui_app_server/src/app_server_session.rs | 32 +++-- codex-rs/tui_app_server/src/lib.rs | 49 ++++++-- .../tui_app_server/src/onboarding/auth.rs | 117 +++++++++++++----- .../onboarding/auth/headless_chatgpt_login.rs | 2 + .../src/onboarding/onboarding_screen.rs | 9 ++ 9 files changed, 213 insertions(+), 59 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c4c2799ca8..2153bcf62f 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1911,6 +1911,13 @@ impl CodexMessageProcessor { } } + pub(crate) async fn cancel_active_login(&self) { + let mut guard = self.active_login.lock().await; + if let Some(active_login) = guard.take() { + drop(active_login); + } + } + pub(crate) async fn clear_all_thread_listeners(&self) { self.thread_state_manager.clear_all_listeners().await; } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 0826f33c67..4acc419c5b 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -457,6 +457,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { } processor.clear_runtime_references(); + processor.cancel_active_login().await; processor.connection_closed(IN_PROCESS_CONNECTION_ID).await; processor.clear_all_thread_listeners().await; processor.drain_background_tasks().await; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 4f53bfc1a4..96fe437c05 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -451,6 +451,10 @@ impl MessageProcessor { self.codex_message_processor.drain_background_tasks().await; } + pub(crate) async fn cancel_active_login(&self) { + self.codex_message_processor.cancel_active_login().await; + } + pub(crate) async fn clear_all_thread_listeners(&self) { self.codex_message_processor .clear_all_thread_listeners() diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 5416200691..039148a718 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -609,6 +609,7 @@ async fn run_ratatui_app( terminal.clear()?; let mut tui = Tui::new(terminal); + let mut terminal_restore_guard = TerminalRestoreGuard::new(); #[cfg(not(debug_assertions))] { @@ -619,7 +620,7 @@ async fn run_ratatui_app( match update_prompt::run_update_prompt_if_needed(&mut tui, &initial_config).await? { UpdatePromptOutcome::Continue => {} UpdatePromptOutcome::RunUpdate(action) => { - crate::tui::restore()?; + terminal_restore_guard.restore()?; return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), thread_id: None, @@ -660,7 +661,7 @@ async fn run_ratatui_app( ) .await?; if onboarding_result.should_exit { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); return Ok(AppExitInfo { @@ -701,7 +702,7 @@ async fn run_ratatui_app( let mut missing_session_exit = |id_str: &str, action: &str| { error!("Error finding conversation path: {id_str}"); - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); Ok(AppExitInfo { @@ -773,7 +774,7 @@ async fn run_ratatui_app( error!( "Error reading session metadata from latest rollout: {rollout_path}" ); - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); return Ok(AppExitInfo { @@ -795,7 +796,7 @@ async fn run_ratatui_app( } else if cli.fork_picker { match resume_picker::run_fork_picker(&mut tui, &config, cli.fork_show_all).await? { resume_picker::SessionSelection::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -867,7 +868,7 @@ async fn run_ratatui_app( error!( "Error reading session metadata from latest rollout: {rollout_path}" ); - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); return Ok(AppExitInfo { @@ -887,7 +888,7 @@ async fn run_ratatui_app( } else if cli.resume_picker { match resume_picker::run_resume_picker(&mut tui, &config, cli.resume_show_all).await? { resume_picker::SessionSelection::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -929,7 +930,7 @@ async fn run_ratatui_app( { ResolveCwdOutcome::Continue(cwd) => cwd, ResolveCwdOutcome::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -1003,7 +1004,7 @@ async fn run_ratatui_app( ) .await; - restore(); + terminal_restore_guard.restore_silently(); // Mark the end of the recorded session. session_log::log_session_end(); // ignore error when collecting usage – report underlying error instead @@ -1128,6 +1129,38 @@ fn restore() { } } +struct TerminalRestoreGuard { + active: bool, +} + +impl TerminalRestoreGuard { + fn new() -> Self { + Self { active: true } + } + + #[cfg_attr(debug_assertions, allow(dead_code))] + fn restore(&mut self) -> color_eyre::Result<()> { + if self.active { + crate::tui::restore()?; + self.active = false; + } + Ok(()) + } + + fn restore_silently(&mut self) { + if self.active { + restore(); + self.active = false; + } + } +} + +impl Drop for TerminalRestoreGuard { + fn drop(&mut self) { + self.restore_silently(); + } +} + /// Determine whether to use the terminal's alternate screen buffer. /// /// The alternate screen buffer provides a cleaner fullscreen experience without polluting diff --git a/codex-rs/tui_app_server/src/app_server_session.rs b/codex-rs/tui_app_server/src/app_server_session.rs index 0da85bc617..ecdea588ee 100644 --- a/codex-rs/tui_app_server/src/app_server_session.rs +++ b/codex-rs/tui_app_server/src/app_server_session.rs @@ -80,6 +80,7 @@ use color_eyre::eyre::Result; use color_eyre::eyre::WrapErr; use std::collections::HashMap; use std::path::PathBuf; +use tracing::debug; use crate::bottom_pane::FeedbackAudience; use crate::status::StatusAccountDisplay; @@ -178,16 +179,6 @@ impl AppServerSession { }) .await .wrap_err("model/list failed during TUI bootstrap")?; - let rate_limit_request_id = self.next_request_id(); - let rate_limits: GetAccountRateLimitsResponse = self - .client - .request_typed(ClientRequest::GetAccountRateLimits { - request_id: rate_limit_request_id, - params: None, - }) - .await - .wrap_err("account/rateLimits/read failed during TUI bootstrap")?; - let available_models = models .data .into_iter() @@ -252,6 +243,25 @@ impl AppServerSession { false, ), }; + let rate_limit_snapshots = if account.requires_openai_auth && has_chatgpt_account { + let rate_limit_request_id = self.next_request_id(); + match self + .client + .request_typed(ClientRequest::GetAccountRateLimits { + request_id: rate_limit_request_id, + params: None, + }) + .await + { + Ok(rate_limits) => app_server_rate_limit_snapshots_to_core(rate_limits), + Err(error) => { + debug!(error = ?error, "failed to fetch rate limits during TUI bootstrap"); + Vec::new() + } + } + } else { + Vec::new() + }; Ok(AppServerBootstrap { account_auth_mode, @@ -263,7 +273,7 @@ impl AppServerSession { feedback_audience, has_chatgpt_account, available_models, - rate_limit_snapshots: app_server_rate_limit_snapshots_to_core(rate_limits), + rate_limit_snapshots, }) } diff --git a/codex-rs/tui_app_server/src/lib.rs b/codex-rs/tui_app_server/src/lib.rs index 17e309d5fb..790a7c6f55 100644 --- a/codex-rs/tui_app_server/src/lib.rs +++ b/codex-rs/tui_app_server/src/lib.rs @@ -938,6 +938,7 @@ async fn run_ratatui_app( terminal.clear()?; let mut tui = Tui::new(terminal); + let mut terminal_restore_guard = TerminalRestoreGuard::new(); #[cfg(not(debug_assertions))] { @@ -948,7 +949,7 @@ async fn run_ratatui_app( match update_prompt::run_update_prompt_if_needed(&mut tui, &initial_config).await? { UpdatePromptOutcome::Continue => {} UpdatePromptOutcome::RunUpdate(action) => { - crate::tui::restore()?; + terminal_restore_guard.restore()?; return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), thread_id: None, @@ -1016,7 +1017,7 @@ async fn run_ratatui_app( ) .await?; if onboarding_result.should_exit { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); return Ok(AppExitInfo { @@ -1062,7 +1063,7 @@ async fn run_ratatui_app( let mut missing_session_exit = |id_str: &str, action: &str| { error!("Error finding conversation path: {id_str}"); - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); let _ = tui.terminal.clear(); Ok(AppExitInfo { @@ -1137,7 +1138,7 @@ async fn run_ratatui_app( .await? { resume_picker::SessionSelection::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -1189,7 +1190,7 @@ async fn run_ratatui_app( .await? { resume_picker::SessionSelection::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -1235,7 +1236,7 @@ async fn run_ratatui_app( { ResolveCwdOutcome::Continue(cwd) => cwd, ResolveCwdOutcome::Exit => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Ok(AppExitInfo { token_usage: codex_protocol::protocol::TokenUsage::default(), @@ -1303,7 +1304,7 @@ async fn run_ratatui_app( { Ok(app_server) => app_server, Err(err) => { - restore(); + terminal_restore_guard.restore_silently(); session_log::log_session_end(); return Err(err); } @@ -1326,7 +1327,7 @@ async fn run_ratatui_app( ) .await; - restore(); + terminal_restore_guard.restore_silently(); // Mark the end of the recorded session. session_log::log_session_end(); // ignore error when collecting usage – report underlying error instead @@ -1468,6 +1469,38 @@ fn restore() { } } +struct TerminalRestoreGuard { + active: bool, +} + +impl TerminalRestoreGuard { + fn new() -> Self { + Self { active: true } + } + + #[cfg_attr(debug_assertions, allow(dead_code))] + fn restore(&mut self) -> color_eyre::Result<()> { + if self.active { + crate::tui::restore()?; + self.active = false; + } + Ok(()) + } + + fn restore_silently(&mut self) { + if self.active { + restore(); + self.active = false; + } + } +} + +impl Drop for TerminalRestoreGuard { + fn drop(&mut self) { + self.restore_silently(); + } +} + /// Determine whether to use the terminal's alternate screen buffer. /// /// The alternate screen buffer provides a cleaner fullscreen experience without polluting diff --git a/codex-rs/tui_app_server/src/onboarding/auth.rs b/codex-rs/tui_app_server/src/onboarding/auth.rs index 612fdbe2ac..37f8819b62 100644 --- a/codex-rs/tui_app_server/src/onboarding/auth.rs +++ b/codex-rs/tui_app_server/src/onboarding/auth.rs @@ -163,37 +163,7 @@ impl KeyboardHandler for AuthModeWidget { } KeyCode::Esc => { tracing::info!("Esc pressed"); - let mut sign_in_state = self.sign_in_state.write().unwrap(); - match &*sign_in_state { - SignInState::ChatGptContinueInBrowser(state) => { - let request_handle = self.app_server_request_handle.clone(); - let login_id = state.login_id.clone(); - tokio::spawn(async move { - let _ = request_handle - .request_typed::( - ClientRequest::CancelLoginAccount { - request_id: onboarding_request_id(), - params: CancelLoginAccountParams { login_id }, - }, - ) - .await; - }); - *sign_in_state = SignInState::PickMode; - drop(sign_in_state); - self.set_error(/*message*/ None); - self.request_frame.schedule_frame(); - } - SignInState::ChatGptDeviceCode(state) => { - if let Some(cancel) = &state.cancel { - cancel.notify_one(); - } - *sign_in_state = SignInState::PickMode; - drop(sign_in_state); - self.set_error(/*message*/ None); - self.request_frame.schedule_frame(); - } - _ => {} - } + self.cancel_active_attempt(); } _ => {} } @@ -221,6 +191,36 @@ pub(crate) struct AuthModeWidget { } impl AuthModeWidget { + pub(crate) fn cancel_active_attempt(&self) { + let mut sign_in_state = self.sign_in_state.write().unwrap(); + match &*sign_in_state { + SignInState::ChatGptContinueInBrowser(state) => { + let request_handle = self.app_server_request_handle.clone(); + let login_id = state.login_id.clone(); + tokio::spawn(async move { + let _ = request_handle + .request_typed::( + ClientRequest::CancelLoginAccount { + request_id: onboarding_request_id(), + params: CancelLoginAccountParams { login_id }, + }, + ) + .await; + }); + } + SignInState::ChatGptDeviceCode(state) => { + if let Some(cancel) = &state.cancel { + cancel.notify_one(); + } + } + _ => return, + } + *sign_in_state = SignInState::PickMode; + drop(sign_in_state); + self.set_error(/*message*/ None); + self.request_frame.schedule_frame(); + } + fn set_error(&self, message: Option) { *self.error.write().unwrap() = message; } @@ -771,6 +771,7 @@ impl AuthModeWidget { .await { Ok(LoginAccountResponse::Chatgpt { login_id, auth_url }) => { + maybe_open_auth_url_in_browser(&request_handle, &auth_url); *error.write().unwrap() = None; *sign_in_state.write().unwrap() = SignInState::ChatGptContinueInBrowser(ContinueInBrowserState { @@ -880,6 +881,16 @@ impl WidgetRef for AuthModeWidget { } } +pub(super) fn maybe_open_auth_url_in_browser(request_handle: &AppServerRequestHandle, url: &str) { + if !matches!(request_handle, AppServerRequestHandle::InProcess(_)) { + return; + } + + if let Err(err) = webbrowser::open(url) { + tracing::warn!("failed to open browser for login URL: {err}"); + } +} + #[cfg(test)] mod tests { use super::*; @@ -990,6 +1001,50 @@ mod tests { )); } + #[tokio::test] + async fn cancel_active_attempt_resets_browser_login_state() { + let (widget, _tmp) = widget_forced_chatgpt().await; + *widget.error.write().unwrap() = Some("still logging in".to_string()); + *widget.sign_in_state.write().unwrap() = + SignInState::ChatGptContinueInBrowser(ContinueInBrowserState { + login_id: "login-1".to_string(), + auth_url: "https://auth.example.com".to_string(), + }); + + widget.cancel_active_attempt(); + + assert_eq!(widget.error_message(), None); + assert!(matches!( + &*widget.sign_in_state.read().unwrap(), + SignInState::PickMode + )); + } + + #[tokio::test] + async fn cancel_active_attempt_notifies_device_code_login() { + let (widget, _tmp) = widget_forced_chatgpt().await; + let cancel = Arc::new(Notify::new()); + *widget.error.write().unwrap() = Some("still logging in".to_string()); + *widget.sign_in_state.write().unwrap() = + SignInState::ChatGptDeviceCode(ContinueWithDeviceCodeState { + device_code: None, + cancel: Some(cancel.clone()), + }); + + widget.cancel_active_attempt(); + + assert_eq!(widget.error_message(), None); + assert!(matches!( + &*widget.sign_in_state.read().unwrap(), + SignInState::PickMode + )); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), cancel.notified()) + .await + .is_ok() + ); + } + /// Collects all buffer cell symbols that contain the OSC 8 open sequence /// for the given URL. Returns the concatenated "inner" characters. fn collect_osc8_chars(buf: &Buffer, area: Rect, url: &str) -> String { diff --git a/codex-rs/tui_app_server/src/onboarding/auth/headless_chatgpt_login.rs b/codex-rs/tui_app_server/src/onboarding/auth/headless_chatgpt_login.rs index 33afe740b8..d0834fe166 100644 --- a/codex-rs/tui_app_server/src/onboarding/auth/headless_chatgpt_login.rs +++ b/codex-rs/tui_app_server/src/onboarding/auth/headless_chatgpt_login.rs @@ -29,6 +29,7 @@ use super::ContinueInBrowserState; use super::ContinueWithDeviceCodeState; use super::SignInState; use super::mark_url_hyperlink; +use super::maybe_open_auth_url_in_browser; use super::onboarding_request_id; pub(super) fn start_headless_chatgpt_login(widget: &mut AuthModeWidget) { @@ -289,6 +290,7 @@ async fn fallback_to_browser_login( .await { Ok(LoginAccountResponse::Chatgpt { login_id, auth_url }) => { + maybe_open_auth_url_in_browser(&request_handle, &auth_url); *error.write().unwrap() = None; let _updated = set_device_code_state_for_active_attempt( &sign_in_state, diff --git a/codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs b/codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs index b6afd2cb98..d092d06e2f 100644 --- a/codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs +++ b/codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs @@ -206,6 +206,14 @@ impl OnboardingScreen { self.should_exit } + fn cancel_auth_if_active(&self) { + for step in &self.steps { + if let Step::Auth(widget) = step { + widget.cancel_active_attempt(); + } + } + } + fn auth_widget_mut(&mut self) -> Option<&mut AuthModeWidget> { self.steps.iter_mut().find_map(|step| match step { Step::Auth(widget) => Some(widget), @@ -270,6 +278,7 @@ impl KeyboardHandler for OnboardingScreen { }; if should_quit { if self.is_auth_in_progress() { + self.cancel_auth_if_active(); // If the user cancels the auth menu, exit the app rather than // leave the user at a prompt in an unauthed state. self.should_exit = true;