diff --git a/codex-rs/core/src/model_provider_info.rs b/codex-rs/core/src/model_provider_info.rs index a980211199..ad8a3713f7 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..22547fa2ec 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)), @@ -160,46 +160,70 @@ 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? + 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") } 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 {