From 40efedf369a4f229db8f982cb3b79c0058b43d51 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 24 Mar 2026 13:03:25 -0600 Subject: [PATCH] tui_app_server: cancel active login before Ctrl+C exit --- .../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 + .../tui_app_server/src/onboarding/auth.rs | 106 +++++++++++++----- .../src/onboarding/onboarding_screen.rs | 9 ++ 5 files changed, 96 insertions(+), 31 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index c4c2799ca8d..2153bcf62f6 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 0826f33c678..4acc419c5b1 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 4f53bfc1a49..96fe437c054 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_app_server/src/onboarding/auth.rs b/codex-rs/tui_app_server/src/onboarding/auth.rs index 612fdbe2ace..542fb0555ca 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; } @@ -990,6 +990,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/onboarding_screen.rs b/codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs index b6afd2cb984..d092d06e2f9 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;