From 548466df0928bf031be13088096588a33d8de4a6 Mon Sep 17 00:00:00 2001 From: Dylan Date: Thu, 7 Aug 2025 15:23:31 -0700 Subject: [PATCH 1/6] [client] Tune retries and backoff (#1956) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary 10 is a bit excessive 😅 Also updates our backoff factor to space out requests further. --- codex-rs/core/src/model_provider_info.rs | 2 +- codex-rs/core/src/util.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index a980211199..887917f470 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -15,7 +15,7 @@ use std::time::Duration; use crate::error::EnvVarError; const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000; -const DEFAULT_STREAM_MAX_RETRIES: u64 = 10; +const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; const DEFAULT_REQUEST_MAX_RETRIES: u64 = 4; /// Wire protocol that the provider speaks. Most third-party services only diff --git a/codex-rs/core/src/util.rs b/codex-rs/core/src/util.rs index 5ba1e25666..fb5f45de6f 100644 --- a/codex-rs/core/src/util.rs +++ b/codex-rs/core/src/util.rs @@ -7,7 +7,7 @@ use tokio::sync::Notify; use tracing::debug; const INITIAL_DELAY_MS: u64 = 200; -const BACKOFF_FACTOR: f64 = 1.3; +const BACKOFF_FACTOR: f64 = 2.0; /// Make a CancellationToken that is fulfilled when SIGINT occurs. pub fn notify_on_sigint() -> Arc { From db76f3288876b78d65a1d502a8503a8e1bedf78b Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 7 Aug 2025 16:33:29 -0700 Subject: [PATCH 2/6] chore: rename CodexAuth::new() to create_dummy_codex_auth_for_testing() because it is not for general consumption (#1962) `CodexAuth::new()` was the first method listed in `CodexAuth`, but it is only meant to be used by tests. Rename it to `create_dummy_chatgpt_auth_for_testing()` and move it to the end of the implementation. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/1962). * #1971 * #1970 * #1966 * #1965 * __->__ #1962 --- codex-rs/core/tests/client.rs | 24 ++--------------------- codex-rs/login/src/lib.rs | 37 +++++++++++++++++++++-------------- 2 files changed, 24 insertions(+), 37 deletions(-) diff --git a/codex-rs/core/tests/client.rs b/codex-rs/core/tests/client.rs index 2148e874bc..2ea772e3b7 100644 --- a/codex-rs/core/tests/client.rs +++ b/codex-rs/core/tests/client.rs @@ -1,8 +1,5 @@ -#![allow(clippy::expect_used)] -#![allow(clippy::unwrap_used)] -use std::path::PathBuf; +#![allow(clippy::expect_used, clippy::unwrap_used)] -use chrono::Utc; use codex_core::Codex; use codex_core::CodexSpawnOk; use codex_core::ModelProviderInfo; @@ -13,10 +10,7 @@ use codex_core::protocol::InputItem; use codex_core::protocol::Op; use codex_core::protocol::SessionConfiguredEvent; use codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; -use codex_login::AuthDotJson; -use codex_login::AuthMode; use codex_login::CodexAuth; -use codex_login::TokenData; use core_test_support::load_default_config_for_test; use core_test_support::load_sse_fixture_with_id; use core_test_support::wait_for_event; @@ -556,19 +550,5 @@ async fn env_var_overrides_loaded_auth() { } fn create_dummy_codex_auth() -> CodexAuth { - CodexAuth::new( - None, - AuthMode::ChatGPT, - PathBuf::new(), - Some(AuthDotJson { - openai_api_key: None, - tokens: Some(TokenData { - id_token: Default::default(), - access_token: "Access Token".to_string(), - refresh_token: "test".to_string(), - account_id: Some("account_id".to_string()), - }), - last_refresh: Some(Utc::now()), - }), - ) + CodexAuth::create_dummy_chatgpt_auth_for_testing() } diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 0157c77555..3aa1816fcb 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -51,21 +51,6 @@ impl PartialEq for CodexAuth { } impl CodexAuth { - pub fn new( - api_key: Option, - mode: AuthMode, - auth_file: PathBuf, - auth_dot_json: Option, - ) -> Self { - let auth_dot_json = Arc::new(Mutex::new(auth_dot_json)); - Self { - api_key, - mode, - auth_file, - auth_dot_json, - } - } - pub fn from_api_key(api_key: String) -> Self { Self { api_key: Some(api_key), @@ -142,6 +127,28 @@ impl CodexAuth { } } } + + /// Consider this private to integration tests. + pub fn create_dummy_chatgpt_auth_for_testing() -> Self { + let auth_dot_json = AuthDotJson { + openai_api_key: None, + tokens: Some(TokenData { + id_token: Default::default(), + access_token: "Access Token".to_string(), + refresh_token: "test".to_string(), + account_id: Some("account_id".to_string()), + }), + last_refresh: Some(Utc::now()), + }; + + let auth_dot_json = Arc::new(Mutex::new(Some(auth_dot_json))); + Self { + api_key: None, + mode: AuthMode::ChatGPT, + auth_file: PathBuf::new(), + auth_dot_json, + } + } } // Loads the available auth information from the auth.json or OPENAI_API_KEY environment variable. From df37e7e33daf9ed4470a079616b6e5f11ce2a632 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 7 Aug 2025 16:33:45 -0700 Subject: [PATCH 3/6] chore: make CodexAuth::api_key a private field --- codex-rs/cli/src/login.rs | 16 +++++++++------- codex-rs/login/src/lib.rs | 3 ++- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 5f56eb1c4d..3881a1ce3f 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -49,9 +49,9 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { match load_auth(&config.codex_home) { Ok(Some(auth)) => match auth.mode { - AuthMode::ApiKey => { - if let Some(api_key) = auth.api_key.as_deref() { - eprintln!("Logged in using an API key - {}", safe_format_key(api_key)); + AuthMode::ApiKey => match auth.get_token().await { + Ok(api_key) => { + eprintln!("Logged in using an API key - {}", safe_format_key(&api_key)); if let Ok(env_api_key) = env::var(OPENAI_API_KEY_ENV_VAR) { if env_api_key == api_key { @@ -60,11 +60,13 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { ); } } - } else { - eprintln!("Logged in using an API key"); + std::process::exit(0); } - std::process::exit(0); - } + Err(e) => { + eprintln!("Unexpected error retrieving API key: {e}"); + std::process::exit(1); + } + }, AuthMode::ChatGPT => { eprintln!("Logged in using ChatGPT"); std::process::exit(0); diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 3aa1816fcb..36a78aaa22 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -38,8 +38,9 @@ pub enum AuthMode { #[derive(Debug, Clone)] pub struct CodexAuth { - pub api_key: Option, pub mode: AuthMode, + + api_key: Option, auth_dot_json: Arc>>, auth_file: PathBuf, } From 5c1bfa660c4d61376cea56e1669d0af00f4496d1 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 7 Aug 2025 16:33:45 -0700 Subject: [PATCH 4/6] chore: move top-level load_auth() to CodexAuth::from_codex_home() --- codex-rs/chatgpt/src/chatgpt_token.rs | 3 ++- codex-rs/cli/src/login.rs | 4 ++-- codex-rs/cli/src/proto.rs | 4 ++-- codex-rs/core/src/codex_wrapper.rs | 4 ++-- codex-rs/login/src/lib.rs | 21 +++++++++++---------- codex-rs/tui/src/lib.rs | 4 ++-- 6 files changed, 21 insertions(+), 19 deletions(-) diff --git a/codex-rs/chatgpt/src/chatgpt_token.rs b/codex-rs/chatgpt/src/chatgpt_token.rs index 55ebc22a08..c674afbc57 100644 --- a/codex-rs/chatgpt/src/chatgpt_token.rs +++ b/codex-rs/chatgpt/src/chatgpt_token.rs @@ -1,3 +1,4 @@ +use codex_login::CodexAuth; use std::path::Path; use std::sync::LazyLock; use std::sync::RwLock; @@ -18,7 +19,7 @@ pub fn set_chatgpt_token_data(value: TokenData) { /// Initialize the ChatGPT token from auth.json file pub async fn init_chatgpt_token_from_auth(codex_home: &Path) -> std::io::Result<()> { - let auth = codex_login::load_auth(codex_home)?; + let auth = CodexAuth::from_codex_home(codex_home)?; if let Some(auth) = auth { let token_data = auth.get_token_data().await?; set_chatgpt_token_data(token_data); diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 3881a1ce3f..1a70bd27b6 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -4,8 +4,8 @@ use codex_common::CliConfigOverrides; use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_login::AuthMode; +use codex_login::CodexAuth; use codex_login::OPENAI_API_KEY_ENV_VAR; -use codex_login::load_auth; use codex_login::login_with_api_key; use codex_login::login_with_chatgpt; use codex_login::logout; @@ -47,7 +47,7 @@ pub async fn run_login_with_api_key( pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides); - match load_auth(&config.codex_home) { + match CodexAuth::from_codex_home(&config.codex_home) { Ok(Some(auth)) => match auth.mode { AuthMode::ApiKey => match auth.get_token().await { Ok(api_key) => { diff --git a/codex-rs/cli/src/proto.rs b/codex-rs/cli/src/proto.rs index 291e1680f1..6c1de7eaa9 100644 --- a/codex-rs/cli/src/proto.rs +++ b/codex-rs/cli/src/proto.rs @@ -9,7 +9,7 @@ use codex_core::config::Config; use codex_core::config::ConfigOverrides; use codex_core::protocol::Submission; use codex_core::util::notify_on_sigint; -use codex_login::load_auth; +use codex_login::CodexAuth; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tracing::error; @@ -36,7 +36,7 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> { .map_err(anyhow::Error::msg)?; let config = Config::load_with_cli_overrides(overrides_vec, ConfigOverrides::default())?; - let auth = load_auth(&config.codex_home)?; + let auth = CodexAuth::from_codex_home(&config.codex_home)?; let ctrl_c = notify_on_sigint(); let CodexSpawnOk { codex, .. } = Codex::spawn(config, auth, ctrl_c.clone()).await?; let codex = Arc::new(codex); diff --git a/codex-rs/core/src/codex_wrapper.rs b/codex-rs/core/src/codex_wrapper.rs index 1e26a9ebed..dc10ec8d84 100644 --- a/codex-rs/core/src/codex_wrapper.rs +++ b/codex-rs/core/src/codex_wrapper.rs @@ -6,7 +6,7 @@ use crate::config::Config; use crate::protocol::Event; use crate::protocol::EventMsg; use crate::util::notify_on_sigint; -use codex_login::load_auth; +use codex_login::CodexAuth; use tokio::sync::Notify; use uuid::Uuid; @@ -26,7 +26,7 @@ pub struct CodexConversation { /// that callers can surface the information to the UI. pub async fn init_codex(config: Config) -> anyhow::Result { let ctrl_c = notify_on_sigint(); - let auth = load_auth(&config.codex_home)?; + let auth = CodexAuth::from_codex_home(&config.codex_home)?; let CodexSpawnOk { codex, init_id, diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 36a78aaa22..8c8ad75bbc 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -61,6 +61,12 @@ impl CodexAuth { } } + /// Loads the available auth information from the auth.json or + /// OPENAI_API_KEY environment variable. + pub fn from_codex_home(codex_home: &Path) -> std::io::Result> { + load_auth(codex_home, true) + } + pub async fn get_token_data(&self) -> Result { #[expect(clippy::unwrap_used)] let auth_dot_json = self.auth_dot_json.lock().unwrap().clone(); @@ -152,12 +158,7 @@ impl CodexAuth { } } -// Loads the available auth information from the auth.json or OPENAI_API_KEY environment variable. -pub fn load_auth(codex_home: &Path) -> std::io::Result> { - _load_auth(codex_home, true) -} - -fn _load_auth(codex_home: &Path, include_env_var: bool) -> std::io::Result> { +fn load_auth(codex_home: &Path, include_env_var: bool) -> std::io::Result> { let auth_file = get_auth_file(codex_home); let auth_dot_json = try_read_auth_json(&auth_file).ok(); @@ -433,7 +434,7 @@ mod tests { fn writes_api_key_and_loads_auth() { let dir = tempdir().unwrap(); login_with_api_key(dir.path(), "sk-test-key").unwrap(); - let auth = _load_auth(dir.path(), false).unwrap().unwrap(); + let auth = load_auth(dir.path(), false).unwrap().unwrap(); assert_eq!(auth.mode, AuthMode::ApiKey); assert_eq!(auth.api_key.as_deref(), Some("sk-test-key")); } @@ -446,7 +447,7 @@ mod tests { let env_var = std::env::var(OPENAI_API_KEY_ENV_VAR); if let Ok(env_var) = env_var { - let auth = _load_auth(dir.path(), true).unwrap().unwrap(); + let auth = load_auth(dir.path(), true).unwrap().unwrap(); assert_eq!(auth.mode, AuthMode::ApiKey); assert_eq!(auth.api_key, Some(env_var)); } @@ -505,7 +506,7 @@ mod tests { mode, auth_dot_json, auth_file, - } = _load_auth(dir.path(), false).unwrap().unwrap(); + } = load_auth(dir.path(), false).unwrap().unwrap(); assert_eq!(None, api_key); assert_eq!(AuthMode::ChatGPT, mode); assert_eq!(dir.path().join("auth.json"), auth_file); @@ -570,7 +571,7 @@ mod tests { ) .unwrap(); - let auth = _load_auth(dir.path(), false).unwrap().unwrap(); + let auth = load_auth(dir.path(), false).unwrap().unwrap(); assert_eq!(auth.mode, AuthMode::ApiKey); assert_eq!(auth.api_key, Some("sk-test-key".to_string())); diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index a5cfe28393..e15a235a71 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -12,7 +12,7 @@ use codex_core::config::load_config_as_toml_with_cli_overrides; use codex_core::config_types::SandboxMode; use codex_core::protocol::AskForApproval; use codex_core::protocol::SandboxPolicy; -use codex_login::load_auth; +use codex_login::CodexAuth; use codex_ollama::DEFAULT_OSS_MODEL; use log_layer::TuiLogLayer; use std::fs::OpenOptions; @@ -304,7 +304,7 @@ fn should_show_login_screen(config: &Config) -> bool { // Reading the OpenAI API key is an async operation because it may need // to refresh the token. Block on it. let codex_home = config.codex_home.clone(); - match load_auth(&codex_home) { + match CodexAuth::from_codex_home(&codex_home) { Ok(Some(_)) => false, Ok(None) => true, Err(err) => { From c0b22b9c78f8785e215d4487d930cc30820d24c4 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 7 Aug 2025 16:33:45 -0700 Subject: [PATCH 5/6] chore: change CodexAuth::from_api_key() to take &str instead of String --- codex-rs/core/src/model_provider_info.rs | 2 +- codex-rs/core/tests/client.rs | 8 ++++---- codex-rs/core/tests/compact.rs | 2 +- codex-rs/core/tests/stream_no_completed.rs | 2 +- codex-rs/login/src/lib.rs | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index 887917f470..98f07deb1e 100644 --- a/codex-rs/core/src/model_provider_info.rs +++ b/codex-rs/core/src/model_provider_info.rs @@ -96,7 +96,7 @@ impl ModelProviderInfo { auth: &Option, ) -> crate::error::Result { let effective_auth = match self.api_key() { - Ok(Some(key)) => Some(CodexAuth::from_api_key(key)), + Ok(Some(key)) => Some(CodexAuth::from_api_key(&key)), Ok(None) => auth.clone(), Err(err) => { if auth.is_some() { diff --git a/codex-rs/core/tests/client.rs b/codex-rs/core/tests/client.rs index 2ea772e3b7..bd7ad2feef 100644 --- a/codex-rs/core/tests/client.rs +++ b/codex-rs/core/tests/client.rs @@ -93,7 +93,7 @@ async fn includes_session_id_and_model_headers_in_request() { let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new()); let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("Test API Key".to_string())), + Some(CodexAuth::from_api_key("Test API Key")), ctrl_c.clone(), ) .await @@ -167,7 +167,7 @@ async fn includes_base_instructions_override_in_request() { let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new()); let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("Test API Key".to_string())), + Some(CodexAuth::from_api_key("Test API Key")), ctrl_c.clone(), ) .await @@ -226,7 +226,7 @@ async fn originator_config_override_is_used() { let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new()); let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("Test API Key".to_string())), + Some(CodexAuth::from_api_key("Test API Key")), ctrl_c.clone(), ) .await @@ -364,7 +364,7 @@ async fn includes_user_instructions_message_in_request() { let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new()); let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("Test API Key".to_string())), + Some(CodexAuth::from_api_key("Test API Key")), ctrl_c.clone(), ) .await diff --git a/codex-rs/core/tests/compact.rs b/codex-rs/core/tests/compact.rs index bf47feed1f..fa5c81d883 100644 --- a/codex-rs/core/tests/compact.rs +++ b/codex-rs/core/tests/compact.rs @@ -145,7 +145,7 @@ async fn summarize_context_three_requests_and_instructions() { let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new()); let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("dummy".to_string())), + Some(CodexAuth::from_api_key("dummy")), ctrl_c.clone(), ) .await diff --git a/codex-rs/core/tests/stream_no_completed.rs b/codex-rs/core/tests/stream_no_completed.rs index 8a4216b129..0ded3337ab 100644 --- a/codex-rs/core/tests/stream_no_completed.rs +++ b/codex-rs/core/tests/stream_no_completed.rs @@ -99,7 +99,7 @@ async fn retries_on_early_close() { config.model_provider = model_provider; let CodexSpawnOk { codex, .. } = Codex::spawn( config, - Some(CodexAuth::from_api_key("Test API Key".to_string())), + Some(CodexAuth::from_api_key("Test API Key")), ctrl_c, ) .await diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 8c8ad75bbc..8571abcfb6 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -52,9 +52,9 @@ impl PartialEq for CodexAuth { } impl CodexAuth { - pub fn from_api_key(api_key: String) -> Self { + pub fn from_api_key(api_key: &str) -> Self { Self { - api_key: Some(api_key), + api_key: Some(api_key.to_owned()), mode: AuthMode::ApiKey, auth_file: PathBuf::new(), auth_dot_json: Arc::new(Mutex::new(None)), From 77e14c5019f9390bdb5eb11c4f087cd88189e96e Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 7 Aug 2025 16:33:45 -0700 Subject: [PATCH 6/6] fix: default to credits from ChatGPT auth, when possible --- codex-rs/login/src/lib.rs | 287 +++++++++++++++++++++---------- codex-rs/login/src/token_data.rs | 13 ++ 2 files changed, 211 insertions(+), 89 deletions(-) diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 8571abcfb6..b2541ddf44 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -160,46 +160,72 @@ impl CodexAuth { fn load_auth(codex_home: &Path, include_env_var: bool) -> std::io::Result> { let auth_file = get_auth_file(codex_home); - - let auth_dot_json = try_read_auth_json(&auth_file).ok(); - - let auth_json_api_key = auth_dot_json - .as_ref() - .and_then(|a| a.openai_api_key.clone()) - .filter(|s| !s.is_empty()); - - let openai_api_key = if include_env_var { - env::var(OPENAI_API_KEY_ENV_VAR) - .ok() - .filter(|s| !s.is_empty()) - .or(auth_json_api_key) - } else { - auth_json_api_key + let auth_dot_json = match try_read_auth_json(&auth_file) { + Ok(auth) => auth, + // If auth.json does not exist, try to read the OPENAI_API_KEY from the + // environment variable. + Err(e) if e.kind() == std::io::ErrorKind::NotFound && include_env_var => { + return match read_openai_api_key_from_env() { + Some(api_key) => Ok(Some(CodexAuth::from_api_key(&api_key))), + None => Ok(None), + }; + } + // Though if auth.json exists but is malformed, do not fall back to the + // env var because the user may be expecting to use AuthMode::ChatGPT. + Err(e) => { + return Err(e); + } }; - let has_tokens = auth_dot_json - .as_ref() - .and_then(|a| a.tokens.as_ref()) - .is_some(); + let AuthDotJson { + openai_api_key: auth_json_api_key, + tokens, + last_refresh, + } = auth_dot_json; - if openai_api_key.is_none() && !has_tokens { - return Ok(None); + // If the auth.json has an API key AND does not appear to be on a plan that + // should use prefer AuthMode::ChatGPT, use AuthMode::ApiKey. + if let Some(api_key) = &auth_json_api_key { + // Should any of these by AuthMode::ChatGPT with the api_key set? + // Does AuthMode::ChatGPT indicate that there is an auth.json that is + // "refreshable" even if we are using the API key for auth? + match &tokens { + Some(tokens) => { + if tokens.is_plan_that_should_use_api_key() { + return Ok(Some(CodexAuth::from_api_key(api_key))); + } else { + // Ignore the API key and fall through to ChatGPT auth. + } + } + None => { + // This is a bit suspicious because we have an API key but no + // tokens. Perhaps the user updated auth.json by hand, so let's + // assume they are trying to use their API key. + return Ok(Some(CodexAuth::from_api_key(api_key))); + } + } } - let mode = if openai_api_key.is_some() { - AuthMode::ApiKey - } else { - AuthMode::ChatGPT - }; - + // For the AuthMode::ChatGPT variant, perhaps neither api_key should not + // exist? Ok(Some(CodexAuth { - api_key: openai_api_key, - mode, + api_key: None, + mode: AuthMode::ChatGPT, auth_file, - auth_dot_json: Arc::new(Mutex::new(auth_dot_json)), + auth_dot_json: Arc::new(Mutex::new(Some(AuthDotJson { + openai_api_key: None, + tokens, + last_refresh, + }))), })) } +fn read_openai_api_key_from_env() -> Option { + env::var(OPENAI_API_KEY_ENV_VAR) + .ok() + .filter(|s| !s.is_empty()) +} + pub fn get_auth_file(codex_home: &Path) -> PathBuf { codex_home.join("auth.json") } @@ -423,14 +449,17 @@ pub struct AuthDotJson { #[cfg(test)] mod tests { + #![expect(clippy::expect_used, clippy::unwrap_used)] use super::*; use crate::token_data::IdTokenInfo; use base64::Engine; use pretty_assertions::assert_eq; + use serde_json::json; use tempfile::tempdir; + const LAST_REFRESH: &str = "2025-08-06T20:41:36.232376Z"; + #[test] - #[expect(clippy::unwrap_used)] fn writes_api_key_and_loads_auth() { let dir = tempdir().unwrap(); login_with_api_key(dir.path(), "sk-test-key").unwrap(); @@ -440,7 +469,6 @@ mod tests { } #[test] - #[expect(clippy::unwrap_used)] fn loads_from_env_var_if_env_var_exists() { let dir = tempdir().unwrap(); @@ -454,10 +482,132 @@ mod tests { } #[tokio::test] - #[expect(clippy::expect_used, clippy::unwrap_used)] - async fn loads_token_data_from_auth_json() { - let dir = tempdir().unwrap(); - let auth_file = dir.path().join("auth.json"); + async fn pro_account_with_no_api_key_uses_chatgpt_auth() { + let codex_home = tempdir().unwrap(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: "pro".to_string(), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let CodexAuth { + api_key, + mode, + auth_dot_json, + auth_file: _, + } = load_auth(codex_home.path(), false).unwrap().unwrap(); + assert_eq!(None, api_key); + assert_eq!(AuthMode::ChatGPT, mode); + + let guard = auth_dot_json.lock().unwrap(); + let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); + assert_eq!( + &AuthDotJson { + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: Some("user@example.com".to_string()), + chatgpt_plan_type: Some("pro".to_string()), + }, + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + }), + last_refresh: Some( + DateTime::parse_from_rfc3339(LAST_REFRESH) + .unwrap() + .with_timezone(&Utc) + ), + }, + auth_dot_json + ) + } + + /// Even if the OPENAI_API_KEY is set in auth.json, if the plan is not in + /// [`TokenData::is_plan_that_should_use_api_key`], it should use + /// [`AuthMode::ChatGPT`]. + #[tokio::test] + async fn pro_account_with_api_key_still_uses_chatgpt_auth() { + let codex_home = tempdir().unwrap(); + write_auth_file( + AuthFileParams { + openai_api_key: Some("sk-test-key".to_string()), + chatgpt_plan_type: "pro".to_string(), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let CodexAuth { + api_key, + mode, + auth_dot_json, + auth_file: _, + } = load_auth(codex_home.path(), false).unwrap().unwrap(); + assert_eq!(None, api_key); + assert_eq!(AuthMode::ChatGPT, mode); + + let guard = auth_dot_json.lock().unwrap(); + let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); + assert_eq!( + &AuthDotJson { + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: Some("user@example.com".to_string()), + chatgpt_plan_type: Some("pro".to_string()), + }, + access_token: "test-access-token".to_string(), + refresh_token: "test-refresh-token".to_string(), + account_id: None, + }), + last_refresh: Some( + DateTime::parse_from_rfc3339(LAST_REFRESH) + .unwrap() + .with_timezone(&Utc) + ), + }, + auth_dot_json + ) + } + + /// If the OPENAI_API_KEY is set in auth.json and it is an enterprise + /// account, then it should use [`AuthMode::ApiKey`]. + #[tokio::test] + async fn enterprise_account_with_api_key_uses_chatgpt_auth() { + let codex_home = tempdir().unwrap(); + write_auth_file( + AuthFileParams { + openai_api_key: Some("sk-test-key".to_string()), + chatgpt_plan_type: "enterprise".to_string(), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + + let CodexAuth { + api_key, + mode, + auth_dot_json, + auth_file: _, + } = load_auth(codex_home.path(), false).unwrap().unwrap(); + assert_eq!(Some("sk-test-key".to_string()), api_key); + assert_eq!(AuthMode::ApiKey, mode); + + let guard = auth_dot_json.lock().expect("should unwrap"); + assert!(guard.is_none(), "auth_dot_json should be None"); + } + + struct AuthFileParams { + openai_api_key: Option, + chatgpt_plan_type: String, + } + + fn write_auth_file(params: AuthFileParams, codex_home: &Path) -> std::io::Result<()> { + let auth_file = get_auth_file(codex_home); // Create a minimal valid JWT for the id_token field. #[derive(Serialize)] struct Header { @@ -473,71 +623,31 @@ mod tests { "email_verified": true, "https://api.openai.com/auth": { "chatgpt_account_id": "bc3618e3-489d-4d49-9362-1561dc53ba53", - "chatgpt_plan_type": "pro", + "chatgpt_plan_type": params.chatgpt_plan_type, "chatgpt_user_id": "user-12345", "user_id": "user-12345", } }); let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b); - let header_b64 = b64(&serde_json::to_vec(&header).unwrap()); - let payload_b64 = b64(&serde_json::to_vec(&payload).unwrap()); + let header_b64 = b64(&serde_json::to_vec(&header)?); + let payload_b64 = b64(&serde_json::to_vec(&payload)?); let signature_b64 = b64(b"sig"); let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); - std::fs::write( - auth_file, - format!( - r#" - {{ - "OPENAI_API_KEY": null, - "tokens": {{ - "id_token": "{fake_jwt}", + + let auth_json_data = json!({ + "OPENAI_API_KEY": params.openai_api_key, + "tokens": { + "id_token": fake_jwt, "access_token": "test-access-token", "refresh_token": "test-refresh-token" - }}, - "last_refresh": "2025-08-06T20:41:36.232376Z" - }} - "#, - ), - ) - .unwrap(); - - let CodexAuth { - api_key, - mode, - auth_dot_json, - auth_file, - } = load_auth(dir.path(), false).unwrap().unwrap(); - assert_eq!(None, api_key); - assert_eq!(AuthMode::ChatGPT, mode); - assert_eq!(dir.path().join("auth.json"), auth_file); - - let guard = auth_dot_json.lock().unwrap(); - let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); - - assert_eq!( - &AuthDotJson { - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: Some("user@example.com".to_string()), - chatgpt_plan_type: Some("pro".to_string()), - }, - access_token: "test-access-token".to_string(), - refresh_token: "test-refresh-token".to_string(), - account_id: None, - }), - last_refresh: Some( - DateTime::parse_from_rfc3339("2025-08-06T20:41:36.232376Z") - .unwrap() - .with_timezone(&Utc) - ), }, - auth_dot_json - ) + "last_refresh": LAST_REFRESH, + }); + let auth_json = serde_json::to_string_pretty(&auth_json_data)?; + std::fs::write(auth_file, auth_json) } #[test] - #[expect(clippy::expect_used, clippy::unwrap_used)] fn id_token_info_handles_missing_fields() { // Payload without email or plan should yield None values. let header = serde_json::json!({"alg": "none", "typ": "JWT"}); @@ -555,7 +665,6 @@ mod tests { } #[tokio::test] - #[expect(clippy::unwrap_used)] async fn loads_api_key_from_auth_json() { let dir = tempdir().unwrap(); let auth_file = dir.path().join("auth.json"); diff --git a/codex-rs/login/src/token_data.rs b/codex-rs/login/src/token_data.rs index 55b51b9d44..896d203809 100644 --- a/codex-rs/login/src/token_data.rs +++ b/codex-rs/login/src/token_data.rs @@ -17,6 +17,19 @@ pub struct TokenData { pub account_id: Option, } +impl TokenData { + /// Returns true if this is a plan that should use the traditional + /// "metered" billing via an API key. + pub(crate) fn is_plan_that_should_use_api_key(&self) -> bool { + match self.id_token.chatgpt_plan_type.as_deref() { + // TODO: Verify this is a comprehensive list of plans that + // should NOT use the API key. + Some("free") | Some("plus") | Some("pro") | Some("team") => false, + _ => true, + } + } +} + /// Flat subset of useful claims in id_token from auth.json. #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize)] pub struct IdTokenInfo {