diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 448a389851..3d1e8b9a25 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -85,6 +85,7 @@ use codex_core::resolve_installation_id; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_login::AuthManager; +use codex_login::AuthStores; use codex_protocol::protocol::SessionSource; pub use codex_rollout::StateDbHandle; pub use codex_state::log_db::LogDbLayer; @@ -348,8 +349,23 @@ impl InProcessClientHandle { /// the handle, so callers receive a ready-to-use runtime. If initialize fails, /// the runtime is shut down and an `InvalidData` error is returned. pub async fn start(args: InProcessStartArgs) -> IoResult { + start_with_optional_auth_stores(args, /*auth_stores*/ None).await +} + +/// Starts an in-process runtime backed by caller-provided auth stores. +pub async fn start_with_auth_stores( + args: InProcessStartArgs, + auth_stores: AuthStores, +) -> IoResult { + start_with_optional_auth_stores(args, Some(auth_stores)).await +} + +async fn start_with_optional_auth_stores( + args: InProcessStartArgs, + auth_stores: Option, +) -> IoResult { let initialize = args.initialize.clone(); - let client = start_uninitialized(args).await?; + let client = start_uninitialized(args, auth_stores).await?; let initialize_response = client .request(ClientRequest::Initialize { @@ -369,7 +385,10 @@ pub async fn start(args: InProcessStartArgs) -> IoResult Ok(client) } -async fn start_uninitialized(args: InProcessStartArgs) -> IoResult { +async fn start_uninitialized( + args: InProcessStartArgs, + auth_stores: Option, +) -> IoResult { let channel_capacity = args.channel_capacity.max(1); let installation_id = resolve_installation_id(&args.config.codex_home).await?; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); @@ -377,9 +396,22 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult(channel_capacity); - let auth_manager = - AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) - .await; + let auth_manager = match auth_stores { + Some(auth_stores) => { + AuthManager::shared_with_stores( + auth_stores, + args.enable_codex_api_key_env, + Some(args.config.chatgpt_base_url.clone()), + ) + .await + } + None => { + AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) + .await + } + }; + auth_manager + .set_forced_chatgpt_workspace_id(args.config.forced_chatgpt_workspace_id.clone()); let analytics_events_client = analytics_events_client_from_config(Arc::clone(&auth_manager), args.config.as_ref()); let outgoing_message_sender = Arc::new(OutgoingMessageSender::new( diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 36008ce0b9..4fc32a28e9 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -341,9 +341,7 @@ use codex_login::CLIENT_ID; use codex_login::CodexAuth; use codex_login::ServerOptions as LoginServerOptions; use codex_login::ShutdownHandle; -use codex_login::auth::login_with_chatgpt_auth_tokens; use codex_login::complete_device_code_login; -use codex_login::login_with_api_key; use codex_login::request_device_code; use codex_login::run_login_server; use codex_mcp::McpRuntimeContext; diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 8ebfc037e1..1b66e66fc3 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -273,11 +273,7 @@ impl AccountRequestProcessor { } } - match login_with_api_key( - &self.config.codex_home, - ¶ms.api_key, - self.config.cli_auth_credentials_store_mode, - ) { + match self.auth_manager.login_with_api_key(¶ms.api_key) { Ok(()) => { self.auth_manager.reload().await; Ok(()) @@ -326,6 +322,7 @@ impl AccountRequestProcessor { config.forced_chatgpt_workspace_id.clone(), config.cli_auth_credentials_store_mode, ) + .with_configured_auth_store(self.auth_manager.configured_auth_store()) }; #[cfg(debug_assertions)] let opts = { @@ -578,13 +575,13 @@ impl AccountRequestProcessor { ))); } - login_with_chatgpt_auth_tokens( - &self.config.codex_home, - &access_token, - &chatgpt_account_id, - chatgpt_plan_type.as_deref(), - ) - .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; + self.auth_manager + .login_with_chatgpt_auth_tokens( + &access_token, + &chatgpt_account_id, + chatgpt_plan_type.as_deref(), + ) + .map_err(|err| internal_error(format!("failed to set external auth: {err}")))?; self.auth_manager.reload().await; self.config_manager.replace_cloud_config_bundle_loader( self.auth_manager.clone(), diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index d581edc583..4180269a92 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1628,6 +1628,11 @@ impl AuthManager { Arc::new(Self::new_with_stores(stores, enable_codex_api_key_env, chatgpt_base_url).await) } + /// Returns the configured credential store for managed auth writers. + pub fn configured_auth_store(&self) -> Arc { + Arc::clone(&self.stores.configured) + } + /// Convenience constructor returning an `Arc` wrapper from resolved config. pub async fn shared_from_config( config: &impl AuthManagerConfig, diff --git a/codex-rs/login/src/device_code_auth.rs b/codex-rs/login/src/device_code_auth.rs index 4b9cb7c321..d31d3b3da6 100644 --- a/codex-rs/login/src/device_code_auth.rs +++ b/codex-rs/login/src/device_code_auth.rs @@ -210,13 +210,12 @@ pub async fn complete_device_code_login( return Err(io::Error::new(io::ErrorKind::PermissionDenied, message)); } - crate::server::persist_tokens_async( - &opts.codex_home, + crate::server::persist_tokens_to_store_async( + opts.configured_auth_store(), /*api_key*/ None, tokens.id_token, tokens.access_token, tokens.refresh_token, - opts.cli_auth_credentials_store_mode, ) .await } diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index b72bc946f2..1e180be8b9 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -11,6 +11,7 @@ //! This module therefore keeps the user-facing error path and the structured-log path separate. //! Returned `io::Error` values still carry the detail needed by CLI/browser callers, while //! structured logs only emit explicitly reviewed fields plus redacted URL/error values. +use std::fmt::Debug; use std::io::Cursor; use std::io::Read; use std::io::Write; @@ -24,10 +25,10 @@ use std::sync::LazyLock; use std::thread; use std::time::Duration; +use crate::auth::AuthCredentialStore; use crate::auth::AuthDotJson; -use crate::auth::load_auth_dot_json; +use crate::auth::AuthStores; use crate::auth::revoke_auth_tokens; -use crate::auth::save_auth; use crate::auth::should_revoke_auth_tokens; use crate::default_client::originator; use crate::pkce::PkceCodes; @@ -61,7 +62,7 @@ static LOGIN_ERROR_PAGE_TEMPLATE: LazyLock