Merge f464733e2d into sapling-pr-archive-bolinfest

codex-rs/core/src/model_provider_info.rs

@@ -96,7 +96,7 @@ impl ModelProviderInfo {
         auth: &Option<CodexAuth>,
     ) -> crate::error::Result<reqwest::RequestBuilder> {
         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() {

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

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

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

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<Option<CodexAuth>> {
     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<String> {
+    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")
 }

codex-rs/login/src/token_data.rs

@@ -17,6 +17,19 @@ pub struct TokenData {
     pub account_id: Option<String>,
 }
 
+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 {