diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index f3256ee6..e8af09fb 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 }, }) }