From d699a3af7539e46ac77ae2efd6f4baa3225d5cea Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Sun, 17 Aug 2025 14:12:18 -0700 Subject: [PATCH 1/2] fix: change `shutdown_flag` from `Arc` to `tokio::sync::Notify` --- codex-rs/login/src/server.rs | 126 ++++++++++++++-------------- codex-rs/tui/src/onboarding/auth.rs | 7 +- 2 files changed, 65 insertions(+), 68 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 060b333c6e..e5891ac70b 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -52,7 +52,7 @@ impl ServerOptions { pub struct LoginServer { pub auth_url: String, pub actual_port: u16, - pub shutdown_flag: Arc, + pub shutdown_flag: Arc, server_handle: tokio::task::JoinHandle>, server: Arc, } @@ -70,7 +70,7 @@ impl LoginServer { pub fn cancel_handle(&self) -> ShutdownHandle { ShutdownHandle { - shutdown_flag: self.shutdown_flag.clone(), + shutdown_notify: self.shutdown_flag.clone(), server: self.server.clone(), } } @@ -78,24 +78,24 @@ impl LoginServer { #[derive(Clone)] pub struct ShutdownHandle { - shutdown_flag: Arc, + shutdown_notify: Arc, server: Arc, } impl ShutdownHandle { pub fn cancel(&self) { - shutdown(&self.shutdown_flag, &self.server); + shutdown(&self.shutdown_notify, &self.server); } } -pub fn shutdown(shutdown_flag: &AtomicBool, server: &Server) { - shutdown_flag.store(true, Ordering::SeqCst); +pub fn shutdown(shutdown_notify: &tokio::sync::Notify, server: &Server) { + shutdown_notify.notify_waiters(); server.unblock(); } pub fn run_login_server( opts: ServerOptions, - shutdown_flag: Option>, + shutdown_flag: Option>, ) -> io::Result { let pkce = generate_pkce(); let state = opts.force_state.clone().unwrap_or_else(generate_state); @@ -118,9 +118,9 @@ pub fn run_login_server( if opts.open_browser { let _ = webbrowser::open(&auth_url); } - let shutdown_flag: Arc = - shutdown_flag.unwrap_or_else(|| Arc::new(AtomicBool::new(false))); - let shutdown_flag_clone = shutdown_flag.clone(); + let shutdown_notify: Arc = + shutdown_flag.unwrap_or_else(|| Arc::new(tokio::sync::Notify::new())); + let shutdown_notify_clone = shutdown_notify.clone(); let timeout_flag = Arc::new(AtomicBool::new(false)); // Channel used to signal completion to timeout watcher. @@ -130,7 +130,7 @@ pub fn run_login_server( spawn_timeout_watcher( done_rx, timeout, - shutdown_flag.clone(), + shutdown_notify.clone(), timeout_flag.clone(), server.clone(), ); @@ -139,61 +139,62 @@ pub fn run_login_server( let (tx, mut rx) = tokio::sync::mpsc::channel::(16); let _server_handle = { let server = server.clone(); - let shutdown_flag = shutdown_flag.clone(); - thread::spawn(move || { - while !shutdown_flag.load(Ordering::SeqCst) { - match server.recv() { - Ok(request) => tx.blocking_send(request).map_err(|e| { - eprintln!("Failed to send request to channel: {e}"); - io::Error::other("Failed to send request to channel") - })?, - Err(e) => { - // If we've been asked to shut down, break gracefully so that - // we can report timeout or cancellation status uniformly. - if shutdown_flag.load(Ordering::SeqCst) { - break; - } else { - return Err(io::Error::other(e)); - } - } - }; + thread::spawn(move || -> io::Result<()> { + while let Ok(request) = server.recv() { + tx.blocking_send(request).map_err(|e| { + eprintln!("Failed to send request to channel: {e}"); + io::Error::other("Failed to send request to channel") + })?; } Ok(()) }) }; + let server_for_task = server.clone(); let server_handle = tokio::spawn(async move { - while let Some(req) = rx.recv().await { - let url_raw = req.url().to_string(); - let response = - process_request(&url_raw, &opts, &redirect_uri, &pkce, actual_port, &state).await; - - let is_login_complete = matches!(response, HandledRequest::ResponseAndExit(_)); - match response { - HandledRequest::Response(r) | HandledRequest::ResponseAndExit(r) => { - let _ = tokio::task::spawn_blocking(move || req.respond(r)).await; + loop { + tokio::select! { + _ = shutdown_notify.notified() => { + let _ = done_tx.send(()); + if timeout_flag.load(Ordering::SeqCst) { + return Err(io::Error::other("Login timed out")); + } else { + return Err(io::Error::other("Login was not completed")); + } } - HandledRequest::RedirectWithHeader(header) => { - let redirect = Response::empty(302).with_header(header); - let _ = tokio::task::spawn_blocking(move || req.respond(redirect)).await; + maybe_req = rx.recv() => { + let Some(req) = maybe_req else { + let _ = done_tx.send(()); + if timeout_flag.load(Ordering::SeqCst) { + return Err(io::Error::other("Login timed out")); + } else { + return Err(io::Error::other("Login was not completed")); + } + }; + + let url_raw = req.url().to_string(); + let response = + process_request(&url_raw, &opts, &redirect_uri, &pkce, actual_port, &state).await; + + let is_login_complete = matches!(response, HandledRequest::ResponseAndExit(_)); + match response { + HandledRequest::Response(r) | HandledRequest::ResponseAndExit(r) => { + let _ = tokio::task::spawn_blocking(move || req.respond(r)).await; + } + HandledRequest::RedirectWithHeader(header) => { + let redirect = Response::empty(302).with_header(header); + let _ = tokio::task::spawn_blocking(move || req.respond(redirect)).await; + } + } + + if is_login_complete { + shutdown_notify.notify_waiters(); + let _ = done_tx.send(()); + server_for_task.unblock(); + return Ok(()); + } } } - - if is_login_complete { - shutdown_flag.store(true, Ordering::SeqCst); - // Login has succeeded, so disarm the timeout watcher. - let _ = done_tx.send(()); - return Ok(()); - } - } - - // Login has failed or timed out, so disarm the timeout watcher. - let _ = done_tx.send(()); - - if timeout_flag.load(Ordering::SeqCst) { - Err(io::Error::other("Login timed out")) - } else { - Err(io::Error::other("Login was not completed")) } }); @@ -201,7 +202,7 @@ pub fn run_login_server( auth_url: auth_url.clone(), actual_port, server_handle, - shutdown_flag: shutdown_flag_clone, + shutdown_flag: shutdown_notify_clone, server, }) } @@ -317,17 +318,14 @@ async fn process_request( fn spawn_timeout_watcher( done_rx: mpsc::Receiver<()>, timeout: Duration, - shutdown_flag: Arc, + shutdown_notify: Arc, timeout_flag: Arc, server: Arc, ) { thread::spawn(move || { - if done_rx.recv_timeout(timeout).is_err() - && shutdown_flag - .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) - .is_ok() - { + if done_rx.recv_timeout(timeout).is_err() { timeout_flag.store(true, Ordering::SeqCst); + shutdown_notify.notify_waiters(); server.unblock(); } }); diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 046fb738e3..eecd088d21 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -26,8 +26,7 @@ use crate::onboarding::onboarding_screen::StepStateProvider; use crate::shimmer::shimmer_spans; use std::path::PathBuf; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; +use tokio::sync::Notify; use super::onboarding_screen::StepState; // no additional imports @@ -46,13 +45,13 @@ pub(crate) enum SignInState { /// Used to manage the lifecycle of SpawnedLogin and ensure it gets cleaned up. pub(crate) struct ContinueInBrowserState { auth_url: String, - shutdown_flag: Option>, + shutdown_flag: Option>, _login_wait_handle: Option>, } impl Drop for ContinueInBrowserState { fn drop(&mut self) { if let Some(flag) = &self.shutdown_flag { - flag.store(true, Ordering::SeqCst); + flag.notify_waiters(); } } } From 173fa5fab45d65caa660c7e8adccf17745661572 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Sun, 17 Aug 2025 14:26:24 -0700 Subject: [PATCH 2/2] fix: eliminate ServerOptions.login_timeout and have caller use tokio::time::timeout() instead --- codex-rs/login/src/server.rs | 55 +------------------ codex-rs/login/tests/login_server_e2e.rs | 2 - .../mcp-server/src/codex_message_processor.rs | 20 +++++-- 3 files changed, 17 insertions(+), 60 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index e5891ac70b..9eae4be348 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -3,11 +3,7 @@ use std::io::{self}; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::sync::mpsc; use std::thread; -use std::time::Duration; use crate::AuthDotJson; use crate::get_auth_file; @@ -32,7 +28,6 @@ pub struct ServerOptions { pub port: u16, pub open_browser: bool, pub force_state: Option, - pub login_timeout: Option, } impl ServerOptions { @@ -44,7 +39,6 @@ impl ServerOptions { port: DEFAULT_PORT, open_browser: true, force_state: None, - login_timeout: None, } } } @@ -121,20 +115,6 @@ pub fn run_login_server( let shutdown_notify: Arc = shutdown_flag.unwrap_or_else(|| Arc::new(tokio::sync::Notify::new())); let shutdown_notify_clone = shutdown_notify.clone(); - let timeout_flag = Arc::new(AtomicBool::new(false)); - - // Channel used to signal completion to timeout watcher. - let (done_tx, done_rx) = mpsc::channel::<()>(); - - if let Some(timeout) = opts.login_timeout { - spawn_timeout_watcher( - done_rx, - timeout, - shutdown_notify.clone(), - timeout_flag.clone(), - server.clone(), - ); - } let (tx, mut rx) = tokio::sync::mpsc::channel::(16); let _server_handle = { @@ -155,21 +135,11 @@ pub fn run_login_server( loop { tokio::select! { _ = shutdown_notify.notified() => { - let _ = done_tx.send(()); - if timeout_flag.load(Ordering::SeqCst) { - return Err(io::Error::other("Login timed out")); - } else { - return Err(io::Error::other("Login was not completed")); - } + return Err(io::Error::other("Login was not completed")); } maybe_req = rx.recv() => { let Some(req) = maybe_req else { - let _ = done_tx.send(()); - if timeout_flag.load(Ordering::SeqCst) { - return Err(io::Error::other("Login timed out")); - } else { - return Err(io::Error::other("Login was not completed")); - } + return Err(io::Error::other("Login was not completed")); }; let url_raw = req.url().to_string(); @@ -189,7 +159,6 @@ pub fn run_login_server( if is_login_complete { shutdown_notify.notify_waiters(); - let _ = done_tx.send(()); server_for_task.unblock(); return Ok(()); } @@ -311,26 +280,6 @@ async fn process_request( } } -/// Spawns a detached thread that waits for either a completion signal on `done_rx` -/// or the specified `timeout` to elapse. If the timeout elapses first it marks -/// the `shutdown_flag`, records `timeout_flag`, and unblocks the HTTP server so -/// that the main server loop can exit promptly. -fn spawn_timeout_watcher( - done_rx: mpsc::Receiver<()>, - timeout: Duration, - shutdown_notify: Arc, - timeout_flag: Arc, - server: Arc, -) { - thread::spawn(move || { - if done_rx.recv_timeout(timeout).is_err() { - timeout_flag.store(true, Ordering::SeqCst); - shutdown_notify.notify_waiters(); - server.unblock(); - } - }); -} - fn build_authorize_url( issuer: &str, client_id: &str, diff --git a/codex-rs/login/tests/login_server_e2e.rs b/codex-rs/login/tests/login_server_e2e.rs index 09a447d565..ef387f575e 100644 --- a/codex-rs/login/tests/login_server_e2e.rs +++ b/codex-rs/login/tests/login_server_e2e.rs @@ -100,7 +100,6 @@ async fn end_to_end_login_flow_persists_auth_json() { port: 0, open_browser: false, force_state: Some(state), - login_timeout: None, }; let server = run_login_server(opts, None).unwrap(); let login_port = server.actual_port; @@ -159,7 +158,6 @@ async fn creates_missing_codex_home_dir() { port: 0, open_browser: false, force_state: Some(state), - login_timeout: None, }; let server = run_login_server(opts, None).unwrap(); let login_port = server.actual_port; diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index 73a5b40c00..70341a0ab9 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -143,7 +143,6 @@ impl CodexMessageProcessor { let opts = LoginServerOptions { open_browser: false, - login_timeout: Some(LOGIN_CHATGPT_TIMEOUT), ..LoginServerOptions::new(config.codex_home.clone(), CLIENT_ID.to_string()) }; @@ -155,6 +154,7 @@ impl CodexMessageProcessor { let reply = match run_login_server(opts, None) { Ok(server) => { let login_id = Uuid::new_v4(); + let shutdown_handle = server.cancel_handle(); // Replace active login if present. { @@ -163,7 +163,7 @@ impl CodexMessageProcessor { existing.drop(); } *guard = Some(ActiveLogin { - shutdown_handle: server.cancel_handle(), + shutdown_handle: shutdown_handle.clone(), login_id, }); } @@ -177,9 +177,19 @@ impl CodexMessageProcessor { let outgoing_clone = self.outgoing.clone(); let active_login = self.active_login.clone(); tokio::spawn(async move { - let (success, error_msg) = match server.block_until_done().await { - Ok(()) => (true, None), - Err(err) => (false, Some(format!("Login server error: {err}"))), + let (success, error_msg) = match tokio::time::timeout( + LOGIN_CHATGPT_TIMEOUT, + server.block_until_done(), + ) + .await + { + Ok(Ok(())) => (true, None), + Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), + Err(_elapsed) => { + // Timeout: cancel server and report + shutdown_handle.cancel(); + (false, Some("Login timed out".to_string())) + } }; let notification = LoginChatGptCompleteNotification { login_id,