Merge 8a46a47f42 into sapling-pr-archive-bolinfest

This commit is contained in:
Michael Bolin
2025-08-18 09:49:51 -07:00
committed by GitHub
5 changed files with 177 additions and 206 deletions

View File

@@ -14,14 +14,14 @@ 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{}",
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) -> ! {

View File

@@ -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<String>,
pub login_timeout: Option<Duration>,
}
impl ServerOptions {
@@ -44,7 +39,6 @@ impl ServerOptions {
port: DEFAULT_PORT,
open_browser: true,
force_state: None,
login_timeout: None,
}
}
}
@@ -52,52 +46,38 @@ impl ServerOptions {
pub struct LoginServer {
pub auth_url: String,
pub actual_port: u16,
pub server_handle: thread::JoinHandle<io::Result<()>>,
pub shutdown_flag: Arc<AtomicBool>,
pub server: Arc<Server>,
server_handle: tokio::task::JoinHandle<io::Result<()>>,
shutdown_handle: ShutdownHandle,
}
impl LoginServer {
pub fn block_until_done(self) -> io::Result<()> {
#[expect(clippy::expect_used)]
pub async fn block_until_done(self) -> io::Result<()> {
self.server_handle
.join()
.expect("can't join on the server thread")
.await
.map_err(|err| io::Error::other(format!("login server thread panicked: {err:?}")))?
}
pub fn cancel(&self) {
shutdown(&self.shutdown_flag, &self.server);
self.shutdown_handle.shutdown();
}
pub fn cancel_handle(&self) -> ShutdownHandle {
ShutdownHandle {
shutdown_flag: self.shutdown_flag.clone(),
server: self.server.clone(),
}
self.shutdown_handle.clone()
}
}
#[derive(Clone)]
#[derive(Clone, Debug)]
pub struct ShutdownHandle {
shutdown_flag: Arc<AtomicBool>,
server: Arc<Server>,
shutdown_notify: Arc<tokio::sync::Notify>,
}
impl ShutdownHandle {
pub fn cancel(&self) {
shutdown(&self.shutdown_flag, &self.server);
pub fn shutdown(&self) {
self.shutdown_notify.notify_waiters();
}
}
pub fn shutdown(shutdown_flag: &AtomicBool, server: &Server) {
shutdown_flag.store(true, Ordering::SeqCst);
server.unblock();
}
pub fn run_login_server(
opts: ServerOptions,
shutdown_flag: Option<Arc<AtomicBool>>,
) -> io::Result<LoginServer> {
pub fn run_login_server(opts: ServerOptions) -> io::Result<LoginServer> {
let pkce = generate_pkce();
let state = opts.force_state.clone().unwrap_or_else(generate_state);
@@ -119,75 +99,71 @@ 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_clone = shutdown_flag.clone();
let timeout_flag = Arc::new(AtomicBool::new(false));
// Channel used to signal completion to timeout watcher.
let (done_tx, done_rx) = mpsc::channel::<()>();
// Map blocking reads from server.recv() to an async channel.
let (tx, mut rx) = tokio::sync::mpsc::channel::<Request>(16);
let _server_handle = {
let server = server.clone();
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(())
})
};
if let Some(timeout) = opts.login_timeout {
spawn_timeout_watcher(
done_rx,
timeout,
shutdown_flag.clone(),
timeout_flag.clone(),
server.clone(),
);
}
let shutdown_notify = Arc::new(tokio::sync::Notify::new());
let server_handle = {
let shutdown_notify = shutdown_notify.clone();
let server = server.clone();
tokio::spawn(async move {
let result = loop {
tokio::select! {
_ = shutdown_notify.notified() => {
break Err(io::Error::other("Login was not completed"));
}
maybe_req = rx.recv() => {
let Some(req) = maybe_req else {
break Err(io::Error::other("Login was not completed"));
};
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 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 {
break Ok(());
}
}
}
};
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);
}
HandledRequest::RedirectWithHeader(header) => {
let redirect = Response::empty(302).with_header(header);
let _ = req.respond(redirect);
}
}
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"))
}
});
// Ensure that the server is unblocked so the thread dedicated to
// running `server.recv()` in a loop exits cleanly.
server.unblock();
result
})
};
Ok(LoginServer {
auth_url: auth_url.clone(),
auth_url,
actual_port,
server_handle,
shutdown_flag: shutdown_flag_clone,
server,
shutdown_handle: ShutdownHandle { shutdown_notify },
})
}
@@ -197,15 +173,14 @@ enum HandledRequest {
ResponseAndExit(Response<Cursor<Vec<u8>>>),
}
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) => {
@@ -236,18 +211,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}"))
@@ -292,29 +271,6 @@ 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_flag: Arc<AtomicBool>,
timeout_flag: Arc<AtomicBool>,
server: Arc<Server>,
) {
thread::spawn(move || {
if done_rx.recv_timeout(timeout).is_err()
&& shutdown_flag
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.is_ok()
{
timeout_flag.store(true, Ordering::SeqCst);
server.unblock();
}
});
}
fn build_authorize_url(
issuer: &str,
client_id: &str,
@@ -353,7 +309,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,
@@ -367,7 +323,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")
@@ -379,6 +335,7 @@ fn exchange_code_for_tokens(
urlencoding::encode(&pkce.code_verifier)
))
.send()
.await
.map_err(io::Error::other)?;
if !resp.status().is_success() {
@@ -388,7 +345,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,
@@ -396,43 +353,49 @@ fn exchange_code_for_tokens(
})
}
fn persist_tokens(
async fn persist_tokens_async(
codex_home: &Path,
api_key: Option<String>,
id_token: String,
access_token: Option<String>,
refresh_token: Option<String>,
) -> 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 {
@@ -525,13 +488,13 @@ fn jwt_auth_claims(jwt: &str) -> serde_json::Map<String, serde_json::Value> {
serde_json::Map::new()
}
fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result<String> {
async fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result<String> {
// 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")
@@ -544,6 +507,7 @@ fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result<S
urlencoding::encode("urn:ietf:params:oauth:token-type:id_token")
))
.send()
.await
.map_err(io::Error::other)?;
if !resp.status().is_success() {
return Err(io::Error::other(format!(
@@ -551,6 +515,6 @@ fn obtain_api_key(issuer: &str, client_id: &str, id_token: &str) -> io::Result<S
resp.status()
)));
}
let body: ExchangeResp = resp.json().map_err(io::Error::other)?;
let body: ExchangeResp = resp.json().await.map_err(io::Error::other)?;
Ok(body.access_token)
}

View File

@@ -73,8 +73,8 @@ fn start_mock_issuer() -> (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."
@@ -100,22 +100,21 @@ 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 server = run_login_server(opts).unwrap();
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 +132,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."
@@ -159,17 +158,16 @@ 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 server = run_login_server(opts).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!(

View File

@@ -66,7 +66,7 @@ struct ActiveLogin {
impl ActiveLogin {
fn drop(&self) {
self.shutdown_handle.cancel();
self.shutdown_handle.shutdown();
}
}
@@ -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())
};
@@ -155,9 +154,10 @@ 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();
// 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,15 +180,19 @@ 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 {
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(join_err) => (
false,
Some(format!("failed to join login server thread: {join_err}")),
),
Err(_elapsed) => {
// Timeout: cancel server and report
shutdown_handle.shutdown();
(false, Some("Login timed out".to_string()))
}
};
let notification = LoginChatGptCompleteNotification {
login_id,

View File

@@ -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,10 +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 std::thread::JoinHandle;
use super::onboarding_screen::StepState;
// no additional imports
@@ -46,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<Arc<AtomicBool>>,
_login_wait_handle: Option<JoinHandle<()>>,
shutdown_handle: Option<ShutdownHandle>,
_login_wait_handle: Option<tokio::task::JoinHandle<()>>,
}
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.shutdown();
}
}
}
@@ -283,16 +281,21 @@ 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();
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 {
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)),
shutdown_handle: Some(shutdown_handle),
_login_wait_handle: Some(join_handle),
});
self.event_tx.send(AppEvent::RequestRedraw);
}
@@ -313,19 +316,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 {