From 377ab0c77cd8609637f2fa6d7c0230fa69ce3acc Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 30 Jan 2026 09:33:23 -0800 Subject: [PATCH] feat: refactor CodexAuth so invalid state cannot be represented (#10208) Previously, `CodexAuth` was defined as follows: https://github.com/openai/codex/blob/d550fbf41afc09d7d7b5ac813aea38de07b2a73f/codex-rs/core/src/auth.rs#L39-L46 But if you looked at its constructors, we had creation for `AuthMode::ApiKey` where `storage` was built using a nonsensical path (`PathBuf::new()`) and `auth_dot_json` was `None`: https://github.com/openai/codex/blob/d550fbf41afc09d7d7b5ac813aea38de07b2a73f/codex-rs/core/src/auth.rs#L212-L220 By comparison, when `AuthMode::ChatGPT` was used, `api_key` was always `None`: https://github.com/openai/codex/blob/d550fbf41afc09d7d7b5ac813aea38de07b2a73f/codex-rs/core/src/auth.rs#L665-L671 https://github.com/openai/codex/pull/10012 took things further because it introduced a new `ChatgptAuthTokens` variant to `AuthMode`, which is important in when invoking `account/login/start` via the app server, but most logic _internal_ to the app server should just reason about two `AuthMode` variants: `ApiKey` and `ChatGPT`. This PR tries to clean things up as follows: - `LoginAccountParams` and `AuthMode` in `codex-rs/app-server-protocol/` both continue to have the `ChatgptAuthTokens` variant, though it is used exclusively for the on-the-wire messaging. - `codex-rs/core/src/auth.rs` now has its own `AuthMode` enum, which only has two variants: `ApiKey` and `ChatGPT`. - `CodexAuth` has been changed from a struct to an enum. It is a disjoint union where each variant (`ApiKey`, `ChatGpt`, and `ChatGptAuthTokens`) have only the associated fields that make sense for that variant. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/10208). * #10224 * __->__ #10208 --- .../src/protocol/common.rs | 4 - .../app-server/src/codex_message_processor.rs | 39 +- codex-rs/cli/src/login.rs | 2 +- codex-rs/core/src/auth.rs | 340 ++++++++++++------ codex-rs/core/src/client.rs | 14 +- codex-rs/core/src/codex.rs | 2 +- codex-rs/core/src/model_provider_info.rs | 9 +- codex-rs/core/src/models_manager/manager.rs | 14 +- .../core/src/models_manager/model_presets.rs | 2 +- codex-rs/tui/src/app.rs | 2 +- codex-rs/tui/src/chatwidget.rs | 13 +- codex-rs/tui/src/lib.rs | 2 +- codex-rs/tui/src/status/helpers.rs | 8 +- 13 files changed, 286 insertions(+), 165 deletions(-) diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 2e8778bdc..67736374e 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -24,10 +24,6 @@ impl GitSha { } /// Authentication mode for OpenAI-backed providers. -/// -/// This is used internally to determine the base URL for generating responses, and to gate -/// ChatGPT-only behaviors like rate limits and available models (as opposed to API key-based -/// auth). #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, Display, JsonSchema, TS)] #[serde(rename_all = "lowercase")] pub enum AuthMode { diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 3817bce3a..c02f3c736 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -133,6 +133,7 @@ use codex_app_server_protocol::build_turns_from_event_msgs; use codex_backend_client::Client as BackendClient; use codex_chatgpt::connectors; use codex_core::AuthManager; +use codex_core::CodexAuth; use codex_core::CodexThread; use codex_core::Cursor as RolloutCursor; use codex_core::InitialHistory; @@ -686,7 +687,11 @@ impl CodexMessageProcessor { .await; let payload = AuthStatusChangeNotification { - auth_method: self.auth_manager.auth_cached().map(|auth| auth.mode), + auth_method: self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode), }; self.outgoing .send_server_notification(ServerNotification::AuthStatusChange(payload)) @@ -716,7 +721,11 @@ impl CodexMessageProcessor { .await; let payload_v2 = AccountUpdatedNotification { - auth_mode: self.auth_manager.auth_cached().map(|auth| auth.mode), + auth_mode: self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode), }; self.outgoing .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) @@ -812,7 +821,10 @@ impl CodexMessageProcessor { auth_manager.reload(); // Notify clients with the actual current auth mode. - let current_auth_method = auth_manager.auth_cached().map(|a| a.mode); + let current_auth_method = auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode); let payload = AuthStatusChangeNotification { auth_method: current_auth_method, }; @@ -902,7 +914,10 @@ impl CodexMessageProcessor { auth_manager.reload(); // Notify clients with the actual current auth mode. - let current_auth_method = auth_manager.auth_cached().map(|a| a.mode); + let current_auth_method = auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode); let payload_v2 = AccountUpdatedNotification { auth_mode: current_auth_method, }; @@ -1106,7 +1121,11 @@ impl CodexMessageProcessor { } // Reflect the current auth method after logout (likely None). - Ok(self.auth_manager.auth_cached().map(|auth| auth.mode)) + Ok(self + .auth_manager + .auth_cached() + .as_ref() + .map(CodexAuth::api_auth_mode)) } async fn logout_v1(&mut self, request_id: RequestId) { @@ -1178,7 +1197,7 @@ impl CodexMessageProcessor { } else { match self.auth_manager.auth().await { Some(auth) => { - let auth_mode = auth.mode; + let auth_mode = auth.api_auth_mode(); let (reported_auth_method, token_opt) = match auth.get_token() { Ok(token) if !token.is_empty() => { let tok = if include_token { Some(token) } else { None }; @@ -1225,9 +1244,9 @@ impl CodexMessageProcessor { } let account = match self.auth_manager.auth_cached() { - Some(auth) => Some(match auth.mode { - AuthMode::ApiKey => Account::ApiKey {}, - AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { + Some(auth) => Some(match auth { + CodexAuth::ApiKey(_) => Account::ApiKey {}, + CodexAuth::ChatGpt(_) | CodexAuth::ChatGptAuthTokens(_) => { let email = auth.get_account_email(); let plan_type = auth.account_plan_type(); @@ -1286,7 +1305,7 @@ impl CodexMessageProcessor { }); }; - if !matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) { + if !auth.is_chatgpt_auth() { return Err(JSONRPCErrorError { code: INVALID_REQUEST_ERROR_CODE, message: "chatgpt authentication required to read rate limits".to_string(), diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 0f4e4eb39..cee002f33 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -225,7 +225,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode) { - Ok(Some(auth)) => match auth.mode { + Ok(Some(auth)) => match auth.api_auth_mode() { AuthMode::ApiKey => match auth.get_token() { Ok(api_key) => { eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); diff --git a/codex-rs/core/src/auth.rs b/codex-rs/core/src/auth.rs index d4520dc85..4bddbe303 100644 --- a/codex-rs/core/src/auth.rs +++ b/codex-rs/core/src/auth.rs @@ -13,8 +13,9 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex; +use std::sync::RwLock; -use codex_app_server_protocol::AuthMode; +use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; pub use crate::auth::storage::AuthCredentialsStoreMode; @@ -35,19 +36,50 @@ use codex_protocol::account::PlanType as AccountPlanType; use serde_json::Value; use thiserror::Error; -#[derive(Debug, Clone)] -pub struct CodexAuth { - pub mode: AuthMode, +/// Account type for the current user. +/// +/// This is used internally to determine the base URL for generating responses, +/// and to gate ChatGPT-only behaviors like rate limits and available models (as +/// opposed to API key-based auth). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AuthMode { + ApiKey, + ChatGPT, +} - pub(crate) api_key: Option, - pub(crate) auth_dot_json: Arc>>, +/// Authentication mechanism used by the current user. +#[derive(Debug, Clone)] +pub enum CodexAuth { + ApiKey(ApiKeyAuth), + ChatGpt(ChatGptAuth), + ChatGptAuthTokens(ChatGptAuthTokens), +} + +#[derive(Debug, Clone)] +pub struct ApiKeyAuth { + api_key: String, +} + +#[derive(Debug, Clone)] +pub struct ChatGptAuth { + state: ChatGptAuthState, storage: Arc, - pub(crate) client: CodexHttpClient, +} + +#[derive(Debug, Clone)] +pub struct ChatGptAuthTokens { + state: ChatGptAuthState, +} + +#[derive(Debug, Clone)] +struct ChatGptAuthState { + auth_dot_json: Arc>>, + client: CodexHttpClient, } impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { - self.mode == other.mode + self.api_auth_mode() == other.api_auth_mode() } } @@ -114,14 +146,78 @@ impl From for std::io::Error { } impl CodexAuth { + fn from_auth_dot_json( + codex_home: &Path, + auth_dot_json: AuthDotJson, + auth_credentials_store_mode: AuthCredentialsStoreMode, + client: CodexHttpClient, + ) -> std::io::Result { + let auth_mode = auth_dot_json.resolved_mode(); + if auth_mode == ApiAuthMode::ApiKey { + let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { + return Err(std::io::Error::other("API key auth is missing a key.")); + }; + return Ok(CodexAuth::from_api_key_with_client(api_key, client)); + } + + 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::ChatgptAuthTokens => { + Ok(Self::ChatGptAuthTokens(ChatGptAuthTokens { state })) + } + ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), + } + } + /// Loads the available auth information from auth storage. pub fn from_auth_storage( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, - ) -> std::io::Result> { + ) -> std::io::Result> { load_auth(codex_home, false, auth_credentials_store_mode) } + pub fn internal_auth_mode(&self) -> AuthMode { + match self { + Self::ApiKey(_) => AuthMode::ApiKey, + Self::ChatGpt(_) | Self::ChatGptAuthTokens(_) => AuthMode::ChatGPT, + } + } + + pub fn api_auth_mode(&self) -> ApiAuthMode { + match self { + Self::ApiKey(_) => ApiAuthMode::ApiKey, + Self::ChatGpt(_) => ApiAuthMode::ChatGPT, + Self::ChatGptAuthTokens(_) => ApiAuthMode::ChatgptAuthTokens, + } + } + + pub fn is_chatgpt_auth(&self) -> bool { + self.internal_auth_mode() == AuthMode::ChatGPT + } + + pub fn is_external_chatgpt_tokens(&self) -> bool { + matches!(self, Self::ChatGptAuthTokens(_)) + } + + /// Returns `None` is `is_internal_auth_mode() != AuthMode::ApiKey`. + pub fn api_key(&self) -> Option<&str> { + match self { + Self::ApiKey(auth) => Some(auth.api_key.as_str()), + Self::ChatGpt(_) | Self::ChatGptAuthTokens(_) => None, + } + } + + /// Returns `Err` if `is_chatgpt_auth()` is false. pub fn get_token_data(&self) -> Result { let auth_dot_json: Option = self.get_current_auth_json(); match auth_dot_json { @@ -134,20 +230,23 @@ impl CodexAuth { } } + /// Returns the token string used for bearer authentication. pub fn get_token(&self) -> Result { - match self.mode { - AuthMode::ApiKey => Ok(self.api_key.clone().unwrap_or_default()), - AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { + match self { + Self::ApiKey(auth) => Ok(auth.api_key.clone()), + Self::ChatGpt(_) | Self::ChatGptAuthTokens(_) => { let access_token = self.get_token_data()?.access_token; Ok(access_token) } } } + /// Returns `None` if `is_chatgpt_auth()` is false. pub fn get_account_id(&self) -> Option { self.get_current_token_data().and_then(|t| t.account_id) } + /// Returns `None` if `is_chatgpt_auth()` is false. pub fn get_account_email(&self) -> Option { self.get_current_token_data().and_then(|t| t.id_token.email) } @@ -176,11 +275,18 @@ impl CodexAuth { }) } + /// Returns `None` if `is_chatgpt_auth()` is false. fn get_current_auth_json(&self) -> Option { + let state = match self { + Self::ChatGpt(auth) => &auth.state, + Self::ChatGptAuthTokens(auth) => &auth.state, + Self::ApiKey(_) => return None, + }; #[expect(clippy::unwrap_used)] - self.auth_dot_json.lock().unwrap().clone() + state.auth_dot_json.lock().unwrap().clone() } + /// Returns `None` if `is_chatgpt_auth()` is false. fn get_current_token_data(&self) -> Option { self.get_current_auth_json().and_then(|t| t.tokens) } @@ -188,7 +294,7 @@ impl CodexAuth { /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { - auth_mode: Some(AuthMode::ChatGPT), + auth_mode: Some(ApiAuthMode::ChatGPT), openai_api_key: None, tokens: Some(TokenData { id_token: Default::default(), @@ -199,24 +305,19 @@ impl CodexAuth { last_refresh: Some(Utc::now()), }; - let auth_dot_json = Arc::new(Mutex::new(Some(auth_dot_json))); - Self { - api_key: None, - mode: AuthMode::ChatGPT, - storage: create_auth_storage(PathBuf::new(), AuthCredentialsStoreMode::File), - auth_dot_json, - client: crate::default_client::create_client(), - } + let client = crate::default_client::create_client(); + let state = ChatGptAuthState { + auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), + client, + }; + let storage = create_auth_storage(PathBuf::new(), AuthCredentialsStoreMode::File); + Self::ChatGpt(ChatGptAuth { state, storage }) } - fn from_api_key_with_client(api_key: &str, client: CodexHttpClient) -> Self { - Self { - api_key: Some(api_key.to_owned()), - mode: AuthMode::ApiKey, - storage: create_auth_storage(PathBuf::new(), AuthCredentialsStoreMode::File), - auth_dot_json: Arc::new(Mutex::new(None)), - client, - } + fn from_api_key_with_client(api_key: &str, _client: CodexHttpClient) -> Self { + Self::ApiKey(ApiKeyAuth { + api_key: api_key.to_owned(), + }) } pub fn from_api_key(api_key: &str) -> Self { @@ -224,6 +325,25 @@ impl CodexAuth { } } +impl ChatGptAuth { + fn current_auth_json(&self) -> Option { + #[expect(clippy::unwrap_used)] + self.state.auth_dot_json.lock().unwrap().clone() + } + + fn current_token_data(&self) -> Option { + self.current_auth_json().and_then(|auth| auth.tokens) + } + + fn storage(&self) -> &Arc { + &self.storage + } + + fn client(&self) -> &CodexHttpClient { + &self.state.client + } +} + pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; @@ -258,7 +378,7 @@ pub fn login_with_api_key( auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result<()> { let auth_dot_json = AuthDotJson { - auth_mode: Some(AuthMode::ApiKey), + auth_mode: Some(ApiAuthMode::ApiKey), openai_api_key: Some(api_key.to_string()), tokens: None, last_refresh: None, @@ -314,10 +434,10 @@ pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { }; if let Some(required_method) = config.forced_login_method { - let method_violation = match (required_method, auth.mode) { + let method_violation = match (required_method, auth.internal_auth_mode()) { (ForcedLoginMethod::Api, AuthMode::ApiKey) => None, - (ForcedLoginMethod::Chatgpt, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) => None, - (ForcedLoginMethod::Api, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) => Some( + (ForcedLoginMethod::Chatgpt, AuthMode::ChatGPT) => None, + (ForcedLoginMethod::Api, AuthMode::ChatGPT) => Some( "API key login is required, but ChatGPT is currently being used. Logging out." .to_string(), ), @@ -337,7 +457,7 @@ pub fn enforce_login_restrictions(config: &Config) -> std::io::Result<()> { } if let Some(expected_account_id) = config.forced_chatgpt_workspace_id.as_deref() { - if !matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) { + if !auth.is_chatgpt_auth() { return Ok(()); } @@ -607,7 +727,7 @@ impl AuthDotJson { }; Self { - auth_mode: Some(AuthMode::ChatgptAuthTokens), + auth_mode: Some(ApiAuthMode::ChatgptAuthTokens), openai_api_key: None, tokens: Some(tokens), last_refresh: Some(Utc::now()), @@ -623,21 +743,21 @@ impl AuthDotJson { Ok(Self::from_external_tokens(&external, id_token_info)) } - fn resolved_mode(&self) -> AuthMode { + fn resolved_mode(&self) -> ApiAuthMode { if let Some(mode) = self.auth_mode { return mode; } if self.openai_api_key.is_some() { - return AuthMode::ApiKey; + return ApiAuthMode::ApiKey; } - AuthMode::ChatGPT + ApiAuthMode::ChatGPT } fn storage_mode( &self, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> AuthCredentialsStoreMode { - if self.resolved_mode() == AuthMode::ChatgptAuthTokens { + if self.resolved_mode() == ApiAuthMode::ChatgptAuthTokens { AuthCredentialsStoreMode::Ephemeral } else { auth_credentials_store_mode @@ -645,35 +765,6 @@ impl AuthDotJson { } } -impl CodexAuth { - fn from_auth_dot_json( - codex_home: &Path, - auth_dot_json: AuthDotJson, - auth_credentials_store_mode: AuthCredentialsStoreMode, - client: CodexHttpClient, - ) -> std::io::Result { - let auth_mode = auth_dot_json.resolved_mode(); - if auth_mode == AuthMode::ApiKey { - let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { - return Err(std::io::Error::other("API key auth is missing a key.")); - }; - return Ok(CodexAuth::from_api_key_with_client(api_key, client)); - } - - let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); - Ok(Self { - api_key: None, - mode: auth_mode, - storage, - auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, - }) - } -} - -use std::sync::RwLock; - /// Internal cached auth state. #[derive(Clone)] struct CachedAuth { @@ -685,7 +776,10 @@ struct CachedAuth { impl Debug for CachedAuth { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CachedAuth") - .field("auth_mode", &self.auth.as_ref().map(|auth| auth.mode)) + .field( + "auth_mode", + &self.auth.as_ref().map(CodexAuth::api_auth_mode), + ) .field( "external_refresher", &self.external_refresher.as_ref().map(|_| "present"), @@ -736,11 +830,13 @@ impl UnauthorizedRecovery { fn new(manager: Arc) -> Self { let cached_auth = manager.auth_cached(); let expected_account_id = cached_auth.as_ref().and_then(CodexAuth::get_account_id); - let mode = match cached_auth { - Some(auth) if auth.mode == AuthMode::ChatgptAuthTokens => { - UnauthorizedRecoveryMode::External - } - _ => UnauthorizedRecoveryMode::Managed, + let mode = if cached_auth + .as_ref() + .is_some_and(CodexAuth::is_external_chatgpt_tokens) + { + UnauthorizedRecoveryMode::External + } else { + UnauthorizedRecoveryMode::Managed }; let step = match mode { UnauthorizedRecoveryMode::Managed => UnauthorizedRecoveryStep::Reload, @@ -755,9 +851,12 @@ impl UnauthorizedRecovery { } pub fn has_next(&self) -> bool { - if !self.manager.auth_cached().is_some_and(|auth| { - matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) - }) { + if !self + .manager + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_chatgpt_auth) + { return false; } @@ -996,7 +1095,8 @@ impl AuthManager { pub fn is_external_auth_active(&self) -> bool { self.auth_cached() - .is_some_and(|auth| auth.mode == AuthMode::ChatgptAuthTokens) + .as_ref() + .is_some_and(CodexAuth::is_external_chatgpt_tokens) } /// Convenience constructor returning an `Arc` wrapper. @@ -1026,18 +1126,25 @@ impl AuthManager { Some(auth) => auth, None => return Ok(()), }; - if auth.mode == AuthMode::ChatgptAuthTokens { - return self - .refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) - .await; + match auth { + CodexAuth::ChatGptAuthTokens(_) => { + self.refresh_external_auth(ExternalAuthRefreshReason::Unauthorized) + .await + } + CodexAuth::ChatGpt(chatgpt_auth) => { + let token_data = chatgpt_auth.current_token_data().ok_or_else(|| { + RefreshTokenError::Transient(std::io::Error::other( + "Token data is not available.", + )) + })?; + self.refresh_tokens(&chatgpt_auth, token_data.refresh_token) + .await?; + // Reload to pick up persisted changes. + self.reload(); + Ok(()) + } + CodexAuth::ApiKey(_) => Ok(()), } - let token_data = auth.get_current_token_data().ok_or_else(|| { - RefreshTokenError::Transient(std::io::Error::other("Token data is not available.")) - })?; - self.refresh_tokens(&auth, token_data.refresh_token).await?; - // Reload to pick up persisted changes. - self.reload(); - Ok(()) } /// Log out by deleting the on‑disk auth.json (if present). Returns Ok(true) @@ -1051,16 +1158,23 @@ impl AuthManager { Ok(removed) } - pub fn get_auth_mode(&self) -> Option { - self.auth_cached().map(|a| a.mode) + pub fn get_auth_mode(&self) -> Option { + self.auth_cached().as_ref().map(CodexAuth::api_auth_mode) + } + + pub fn get_internal_auth_mode(&self) -> Option { + self.auth_cached() + .as_ref() + .map(CodexAuth::internal_auth_mode) } async fn refresh_if_stale(&self, auth: &CodexAuth) -> Result { - if auth.mode != AuthMode::ChatGPT { - return Ok(false); - } + let chatgpt_auth = match auth { + CodexAuth::ChatGpt(chatgpt_auth) => chatgpt_auth, + _ => return Ok(false), + }; - let auth_dot_json = match auth.get_current_auth_json() { + let auth_dot_json = match chatgpt_auth.current_auth_json() { Some(auth_dot_json) => auth_dot_json, None => return Ok(false), }; @@ -1075,7 +1189,8 @@ impl AuthManager { if last_refresh >= Utc::now() - chrono::Duration::days(TOKEN_REFRESH_INTERVAL) { return Ok(false); } - self.refresh_tokens(auth, tokens.refresh_token).await?; + self.refresh_tokens(chatgpt_auth, tokens.refresh_token) + .await?; self.reload(); Ok(true) } @@ -1135,13 +1250,13 @@ impl AuthManager { async fn refresh_tokens( &self, - auth: &CodexAuth, + auth: &ChatGptAuth, refresh_token: String, ) -> Result<(), RefreshTokenError> { - let refresh_response = try_refresh_token(refresh_token, &auth.client).await?; + let refresh_response = try_refresh_token(refresh_token, auth.client()).await?; update_tokens( - &auth.storage, + auth.storage(), refresh_response.id_token, refresh_response.access_token, refresh_response.refresh_token, @@ -1254,26 +1369,21 @@ mod tests { ) .expect("failed to write auth file"); - let CodexAuth { - api_key, - mode, - auth_dot_json, - storage: _, - .. - } = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) + let auth = super::load_auth(codex_home.path(), false, AuthCredentialsStoreMode::File) .unwrap() .unwrap(); - assert_eq!(None, api_key); - assert_eq!(AuthMode::ChatGPT, mode); + assert_eq!(None, auth.api_key()); + assert_eq!(AuthMode::ChatGPT, auth.internal_auth_mode()); - let guard = auth_dot_json.lock().unwrap(); - let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); + let auth_dot_json = auth + .get_current_auth_json() + .expect("AuthDotJson should exist"); let last_refresh = auth_dot_json .last_refresh .expect("last_refresh should be recorded"); assert_eq!( - &AuthDotJson { + AuthDotJson { auth_mode: None, openai_api_key: None, tokens: Some(TokenData { @@ -1308,8 +1418,8 @@ mod tests { let auth = super::load_auth(dir.path(), false, AuthCredentialsStoreMode::File) .unwrap() .unwrap(); - assert_eq!(auth.mode, AuthMode::ApiKey); - assert_eq!(auth.api_key, Some("sk-test-key".to_string())); + assert_eq!(auth.internal_auth_mode(), AuthMode::ApiKey); + assert_eq!(auth.api_key(), Some("sk-test-key")); assert!(auth.get_token_data().is_err()); } @@ -1318,7 +1428,7 @@ mod tests { fn logout_removes_auth_file() -> Result<(), std::io::Error> { let dir = tempdir()?; let auth_dot_json = AuthDotJson { - auth_mode: Some(AuthMode::ApiKey), + auth_mode: Some(ApiAuthMode::ApiKey), openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 0f05f58c4..f01c145e5 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -27,7 +27,6 @@ use codex_api::common::ResponsesWsRequest; use codex_api::create_text_param_for_request; use codex_api::error::ApiError; use codex_api::requests::responses::Compression; -use codex_app_server_protocol::AuthMode; use codex_otel::OtelManager; use codex_protocol::ThreadId; @@ -50,6 +49,7 @@ use tokio::sync::mpsc; use tracing::warn; use crate::AuthManager; +use crate::auth::CodexAuth; use crate::auth::RefreshTokenError; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; @@ -220,7 +220,7 @@ impl ModelClient { let api_provider = self .state .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = self.build_request_telemetry(); @@ -469,9 +469,7 @@ impl ModelClientSession { .config .features .enabled(Feature::EnableRequestCompression) - && auth.is_some_and(|auth| { - matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) - }) + && auth.is_some_and(CodexAuth::is_chatgpt_auth) && self.state.provider.is_openai() { Compression::Zstd @@ -509,7 +507,7 @@ impl ModelClientSession { let api_provider = self .state .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); @@ -565,7 +563,7 @@ impl ModelClientSession { let api_provider = self .state .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); let (request_telemetry, sse_telemetry) = self.build_streaming_telemetry(); @@ -611,7 +609,7 @@ impl ModelClientSession { let api_provider = self .state .provider - .to_api_provider(auth.as_ref().map(|a| a.mode))?; + .to_api_provider(auth.as_ref().map(CodexAuth::internal_auth_mode))?; let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; let compression = self.responses_request_compression(auth.as_ref()); diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 7eefd04ec..324f41a04 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -838,7 +838,7 @@ impl Session { session_configuration.collaboration_mode.model(), auth.and_then(CodexAuth::get_account_id), auth.and_then(CodexAuth::get_account_email), - auth.map(|a| a.mode), + auth.map(CodexAuth::api_auth_mode), config.otel.log_user_prompt, terminal::user_agent(), session_configuration.session_source.clone(), diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 0f3551480..0f5a8643a 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -5,10 +5,11 @@ //! 2. User-defined entries inside `~/.codex/config.toml` under the `model_providers` //! key. These override or extend the defaults at runtime. +use crate::auth::AuthMode; +use crate::error::EnvVarError; use codex_api::Provider as ApiProvider; use codex_api::WireApi as ApiWireApi; use codex_api::provider::RetryConfig as ApiRetryConfig; -use codex_app_server_protocol::AuthMode; use http::HeaderMap; use http::header::HeaderName; use http::header::HeaderValue; @@ -19,7 +20,6 @@ use std::collections::HashMap; use std::env::VarError; use std::time::Duration; -use crate::error::EnvVarError; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; @@ -137,10 +137,7 @@ impl ModelProviderInfo { &self, auth_mode: Option, ) -> crate::error::Result { - let default_base_url = if matches!( - auth_mode, - Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) - ) { + let default_base_url = if matches!(auth_mode, Some(AuthMode::ChatGPT)) { "https://chatgpt.com/backend-api/codex" } else { "https://api.openai.com/v1" diff --git a/codex-rs/core/src/models_manager/manager.rs b/codex-rs/core/src/models_manager/manager.rs index c993032ea..40f9c8ccd 100644 --- a/codex-rs/core/src/models_manager/manager.rs +++ b/codex-rs/core/src/models_manager/manager.rs @@ -2,6 +2,7 @@ use super::cache::ModelsCacheManager; use crate::api_bridge::auth_provider_from_auth; use crate::api_bridge::map_api_error; use crate::auth::AuthManager; +use crate::auth::AuthMode; use crate::config::Config; use crate::default_client::build_reqwest_client; use crate::error::CodexErr; @@ -13,7 +14,6 @@ use crate::models_manager::model_info; use crate::models_manager::model_presets::builtin_model_presets; use codex_api::ModelsClient; use codex_api::ReqwestTransport; -use codex_app_server_protocol::AuthMode; use codex_protocol::config_types::CollaborationModeMask; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelPreset; @@ -61,7 +61,7 @@ impl ModelsManager { let cache_path = codex_home.join(MODEL_CACHE_FILE); let cache_manager = ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL); Self { - local_models: builtin_model_presets(auth_manager.get_auth_mode()), + local_models: builtin_model_presets(auth_manager.get_internal_auth_mode()), remote_models: RwLock::new(Self::load_remote_models_from_file().unwrap_or_default()), auth_manager, etag: RwLock::new(None), @@ -175,7 +175,7 @@ impl ModelsManager { refresh_strategy: RefreshStrategy, ) -> CoreResult<()> { if !config.features.enabled(Feature::RemoteModels) - || self.auth_manager.get_auth_mode() == Some(AuthMode::ApiKey) + || self.auth_manager.get_internal_auth_mode() == Some(AuthMode::ApiKey) { return Ok(()); } @@ -204,7 +204,7 @@ impl ModelsManager { let _timer = codex_otel::start_global_timer("codex.remote_models.fetch_update.duration_ms", &[]); let auth = self.auth_manager.auth().await; - let auth_mode = self.auth_manager.get_auth_mode(); + let auth_mode = self.auth_manager.get_internal_auth_mode(); let api_provider = self.provider.to_api_provider(auth_mode)?; let api_auth = auth_provider_from_auth(auth.clone(), &self.provider)?; let transport = ReqwestTransport::new(build_reqwest_client()); @@ -273,8 +273,8 @@ impl ModelsManager { let existing_presets = self.local_models.clone(); let mut merged_presets = ModelPreset::merge(remote_presets, existing_presets); let chatgpt_mode = matches!( - self.auth_manager.get_auth_mode(), - Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + self.auth_manager.get_internal_auth_mode(), + Some(AuthMode::ChatGPT) ); merged_presets = ModelPreset::filter_by_auth(merged_presets, chatgpt_mode); @@ -319,7 +319,7 @@ impl ModelsManager { let cache_path = codex_home.join(MODEL_CACHE_FILE); let cache_manager = ModelsCacheManager::new(cache_path, DEFAULT_MODEL_CACHE_TTL); Self { - local_models: builtin_model_presets(auth_manager.get_auth_mode()), + local_models: builtin_model_presets(auth_manager.get_internal_auth_mode()), remote_models: RwLock::new(Self::load_remote_models_from_file().unwrap_or_default()), auth_manager, etag: RwLock::new(None), diff --git a/codex-rs/core/src/models_manager/model_presets.rs b/codex-rs/core/src/models_manager/model_presets.rs index a8f4931f8..f9105c644 100644 --- a/codex-rs/core/src/models_manager/model_presets.rs +++ b/codex-rs/core/src/models_manager/model_presets.rs @@ -1,4 +1,4 @@ -use codex_app_server_protocol::AuthMode; +use crate::auth::AuthMode; use codex_protocol::openai_models::ModelPreset; use codex_protocol::openai_models::ModelUpgrade; use codex_protocol::openai_models::ReasoningEffort; diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 107e9f6ed..c65ce4498 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -984,7 +984,7 @@ impl App { model.as_str(), auth_ref.and_then(CodexAuth::get_account_id), auth_ref.and_then(CodexAuth::get_account_email), - auth_ref.map(|auth| auth.mode), + auth_ref.map(CodexAuth::api_auth_mode), config.otel.log_user_prompt, codex_core::terminal::user_agent(), SessionSource::Cli, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 8223c40ce..491ea3add 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -30,7 +30,6 @@ use std::time::Duration; use std::time::Instant; use crate::version::CODEX_CLI_VERSION; -use codex_app_server_protocol::AuthMode; use codex_backend_client::Client as BackendClient; use codex_chatgpt::connectors; use codex_core::config::Config; @@ -3584,10 +3583,12 @@ impl ChatWidget { fn prefetch_rate_limits(&mut self) { self.stop_rate_limit_poller(); - if !matches!( - self.auth_manager.auth_cached().map(|auth| auth.mode), - Some(AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) - ) { + if !self + .auth_manager + .auth_cached() + .as_ref() + .is_some_and(CodexAuth::is_chatgpt_auth) + { return; } @@ -3600,7 +3601,7 @@ impl ChatWidget { loop { if let Some(auth) = auth_manager.auth().await - && matches!(auth.mode, AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens) + && auth.is_chatgpt_auth() && let Some(snapshot) = fetch_rate_limits(base_url.clone(), auth).await { app_event_tx.send(AppEvent::RateLimitSnapshotFetched(snapshot)); diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index a8c6e2b88..1f34ed604 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -833,7 +833,7 @@ fn get_login_status(config: &Config) -> LoginStatus { // to refresh the token. Block on it. let codex_home = config.codex_home.clone(); match CodexAuth::from_auth_storage(&codex_home, config.cli_auth_credentials_store_mode) { - Ok(Some(auth)) => LoginStatus::AuthMode(auth.mode), + Ok(Some(auth)) => LoginStatus::AuthMode(auth.api_auth_mode()), Ok(None) => LoginStatus::NotAuthenticated, Err(err) => { error!("Failed to read auth.json: {err}"); diff --git a/codex-rs/tui/src/status/helpers.rs b/codex-rs/tui/src/status/helpers.rs index 91f958a9c..0a801227f 100644 --- a/codex-rs/tui/src/status/helpers.rs +++ b/codex-rs/tui/src/status/helpers.rs @@ -2,8 +2,8 @@ use crate::exec_command::relativize_to_home; use crate::text_formatting; use chrono::DateTime; use chrono::Local; -use codex_app_server_protocol::AuthMode; use codex_core::AuthManager; +use codex_core::CodexAuth; use codex_core::config::Config; use codex_core::project_doc::discover_project_doc_paths; use codex_protocol::account::PlanType; @@ -90,15 +90,15 @@ pub(crate) fn compose_account_display( ) -> Option { let auth = auth_manager.auth_cached()?; - match auth.mode { - AuthMode::ChatGPT | AuthMode::ChatgptAuthTokens => { + match auth { + CodexAuth::ChatGpt(_) | CodexAuth::ChatGptAuthTokens(_) => { let email = auth.get_account_email(); let plan = plan .map(|plan_type| title_case(format!("{plan_type:?}").as_str())) .or_else(|| Some("Unknown".to_string())); Some(StatusAccountDisplay::ChatGpt { email, plan }) } - AuthMode::ApiKey => Some(StatusAccountDisplay::ApiKey), + CodexAuth::ApiKey(_) => Some(StatusAccountDisplay::ApiKey), } }