From b887124e72b6247b2beed4dc7f3e339a1611503f Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 1/7] chore: prefer returning Err to expect() --- codex-rs/login/src/server.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index ef85df69ea..19ef4c1cfe 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -59,10 +59,9 @@ pub struct LoginServer { impl LoginServer { pub fn block_until_done(self) -> io::Result<()> { - #[expect(clippy::expect_used)] self.server_handle .join() - .expect("can't join on the server thread") + .map_err(|err| io::Error::other(format!("login server thread panicked: {err:?}")))? } pub fn cancel(&self) { From 74ccd9dba103d98c709da773a5fea722e76af132 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 2/7] fix: async-ify login flow --- codex-rs/cli/src/login.rs | 2 +- codex-rs/login/src/server.rs | 155 ++++++++++-------- codex-rs/login/tests/login_server_e2e.rs | 20 +-- .../mcp-server/src/codex_message_processor.rs | 12 +- codex-rs/tui/src/onboarding/auth.rs | 36 ++-- 5 files changed, 126 insertions(+), 99 deletions(-) diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 5f9dc5f908..fc40a0271f 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -21,7 +21,7 @@ pub async fn login_with_chatgpt(codex_home: PathBuf) -> std::io::Result<()> { server.actual_port, server.auth_url, ); - server.block_until_done() + server.block_until_done().await } pub async fn run_login_with_chatgpt(cli_config_overrides: CliConfigOverrides) -> ! { diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 19ef4c1cfe..060b333c6e 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -52,15 +52,15 @@ impl ServerOptions { pub struct LoginServer { pub auth_url: String, pub actual_port: u16, - pub server_handle: thread::JoinHandle>, pub shutdown_flag: Arc, - pub server: Arc, + server_handle: tokio::task::JoinHandle>, + server: Arc, } impl LoginServer { - pub fn block_until_done(self) -> io::Result<()> { + pub async fn block_until_done(self) -> io::Result<()> { self.server_handle - .join() + .await .map_err(|err| io::Error::other(format!("login server thread panicked: {err:?}")))? } @@ -118,7 +118,8 @@ pub fn run_login_server( if opts.open_browser { let _ = webbrowser::open(&auth_url); } - let shutdown_flag = shutdown_flag.unwrap_or_else(|| Arc::new(AtomicBool::new(false))); + let shutdown_flag: Arc = + shutdown_flag.unwrap_or_else(|| Arc::new(AtomicBool::new(false))); let shutdown_flag_clone = shutdown_flag.clone(); let timeout_flag = Arc::new(AtomicBool::new(false)); @@ -135,31 +136,46 @@ pub fn run_login_server( ); } - let server_for_thread = server.clone(); - let server_handle = thread::spawn(move || { - while !shutdown_flag.load(Ordering::SeqCst) { - let req = match server_for_thread.recv() { - Ok(r) => r, - 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)); + 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)); + } } - } - }; + }; + } + Ok(()) + }) + }; + + 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 response = process_request(&req, &opts, &redirect_uri, &pkce, actual_port, &state); let is_login_complete = matches!(response, HandledRequest::ResponseAndExit(_)); match response { HandledRequest::Response(r) | HandledRequest::ResponseAndExit(r) => { - let _ = req.respond(r); + let _ = tokio::task::spawn_blocking(move || req.respond(r)).await; } HandledRequest::RedirectWithHeader(header) => { let redirect = Response::empty(302).with_header(header); - let _ = req.respond(redirect); + let _ = tokio::task::spawn_blocking(move || req.respond(redirect)).await; } } @@ -196,15 +212,14 @@ enum HandledRequest { ResponseAndExit(Response>>), } -fn process_request( - req: &Request, +async fn process_request( + url_raw: &str, opts: &ServerOptions, redirect_uri: &str, pkce: &PkceCodes, actual_port: u16, state: &str, ) -> HandledRequest { - let url_raw = req.url().to_string(); let parsed_url = match url::Url::parse(&format!("http://localhost{url_raw}")) { Ok(u) => u, Err(e) => { @@ -235,18 +250,22 @@ fn process_request( }; match exchange_code_for_tokens(&opts.issuer, &opts.client_id, redirect_uri, pkce, &code) + .await { Ok(tokens) => { // Obtain API key via token-exchange and persist - let api_key = - obtain_api_key(&opts.issuer, &opts.client_id, &tokens.id_token).ok(); - if let Err(err) = persist_tokens( + let api_key = obtain_api_key(&opts.issuer, &opts.client_id, &tokens.id_token) + .await + .ok(); + if let Err(err) = persist_tokens_async( &opts.codex_home, api_key.clone(), tokens.id_token.clone(), Some(tokens.access_token.clone()), Some(tokens.refresh_token.clone()), - ) { + ) + .await + { eprintln!("Persist error: {err}"); return HandledRequest::Response( Response::from_string(format!("Unable to persist auth file: {err}")) @@ -352,7 +371,7 @@ struct ExchangedTokens { refresh_token: String, } -fn exchange_code_for_tokens( +async fn exchange_code_for_tokens( issuer: &str, client_id: &str, redirect_uri: &str, @@ -366,7 +385,7 @@ fn exchange_code_for_tokens( refresh_token: String, } - let client = reqwest::blocking::Client::new(); + let client = reqwest::Client::new(); let resp = client .post(format!("{issuer}/oauth/token")) .header("Content-Type", "application/x-www-form-urlencoded") @@ -378,6 +397,7 @@ fn exchange_code_for_tokens( urlencoding::encode(&pkce.code_verifier) )) .send() + .await .map_err(io::Error::other)?; if !resp.status().is_success() { @@ -387,7 +407,7 @@ fn exchange_code_for_tokens( ))); } - let tokens: TokenResponse = resp.json().map_err(io::Error::other)?; + let tokens: TokenResponse = resp.json().await.map_err(io::Error::other)?; Ok(ExchangedTokens { id_token: tokens.id_token, access_token: tokens.access_token, @@ -395,43 +415,49 @@ fn exchange_code_for_tokens( }) } -fn persist_tokens( +async fn persist_tokens_async( codex_home: &Path, api_key: Option, id_token: String, access_token: Option, refresh_token: Option, ) -> io::Result<()> { - let auth_file = get_auth_file(codex_home); - if let Some(parent) = auth_file.parent() { - if !parent.exists() { - std::fs::create_dir_all(parent).map_err(io::Error::other)?; + // Reuse existing synchronous logic but run it off the async runtime. + let codex_home = codex_home.to_path_buf(); + tokio::task::spawn_blocking(move || { + let auth_file = get_auth_file(&codex_home); + if let Some(parent) = auth_file.parent() { + if !parent.exists() { + std::fs::create_dir_all(parent).map_err(io::Error::other)?; + } } - } - let mut auth = read_or_default(&auth_file); - if let Some(key) = api_key { - auth.openai_api_key = Some(key); - } - let tokens = auth - .tokens - .get_or_insert_with(crate::token_data::TokenData::default); - tokens.id_token = crate::token_data::parse_id_token(&id_token).map_err(io::Error::other)?; - // Persist chatgpt_account_id if present in claims - if let Some(acc) = jwt_auth_claims(&id_token) - .get("chatgpt_account_id") - .and_then(|v| v.as_str()) - { - tokens.account_id = Some(acc.to_string()); - } - if let Some(at) = access_token { - tokens.access_token = at; - } - if let Some(rt) = refresh_token { - tokens.refresh_token = rt; - } - auth.last_refresh = Some(Utc::now()); - super::write_auth_json(&auth_file, &auth) + let mut auth = read_or_default(&auth_file); + if let Some(key) = api_key { + auth.openai_api_key = Some(key); + } + let tokens = auth + .tokens + .get_or_insert_with(crate::token_data::TokenData::default); + tokens.id_token = crate::token_data::parse_id_token(&id_token).map_err(io::Error::other)?; + // Persist chatgpt_account_id if present in claims + if let Some(acc) = jwt_auth_claims(&id_token) + .get("chatgpt_account_id") + .and_then(|v| v.as_str()) + { + tokens.account_id = Some(acc.to_string()); + } + if let Some(at) = access_token { + tokens.access_token = at; + } + if let Some(rt) = refresh_token { + tokens.refresh_token = rt; + } + auth.last_refresh = Some(Utc::now()); + super::write_auth_json(&auth_file, &auth) + }) + .await + .map_err(|e| io::Error::other(format!("persist task failed: {e}")))? } fn read_or_default(path: &Path) -> AuthDotJson { @@ -524,13 +550,13 @@ fn jwt_auth_claims(jwt: &str) -> serde_json::Map { serde_json::Map::new() } -fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result { +async fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result { // Token exchange for an API key access token #[derive(serde::Deserialize)] struct ExchangeResp { access_token: String, } - let client = reqwest::blocking::Client::new(); + let client = reqwest::Client::new(); let resp = client .post(format!("{issuer}/oauth/token")) .header("Content-Type", "application/x-www-form-urlencoded") @@ -543,6 +569,7 @@ fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result io::Result (SocketAddr, thread::JoinHandle<()>) { (addr, handle) } -#[test] -fn end_to_end_login_flow_persists_auth_json() { +#[tokio::test] +async fn end_to_end_login_flow_persists_auth_json() { if std::env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { println!( "Skipping test because it cannot execute when network is disabled in a Codex sandbox." @@ -106,16 +106,16 @@ fn end_to_end_login_flow_persists_auth_json() { let login_port = server.actual_port; // Simulate browser callback, and follow redirect to /success - let client = reqwest::blocking::Client::builder() + let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::limited(5)) .build() .unwrap(); let url = format!("http://127.0.0.1:{login_port}/auth/callback?code=abc&state=test_state_123"); - let resp = client.get(&url).send().unwrap(); + let resp = client.get(&url).send().await.unwrap(); assert!(resp.status().is_success()); // Wait for server shutdown - server.block_until_done().unwrap(); + server.block_until_done().await.unwrap(); // Validate auth.json let auth_path = codex_home.join("auth.json"); @@ -133,8 +133,8 @@ fn end_to_end_login_flow_persists_auth_json() { drop(issuer_handle); } -#[test] -fn creates_missing_codex_home_dir() { +#[tokio::test] +async fn creates_missing_codex_home_dir() { if std::env::var(CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() { println!( "Skipping test because it cannot execute when network is disabled in a Codex sandbox." @@ -164,12 +164,12 @@ fn creates_missing_codex_home_dir() { let server = run_login_server(opts, None).unwrap(); let login_port = server.actual_port; - let client = reqwest::blocking::Client::new(); + let client = reqwest::Client::new(); let url = format!("http://127.0.0.1:{login_port}/auth/callback?code=abc&state=state2"); - let resp = client.get(&url).send().unwrap(); + let resp = client.get(&url).send().await.unwrap(); assert!(resp.status().is_success()); - server.block_until_done().unwrap(); + server.block_until_done().await.unwrap(); let auth_path = codex_home.join("auth.json"); assert!( diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index d13bdbf346..7e5da55a32 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -180,15 +180,9 @@ impl CodexMessageProcessor { let outgoing_clone = self.outgoing.clone(); let active_login = self.active_login.clone(); tokio::spawn(async move { - let result = - tokio::task::spawn_blocking(move || server.block_until_done()).await; - let (success, error_msg) = match result { - Ok(Ok(())) => (true, None), - Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), - Err(join_err) => ( - false, - Some(format!("failed to join login server thread: {join_err}")), - ), + let (success, error_msg) = match server.block_until_done().await { + Ok(()) => (true, None), + Err(err) => (false, Some(format!("Login server error: {err}"))), }; let notification = LoginChatGptCompleteNotification { login_id, diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 8407e84e21..7961d75b5a 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -27,7 +27,6 @@ use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; -use std::thread::JoinHandle; use super::onboarding_screen::StepState; // no additional imports @@ -47,7 +46,7 @@ pub(crate) enum SignInState { pub(crate) struct ContinueInBrowserState { auth_url: String, shutdown_flag: Option>, - _login_wait_handle: Option>, + _login_wait_handle: Option>, } impl Drop for ContinueInBrowserState { fn drop(&mut self) { @@ -288,11 +287,16 @@ impl AuthModeWidget { Ok(child) => { let auth_url = child.auth_url.clone(); let shutdown_flag = child.shutdown_flag.clone(); + + let event_tx = self.event_tx.clone(); + let join_handle = tokio::spawn(async move { + spawn_completion_poller(child, event_tx).await; + }); self.sign_in_state = SignInState::ChatGptContinueInBrowser(ContinueInBrowserState { auth_url, shutdown_flag: Some(shutdown_flag), - _login_wait_handle: Some(self.spawn_completion_poller(child)), + _login_wait_handle: Some(join_handle), }); self.event_tx.send(AppEvent::RequestRedraw); } @@ -313,19 +317,21 @@ impl AuthModeWidget { } self.event_tx.send(AppEvent::RequestRedraw); } +} - fn spawn_completion_poller(&self, child: codex_login::LoginServer) -> JoinHandle<()> { - let event_tx = self.event_tx.clone(); - std::thread::spawn(move || { - if let Ok(()) = child.block_until_done() { - event_tx.send(AppEvent::OnboardingAuthComplete(Ok(()))); - } else { - event_tx.send(AppEvent::OnboardingAuthComplete(Err( - "login failed".to_string() - ))); - } - }) - } +async fn spawn_completion_poller( + child: codex_login::LoginServer, + event_tx: AppEventSender, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + if let Ok(()) = child.block_until_done().await { + event_tx.send(AppEvent::OnboardingAuthComplete(Ok(()))); + } else { + event_tx.send(AppEvent::OnboardingAuthComplete(Err( + "login failed".to_string() + ))); + } + }) } impl StepStateProvider for AuthModeWidget { From dc7ef4c896d43c86e2f4b665fa8360eb620c65cf Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 3/7] fix: change `shutdown_flag` from `Arc` to `tokio::sync::Notify` --- codex-rs/login/src/server.rs | 134 +++++++++++++++------------- codex-rs/tui/src/onboarding/auth.rs | 15 ++-- 2 files changed, 77 insertions(+), 72 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 060b333c6e..419874e7a0 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, + 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,32 @@ impl LoginServer { #[derive(Clone)] pub struct ShutdownHandle { - shutdown_flag: Arc, + shutdown_notify: Arc, server: Arc, } +impl std::fmt::Debug for ShutdownHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ShutdownHandle") + .field("shutdown_notify", &self.shutdown_notify) + .finish() + } +} + 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 +126,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 +138,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 +147,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 +210,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 +326,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 7961d75b5a..7166e349c3 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -1,5 +1,6 @@ use codex_login::CLIENT_ID; use codex_login::ServerOptions; +use codex_login::ShutdownHandle; use codex_login::run_login_server; use crossterm::event::KeyCode; use crossterm::event::KeyEvent; @@ -24,9 +25,6 @@ use crate::onboarding::onboarding_screen::KeyboardHandler; 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 super::onboarding_screen::StepState; // no additional imports @@ -45,13 +43,14 @@ 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_handle: 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); + if let Some(flag) = &self.shutdown_handle { + flag.cancel(); } } } @@ -286,7 +285,7 @@ impl AuthModeWidget { match server { Ok(child) => { let auth_url = child.auth_url.clone(); - let shutdown_flag = child.shutdown_flag.clone(); + let shutdown_handle = child.cancel_handle(); let event_tx = self.event_tx.clone(); let join_handle = tokio::spawn(async move { @@ -295,7 +294,7 @@ impl AuthModeWidget { self.sign_in_state = SignInState::ChatGptContinueInBrowser(ContinueInBrowserState { auth_url, - shutdown_flag: Some(shutdown_flag), + shutdown_handle: Some(shutdown_handle), _login_wait_handle: Some(join_handle), }); self.event_tx.send(AppEvent::RequestRedraw); From df8a64494b8cfc4ff4ae2bdfdd24060c17b693ab Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 4/7] fix: eliminate ServerOptions.login_timeout and have caller use tokio::time::timeout() instead --- codex-rs/login/src/server.rs | 115 +++++------------- codex-rs/login/tests/login_server_e2e.rs | 2 - .../mcp-server/src/codex_message_processor.rs | 20 ++- 3 files changed, 48 insertions(+), 89 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 419874e7a0..f33f1ae4fb 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, } } } @@ -126,24 +120,8 @@ pub fn run_login_server( if opts.open_browser { let _ = webbrowser::open(&auth_url); } - 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(), - ); - } + // Map blocking reads from server.recv() to an async channel. let (tx, mut rx) = tokio::sync::mpsc::channel::(16); let _server_handle = { let server = server.clone(); @@ -158,59 +136,52 @@ pub fn run_login_server( }) }; - let server_for_task = server.clone(); - let server_handle = tokio::spawn(async move { - 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 { + let shutdown_notify = shutdown_flag.unwrap_or_else(|| Arc::new(tokio::sync::Notify::new())); + let server_handle = { + let shutdown_notify = shutdown_notify.clone(); + let server = server.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = shutdown_notify.notified() => { 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 { + maybe_req = rx.recv() => { + let Some(req) = maybe_req 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 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; + 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; + } } - 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_notify.notify_waiters(); + server.unblock(); + return Ok(()); + } } } } - } - }); + }) + }; Ok(LoginServer { - auth_url: auth_url.clone(), + auth_url, actual_port, server_handle, - shutdown_flag: shutdown_notify_clone, + shutdown_flag: shutdown_notify, server, }) } @@ -319,26 +290,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 7e5da55a32..00c8717c43 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -146,7 +146,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()) }; @@ -158,6 +157,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. { @@ -166,7 +166,7 @@ impl CodexMessageProcessor { existing.drop(); } *guard = Some(ActiveLogin { - shutdown_handle: server.cancel_handle(), + shutdown_handle: shutdown_handle.clone(), login_id, }); } @@ -180,9 +180,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, From 7a91cf2d7882ca4ddaf474131ab5c5ca2877c4d3 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 5/7] fix: make ShutdownHandle a private field of LoginServer --- codex-rs/login/src/server.rs | 26 +++++++------------ .../mcp-server/src/codex_message_processor.rs | 4 +-- codex-rs/tui/src/onboarding/auth.rs | 2 +- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index f33f1ae4fb..f3256ee6c8 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -46,9 +46,8 @@ impl ServerOptions { pub struct LoginServer { pub auth_url: String, pub actual_port: u16, - shutdown_flag: Arc, server_handle: tokio::task::JoinHandle>, - server: Arc, + shutdown_handle: ShutdownHandle, } impl LoginServer { @@ -59,14 +58,11 @@ impl LoginServer { } pub fn cancel(&self) { - shutdown(&self.shutdown_flag, &self.server); + self.shutdown_handle.shutdown(); } pub fn cancel_handle(&self) -> ShutdownHandle { - ShutdownHandle { - shutdown_notify: self.shutdown_flag.clone(), - server: self.server.clone(), - } + self.shutdown_handle.clone() } } @@ -85,16 +81,12 @@ impl std::fmt::Debug for ShutdownHandle { } impl ShutdownHandle { - pub fn cancel(&self) { - shutdown(&self.shutdown_notify, &self.server); + pub fn shutdown(&self) { + self.shutdown_notify.notify_waiters(); + self.server.unblock(); } } -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>, @@ -181,8 +173,10 @@ pub fn run_login_server( auth_url, actual_port, server_handle, - shutdown_flag: shutdown_notify, - server, + shutdown_handle: ShutdownHandle { + shutdown_notify, + server, + }, }) } diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index 00c8717c43..4f2b3bb693 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -66,7 +66,7 @@ struct ActiveLogin { impl ActiveLogin { fn drop(&self) { - self.shutdown_handle.cancel(); + self.shutdown_handle.shutdown(); } } @@ -190,7 +190,7 @@ impl CodexMessageProcessor { Ok(Err(err)) => (false, Some(format!("Login server error: {err}"))), Err(_elapsed) => { // Timeout: cancel server and report - shutdown_handle.cancel(); + shutdown_handle.shutdown(); (false, Some("Login timed out".to_string())) } }; diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 7166e349c3..490f85bff8 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -50,7 +50,7 @@ pub(crate) struct ContinueInBrowserState { impl Drop for ContinueInBrowserState { fn drop(&mut self) { if let Some(flag) = &self.shutdown_handle { - flag.cancel(); + flag.shutdown(); } } } From 977e586b96449377867df7e803eeb1e78e93469f Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:47:10 -0700 Subject: [PATCH 6/7] fix: reduce references to Server in codex-login crate --- codex-rs/login/src/server.rs | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index f3256ee6c8..e8af09fba0 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -66,24 +66,14 @@ impl LoginServer { } } -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct ShutdownHandle { shutdown_notify: Arc, - server: Arc, -} - -impl std::fmt::Debug for ShutdownHandle { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ShutdownHandle") - .field("shutdown_notify", &self.shutdown_notify) - .finish() - } } impl ShutdownHandle { pub fn shutdown(&self) { self.shutdown_notify.notify_waiters(); - self.server.unblock(); } } @@ -133,14 +123,14 @@ pub fn run_login_server( let shutdown_notify = shutdown_notify.clone(); let server = server.clone(); tokio::spawn(async move { - loop { + let result = loop { tokio::select! { _ = shutdown_notify.notified() => { - return Err(io::Error::other("Login was not completed")); + break Err(io::Error::other("Login was not completed")); } maybe_req = rx.recv() => { let Some(req) = maybe_req else { - return Err(io::Error::other("Login was not completed")); + break Err(io::Error::other("Login was not completed")); }; let url_raw = req.url().to_string(); @@ -159,13 +149,16 @@ pub fn run_login_server( } if is_login_complete { - shutdown_notify.notify_waiters(); - server.unblock(); - return Ok(()); + break Ok(()); } } } - } + }; + + // Ensure that the server is unblocked so the thread dedicated to + // running `server.recv()` in a loop exits cleanly. + server.unblock(); + result }) }; @@ -173,10 +166,7 @@ pub fn run_login_server( auth_url, actual_port, server_handle, - shutdown_handle: ShutdownHandle { - shutdown_notify, - server, - }, + shutdown_handle: ShutdownHandle { shutdown_notify }, }) } From 8a46a47f420e826ddc05d9b7e495f3a370049efc Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 18 Aug 2025 09:49:35 -0700 Subject: [PATCH 7/7] fix: remove shutdown_flag param to run_login_server() --- codex-rs/cli/src/login.rs | 2 +- codex-rs/login/src/server.rs | 7 ++----- codex-rs/login/tests/login_server_e2e.rs | 4 ++-- codex-rs/mcp-server/src/codex_message_processor.rs | 2 +- codex-rs/tui/src/onboarding/auth.rs | 2 +- 5 files changed, 7 insertions(+), 10 deletions(-) diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index fc40a0271f..36bbf2208d 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -14,7 +14,7 @@ use std::path::PathBuf; pub async fn login_with_chatgpt(codex_home: PathBuf) -> std::io::Result<()> { let opts = ServerOptions::new(codex_home, CLIENT_ID.to_string()); - let server = run_login_server(opts, None)?; + let server = run_login_server(opts)?; eprintln!( "Starting local login server on http://localhost:{}.\nIf your browser did not open, navigate to this URL to authenticate:\n\n{}", diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index e8af09fba0..32229484ff 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -77,10 +77,7 @@ impl ShutdownHandle { } } -pub fn run_login_server( - opts: ServerOptions, - shutdown_flag: Option>, -) -> io::Result { +pub fn run_login_server(opts: ServerOptions) -> io::Result { let pkce = generate_pkce(); let state = opts.force_state.clone().unwrap_or_else(generate_state); @@ -118,7 +115,7 @@ pub fn run_login_server( }) }; - let shutdown_notify = shutdown_flag.unwrap_or_else(|| Arc::new(tokio::sync::Notify::new())); + let shutdown_notify = Arc::new(tokio::sync::Notify::new()); let server_handle = { let shutdown_notify = shutdown_notify.clone(); let server = server.clone(); diff --git a/codex-rs/login/tests/login_server_e2e.rs b/codex-rs/login/tests/login_server_e2e.rs index ef387f575e..ceb0a94733 100644 --- a/codex-rs/login/tests/login_server_e2e.rs +++ b/codex-rs/login/tests/login_server_e2e.rs @@ -101,7 +101,7 @@ async fn end_to_end_login_flow_persists_auth_json() { open_browser: false, force_state: Some(state), }; - let server = run_login_server(opts, None).unwrap(); + let server = run_login_server(opts).unwrap(); let login_port = server.actual_port; // Simulate browser callback, and follow redirect to /success @@ -159,7 +159,7 @@ async fn creates_missing_codex_home_dir() { open_browser: false, force_state: Some(state), }; - let server = run_login_server(opts, None).unwrap(); + let server = run_login_server(opts).unwrap(); let login_port = server.actual_port; let client = reqwest::Client::new(); diff --git a/codex-rs/mcp-server/src/codex_message_processor.rs b/codex-rs/mcp-server/src/codex_message_processor.rs index 4f2b3bb693..1decf11da6 100644 --- a/codex-rs/mcp-server/src/codex_message_processor.rs +++ b/codex-rs/mcp-server/src/codex_message_processor.rs @@ -154,7 +154,7 @@ impl CodexMessageProcessor { Error(JSONRPCErrorError), } - let reply = match run_login_server(opts, None) { + let reply = match run_login_server(opts) { Ok(server) => { let login_id = Uuid::new_v4(); let shutdown_handle = server.cancel_handle(); diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 490f85bff8..facd06c64a 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -281,7 +281,7 @@ impl AuthModeWidget { fn start_chatgpt_login(&mut self) { self.error = None; let opts = ServerOptions::new(self.codex_home.clone(), CLIENT_ID.to_string()); - let server = run_login_server(opts, None); + let server = run_login_server(opts); match server { Ok(child) => { let auth_url = child.auth_url.clone();