From 5cf38e09ed0362d3dd70a4475babaaa92c9c38a2 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 3 Jun 2026 00:26:11 -0700 Subject: [PATCH] refactor auth credential store seam --- codex-rs/app-server/src/in_process.rs | 42 +++- codex-rs/app-server/src/request_processors.rs | 2 - .../request_processors/account_processor.rs | 20 +- codex-rs/login/src/auth/auth_tests.rs | 101 +++++++++- codex-rs/login/src/auth/manager.rs | 180 ++++++++++-------- codex-rs/login/src/auth/mod.rs | 2 + codex-rs/login/src/auth/storage.rs | 51 ++++- codex-rs/login/src/lib.rs | 2 + 8 files changed, 294 insertions(+), 106 deletions(-) diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 448a389851..1406d41dcc 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, 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..d3cb4db01f 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(()) @@ -578,13 +574,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/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 63ab3e0c42..2f8729bdb8 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -14,6 +14,7 @@ use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; use std::sync::Arc; +use std::sync::Mutex; use tempfile::TempDir; use tempfile::tempdir; use wiremock::Mock; @@ -26,6 +27,34 @@ const WORKSPACE_ID_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174000"; const WORKSPACE_ID_SECOND_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174001"; const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; +#[derive(Debug, Default)] +struct FakeAuthCredentialStore { + auth: Mutex>, +} + +impl FakeAuthCredentialStore { + fn with_auth(auth: AuthDotJson) -> Self { + Self { + auth: Mutex::new(Some(auth)), + } + } +} + +impl AuthCredentialStore for FakeAuthCredentialStore { + fn load(&self) -> std::io::Result> { + Ok(self.auth.lock().expect("fake store lock").clone()) + } + + fn save(&self, auth: &AuthDotJson) -> std::io::Result<()> { + *self.auth.lock().expect("fake store lock") = Some(auth.clone()); + Ok(()) + } + + fn delete(&self) -> std::io::Result { + Ok(self.auth.lock().expect("fake store lock").take().is_some()) + } +} + #[tokio::test] async fn refresh_without_id_token() { let codex_home = tempdir().unwrap(); @@ -295,6 +324,72 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { Ok(()) } +#[tokio::test] +async fn injected_auth_stores_drive_manager_reads_writes_and_logout() { + let initial_auth = AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-initial".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }; + let configured = Arc::new(FakeAuthCredentialStore::with_auth(initial_auth)); + let ephemeral = Arc::new(FakeAuthCredentialStore::default()); + let manager = AuthManager::shared_with_stores( + AuthStores { + configured: configured.clone(), + ephemeral: ephemeral.clone(), + }, + /*enable_codex_api_key_env*/ false, + /*chatgpt_base_url*/ None, + ) + .await; + + assert_eq!(manager.get_api_auth_mode(), Some(ApiAuthMode::ApiKey)); + + manager + .login_with_api_key("sk-updated") + .expect("api key login should save"); + assert_eq!( + configured.load().expect("configured load"), + Some(AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-updated".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }) + ); + + let access_token = fake_jwt_for_auth_file_params(&AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some(WORKSPACE_ID_ALLOWED.to_string()), + }) + .expect("external token"); + manager + .login_with_chatgpt_auth_tokens(&access_token, WORKSPACE_ID_ALLOWED, Some("pro")) + .expect("external token login should save"); + assert_eq!( + ephemeral.load().expect("ephemeral load"), + Some( + AuthDotJson::from_external_access_token( + &access_token, + WORKSPACE_ID_ALLOWED, + Some("pro"), + ) + .expect("external auth payload") + ) + ); + + manager.reload().await; + assert!(manager.is_external_chatgpt_auth_active()); + assert!(manager.logout().await.expect("logout should delete")); + assert_eq!(configured.load().expect("configured load"), None); + assert_eq!(ephemeral.load().expect("ephemeral load"), None); + assert_eq!(manager.auth_mode(), None); +} + #[tokio::test] async fn unauthorized_recovery_reports_mode_and_step_names() { let dir = tempdir().unwrap(); @@ -358,9 +453,11 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { updated_tokens.access_token = "new-access-token".to_string(); updated_tokens.refresh_token = "new-refresh-token".to_string(); let updated_auth = CodexAuth::from_auth_dot_json( - codex_home.path(), updated_auth_dot_json, - AuthCredentialsStoreMode::File, + create_auth_storage( + codex_home.path().to_path_buf(), + AuthCredentialsStoreMode::File, + ), /*chatgpt_base_url*/ None, ) .await diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 22e2c9f3b9..c1c878a870 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -28,8 +28,9 @@ use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::agent_identity::AgentIdentityAuth; pub use crate::auth::storage::AgentIdentityAuthRecord; +use crate::auth::storage::AuthCredentialStore; pub use crate::auth::storage::AuthDotJson; -use crate::auth::storage::AuthStorageBackend; +use crate::auth::storage::AuthStores; use crate::auth::storage::create_auth_storage; use crate::auth::util::try_parse_error_message; use crate::default_client::build_reqwest_client; @@ -69,7 +70,7 @@ pub struct ApiKeyAuth { #[derive(Debug, Clone)] pub struct ChatgptAuth { state: ChatgptAuthState, - storage: Arc, + storage: Arc, } #[derive(Debug, Clone)] @@ -199,9 +200,8 @@ impl From for std::io::Error { impl CodexAuth { async fn from_auth_dot_json( - codex_home: &Path, auth_dot_json: AuthDotJson, - auth_credentials_store_mode: AuthCredentialsStoreMode, + storage: Arc, chatgpt_base_url: Option<&str>, ) -> std::io::Result { let auth_mode = auth_dot_json.resolved_mode(); @@ -221,17 +221,13 @@ impl CodexAuth { return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await; } - let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); let state = ChatgptAuthState { auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), client, }; match auth_mode { - ApiAuthMode::Chatgpt => { - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); - Ok(Self::Chatgpt(ChatgptAuth { state, storage })) - } + ApiAuthMode::Chatgpt => Ok(Self::Chatgpt(ChatgptAuth { state, storage })), ApiAuthMode::ChatgptAuthTokens => { Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) } @@ -455,7 +451,7 @@ impl ChatgptAuth { self.current_auth_json().and_then(|auth| auth.tokens) } - fn storage(&self) -> &Arc { + fn storage(&self) -> &Arc { &self.storage } @@ -710,7 +706,8 @@ fn logout_with_message( ) -> std::io::Result<()> { // External auth tokens live in the ephemeral store, but persistent auth may still exist // from earlier logins. Clear both so a forced logout truly removes all active auth. - let removal_result = logout_all_stores(codex_home, auth_credentials_store_mode); + let stores = AuthStores::local(codex_home.to_path_buf(), auth_credentials_store_mode); + let removal_result = logout_all_stores(&stores); let error_message = match removal_result { Ok(_) => message, Err(err) => format!("{message}. Failed to remove auth.json: {err}"), @@ -718,15 +715,9 @@ fn logout_with_message( Err(std::io::Error::other(error_message)) } -fn logout_all_stores( - codex_home: &Path, - auth_credentials_store_mode: AuthCredentialsStoreMode, -) -> std::io::Result { - if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { - return logout(codex_home, AuthCredentialsStoreMode::Ephemeral); - } - let removed_ephemeral = logout(codex_home, AuthCredentialsStoreMode::Ephemeral)?; - let removed_managed = logout(codex_home, auth_credentials_store_mode)?; +fn logout_all_stores(stores: &AuthStores) -> std::io::Result { + let removed_ephemeral = stores.ephemeral.delete()?; + let removed_managed = stores.configured.delete()?; Ok(removed_ephemeral || removed_managed) } @@ -735,6 +726,15 @@ async fn load_auth( enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option<&str>, +) -> std::io::Result> { + let stores = AuthStores::local(codex_home.to_path_buf(), auth_credentials_store_mode); + load_auth_from_stores(&stores, enable_codex_api_key_env, chatgpt_base_url).await +} + +async fn load_auth_from_stores( + stores: &AuthStores, + enable_codex_api_key_env: bool, + chatgpt_base_url: Option<&str>, ) -> std::io::Result> { // API key via env var takes precedence over any other auth method. if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() { @@ -743,15 +743,10 @@ async fn load_auth( // External ChatGPT auth tokens live in the in-memory (ephemeral) store. Always check this // first so external auth takes precedence over any persisted credentials. - let ephemeral_storage = create_auth_storage( - codex_home.to_path_buf(), - AuthCredentialsStoreMode::Ephemeral, - ); - if let Some(auth_dot_json) = ephemeral_storage.load()? { + if let Some(auth_dot_json) = stores.ephemeral.load()? { let auth = CodexAuth::from_auth_dot_json( - codex_home, auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, + Arc::clone(&stores.ephemeral), chatgpt_base_url, ) .await?; @@ -759,7 +754,7 @@ async fn load_auth( } // If the caller explicitly requested ephemeral auth, there is no persisted fallback. - if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { + if Arc::ptr_eq(&stores.configured, &stores.ephemeral) { return Ok(None); } @@ -770,16 +765,14 @@ async fn load_auth( } // Fall back to the configured persistent store (file/keyring/auto) for managed auth. - let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); - let auth_dot_json = match storage.load()? { + let auth_dot_json = match stores.configured.load()? { Some(auth) => auth, None => return Ok(None), }; let auth = CodexAuth::from_auth_dot_json( - codex_home, auth_dot_json, - auth_credentials_store_mode, + Arc::clone(&stores.configured), chatgpt_base_url, ) .await?; @@ -788,7 +781,7 @@ async fn load_auth( // Persist refreshed tokens into auth storage and update last_refresh. fn persist_tokens( - storage: &Arc, + storage: &Arc, id_token: Option, access_token: Option, refresh_token: Option, @@ -988,17 +981,6 @@ impl AuthDotJson { } ApiAuthMode::Chatgpt } - - fn storage_mode( - &self, - auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> AuthCredentialsStoreMode { - if self.resolved_mode() == ApiAuthMode::ChatgptAuthTokens { - AuthCredentialsStoreMode::Ephemeral - } else { - auth_credentials_store_mode - } - } } /// Internal cached auth state. @@ -1243,20 +1225,19 @@ impl UnauthorizedRecovery { } } -/// Central manager providing a single source of truth for auth.json derived +/// Central manager providing a single source of truth for stored /// authentication data. It loads once (or on preference change) and then /// hands out cloned `CodexAuth` values so the rest of the program has a /// consistent snapshot. /// -/// External modifications to `auth.json` will NOT be observed until +/// External modifications to auth storage will NOT be observed until /// `reload()` is called explicitly. This matches the design goal of avoiding /// different parts of the program seeing inconsistent auth data mid‑run. pub struct AuthManager { - codex_home: PathBuf, inner: RwLock, auth_change_tx: watch::Sender, enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, + stores: AuthStores, forced_chatgpt_workspace_id: RwLock>>, chatgpt_base_url: Option, refresh_lock: Semaphore, @@ -1286,13 +1267,9 @@ pub trait AuthManagerConfig { impl Debug for AuthManager { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AuthManager") - .field("codex_home", &self.codex_home) .field("inner", &self.inner) .field("enable_codex_api_key_env", &self.enable_codex_api_key_env) - .field( - "auth_credentials_store_mode", - &self.auth_credentials_store_mode, - ) + .field("stores", &self.stores) .field( "forced_chatgpt_workspace_id", &self.forced_chatgpt_workspace_id, @@ -1314,10 +1291,22 @@ impl AuthManager { auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option, ) -> Self { - let managed_auth = load_auth( - &codex_home, + Self::new_with_stores( + AuthStores::local(codex_home, auth_credentials_store_mode), + enable_codex_api_key_env, + chatgpt_base_url, + ) + .await + } + + async fn new_with_stores( + stores: AuthStores, + enable_codex_api_key_env: bool, + chatgpt_base_url: Option, + ) -> Self { + let managed_auth = load_auth_from_stores( + &stores, enable_codex_api_key_env, - auth_credentials_store_mode, chatgpt_base_url.as_deref(), ) .await @@ -1325,14 +1314,13 @@ impl AuthManager { .flatten(); let (auth_change_tx, _auth_change_rx) = watch::channel(0); Self { - codex_home, inner: RwLock::new(CachedAuth { auth: managed_auth, permanent_refresh_failure: None, }), auth_change_tx, enable_codex_api_key_env, - auth_credentials_store_mode, + stores, forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url, refresh_lock: Semaphore::new(/*permits*/ 1), @@ -1349,11 +1337,13 @@ impl AuthManager { let (auth_change_tx, _auth_change_rx) = watch::channel(0); Arc::new(Self { - codex_home: PathBuf::from("non-existent"), inner: RwLock::new(cached), auth_change_tx, enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, + stores: AuthStores::local( + PathBuf::from("non-existent"), + AuthCredentialsStoreMode::File, + ), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), @@ -1369,11 +1359,10 @@ impl AuthManager { }; let (auth_change_tx, _auth_change_rx) = watch::channel(0); Arc::new(Self { - codex_home, inner: RwLock::new(cached), auth_change_tx, enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, + stores: AuthStores::local(codex_home, AuthCredentialsStoreMode::File), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), @@ -1384,14 +1373,16 @@ impl AuthManager { pub fn external_bearer_only(config: ModelProviderAuthInfo) -> Arc { let (auth_change_tx, _auth_change_rx) = watch::channel(0); Arc::new(Self { - codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { auth: None, permanent_refresh_failure: None, }), auth_change_tx, enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, + stores: AuthStores::local( + PathBuf::from("non-existent"), + AuthCredentialsStoreMode::File, + ), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, refresh_lock: Semaphore::new(/*permits*/ 1), @@ -1439,7 +1430,7 @@ impl AuthManager { self.auth_cached() } - /// Force a reload of the auth information from auth.json. Returns + /// Force a reload of the auth information from auth storage. Returns /// whether the auth value changed. pub async fn reload(&self) -> bool { tracing::info!("Reloading auth"); @@ -1531,10 +1522,9 @@ impl AuthManager { } async fn load_auth_from_storage(&self) -> Option { - load_auth( - &self.codex_home, + load_auth_from_stores( + &self.stores, self.enable_codex_api_key_env, - self.auth_credentials_store_mode, self.chatgpt_base_url.as_deref(), ) .await @@ -1621,6 +1611,15 @@ impl AuthManager { ) } + /// Creates a shared manager backed by caller-provided configured and ephemeral stores. + pub async fn shared_with_stores( + stores: AuthStores, + enable_codex_api_key_env: bool, + chatgpt_base_url: Option, + ) -> Arc { + Arc::new(Self::new_with_stores(stores, enable_codex_api_key_env, chatgpt_base_url).await) + } + /// Convenience constructor returning an `Arc` wrapper from resolved config. pub async fn shared_from_config( config: &impl AuthManagerConfig, @@ -1764,12 +1763,12 @@ impl AuthManager { result } - /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) - /// if a file was removed, Ok(false) if no auth file existed. On success, + /// Log out by deleting configured and ephemeral auth records. Returns Ok(true) + /// if a record was removed, Ok(false) if no auth record existed. On success, /// reloads the in‑memory auth cache so callers immediately observe the /// unauthenticated state. pub async fn logout(&self) -> std::io::Result { - let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; + let removed = logout_all_stores(&self.stores)?; // Always reload to clear any cached auth (even if file absent). self.reload().await; Ok(removed) @@ -1782,7 +1781,7 @@ impl AuthManager { if let Err(err) = revoke_auth_tokens(auth_dot_json.as_ref()).await { tracing::warn!("failed to revoke auth tokens during logout: {err}"); } - let result = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; + let result = logout_all_stores(&self.stores)?; // Always reload to clear any cached auth (even if file absent). self.reload().await; Ok(result) @@ -1876,12 +1875,10 @@ impl AuthManager { } let auth_dot_json = AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; - save_auth( - &self.codex_home, - &auth_dot_json, - AuthCredentialsStoreMode::Ephemeral, - ) - .map_err(RefreshTokenError::Transient)?; + self.stores + .ephemeral + .save(&auth_dot_json) + .map_err(RefreshTokenError::Transient)?; self.reload().await; Ok(()) } @@ -1906,6 +1903,33 @@ impl AuthManager { Ok(()) } + + /// Saves API key auth into the configured credential store. + pub fn login_with_api_key(&self, api_key: &str) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some(api_key.to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }; + self.stores.configured.save(&auth_dot_json) + } + + /// Saves externally managed ChatGPT tokens into the ephemeral credential store. + pub fn login_with_chatgpt_auth_tokens( + &self, + access_token: &str, + chatgpt_account_id: &str, + chatgpt_plan_type: Option<&str>, + ) -> std::io::Result<()> { + let auth_dot_json = AuthDotJson::from_external_access_token( + access_token, + chatgpt_account_id, + chatgpt_plan_type, + )?; + self.stores.ephemeral.save(&auth_dot_json) + } } #[cfg(test)] diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 167ccd3f75..353a4b5a4f 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -13,3 +13,5 @@ pub use error::RefreshTokenFailedReason; pub use manager::*; pub(crate) use revoke::revoke_auth_tokens; pub(crate) use revoke::should_revoke_auth_tokens; +pub use storage::AuthCredentialStore; +pub use storage::AuthStores; diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 3a1c8ae6aa..bc6c90cd70 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -94,12 +94,49 @@ pub(super) fn delete_file_if_exists(codex_home: &Path) -> std::io::Result } } -pub(super) trait AuthStorageBackend: Debug + Send + Sync { +/// Stores raw Codex auth payloads for one credential lane. +/// +/// Implementations should persist and return complete [`AuthDotJson`] records. +/// Cloud hosts may provide non-file-backed implementations, but callers should +/// keep configured and ephemeral auth stores separate. +pub trait AuthCredentialStore: Debug + Send + Sync + 'static { + /// Loads the stored auth payload, or `None` when this store is empty. fn load(&self) -> std::io::Result>; + /// Replaces the stored auth payload. fn save(&self, auth: &AuthDotJson) -> std::io::Result<()>; + /// Removes the stored auth payload and reports whether one existed. fn delete(&self) -> std::io::Result; } +/// Configured and ephemeral auth stores used by [`crate::auth::AuthManager`]. +/// +/// `configured` owns managed auth such as API keys and first-party ChatGPT +/// login. `ephemeral` owns externally supplied ChatGPT auth tokens that should +/// not be persisted to the configured store. +#[derive(Clone, Debug)] +pub struct AuthStores { + /// Store for managed auth selected by runtime configuration. + pub configured: Arc, + /// Store for externally supplied ChatGPT auth tokens. + pub ephemeral: Arc, +} + +impl AuthStores { + /// Builds the existing Codex Home backed auth stores for local runtimes. + pub fn local(codex_home: PathBuf, mode: AuthCredentialsStoreMode) -> Self { + let configured = create_auth_storage(codex_home.clone(), mode); + let ephemeral = if mode == AuthCredentialsStoreMode::Ephemeral { + Arc::clone(&configured) + } else { + create_auth_storage(codex_home, AuthCredentialsStoreMode::Ephemeral) + }; + Self { + configured, + ephemeral, + } + } +} + #[derive(Clone, Debug)] pub(super) struct FileAuthStorage { codex_home: PathBuf, @@ -122,7 +159,7 @@ impl FileAuthStorage { } } -impl AuthStorageBackend for FileAuthStorage { +impl AuthCredentialStore for FileAuthStorage { fn load(&self) -> std::io::Result> { let auth_file = get_auth_file(&self.codex_home); let auth_dot_json = match self.try_read_auth_json(&auth_file) { @@ -217,7 +254,7 @@ impl KeyringAuthStorage { } } -impl AuthStorageBackend for KeyringAuthStorage { +impl AuthCredentialStore for KeyringAuthStorage { fn load(&self) -> std::io::Result> { let key = compute_store_key(&self.codex_home)?; self.load_from_keyring(&key) @@ -262,7 +299,7 @@ impl AutoAuthStorage { } } -impl AuthStorageBackend for AutoAuthStorage { +impl AuthCredentialStore for AutoAuthStorage { fn load(&self) -> std::io::Result> { match self.keyring_storage.load() { Ok(Some(auth)) => Ok(Some(auth)), @@ -316,7 +353,7 @@ impl EphemeralAuthStorage { } } -impl AuthStorageBackend for EphemeralAuthStorage { +impl AuthCredentialStore for EphemeralAuthStorage { fn load(&self) -> std::io::Result> { self.with_store(|store, key| Ok(store.get(&key).cloned())) } @@ -336,7 +373,7 @@ impl AuthStorageBackend for EphemeralAuthStorage { pub(super) fn create_auth_storage( codex_home: PathBuf, mode: AuthCredentialsStoreMode, -) -> Arc { +) -> Arc { let keyring_store: Arc = Arc::new(DefaultKeyringStore); create_auth_storage_with_keyring_store(codex_home, mode, keyring_store) } @@ -345,7 +382,7 @@ fn create_auth_storage_with_keyring_store( codex_home: PathBuf, mode: AuthCredentialsStoreMode, keyring_store: Arc, -) -> Arc { +) -> Arc { match mode { AuthCredentialsStoreMode::File => Arc::new(FileAuthStorage::new(codex_home)), AuthCredentialsStoreMode::Keyring => { diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 990cf8b80e..fde253312b 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -18,9 +18,11 @@ pub use server::ShutdownHandle; pub use server::run_login_server; pub use auth::AuthConfig; +pub use auth::AuthCredentialStore; pub use auth::AuthDotJson; pub use auth::AuthManager; pub use auth::AuthManagerConfig; +pub use auth::AuthStores; pub use auth::CLIENT_ID; pub use auth::CODEX_ACCESS_TOKEN_ENV_VAR; pub use auth::CODEX_API_KEY_ENV_VAR;