Support agent identity endpoint overrides (#37356)

## What changed

- Honor `CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL` and
  `CODEX_AGENT_IDENTITY_JWKS_BASE_URL` when registering and verifying agent
  identities.
- Normalize override values by trimming whitespace and trailing slashes.
- Preserve ChatGPT environment validation, while allowing a custom launcher
  configuration when its ChatGPT and JWKS base URLs match and an AuthAPI
  override is explicitly configured.

## Testing

Added coverage for explicit staging overrides, custom launcher compatibility,
and rejection of overrides that would bypass environment validation.

GitOrigin-RevId: c557c8167ec101c2993f65aad3fcb3bf25785d51

codex-rs/login/src/auth/agent_identity.rs

@@ -1,3 +1,4 @@
+use std::env;
 use std::future::Future;
 use std::sync::Arc;
 
@@ -25,16 +26,46 @@ use crate::outbound_proxy::AuthRouteConfig;
 use super::storage::AgentIdentityAuthRecord;
 
 pub(super) const MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS: usize = 3;
+const CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL";
+const CODEX_AGENT_IDENTITY_JWKS_BASE_URL_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_JWKS_BASE_URL";
+
+fn agent_identity_endpoint_override(environment_variable: &str) -> Option<String> {
+    env::var(environment_variable)
+        .ok()
+        .map(|base_url| base_url.trim().trim_end_matches('/').to_string())
+        .filter(|base_url| !base_url.is_empty())
+}
+
+fn agent_identity_jwks_base_url_matches(chatgpt_base_url: &str, jwks_base_url: &str) -> bool {
+    chatgpt_base_url.trim().trim_end_matches('/') == jwks_base_url
+}
 
 pub(super) fn agent_identity_authapi_base_url(
     chatgpt_base_url: Option<&str>,
 ) -> std::io::Result<String> {
     let environment = match chatgpt_base_url {
-        Some(chatgpt_base_url) => ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url)
-            .map_err(std::io::Error::other)?,
-        None => ChatGptEnvironment::default(),
+        Some(chatgpt_base_url) => ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url),
+        None => Ok(ChatGptEnvironment::default()),
     };
-    Ok(environment.agent_identity_authapi_base_url().to_string())
+    let authapi_base_url =
+        agent_identity_endpoint_override(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR);
+    let jwks_base_url =
+        agent_identity_endpoint_override(CODEX_AGENT_IDENTITY_JWKS_BASE_URL_ENV_VAR);
+
+    match (environment, authapi_base_url) {
+        (Ok(_), Some(base_url)) => Ok(base_url),
+        (Ok(environment), None) => Ok(environment.agent_identity_authapi_base_url().to_string()),
+        (Err(_), Some(base_url))
+            if chatgpt_base_url.is_some_and(|chatgpt_base_url| {
+                jwks_base_url.as_deref().is_some_and(|jwks_base_url| {
+                    agent_identity_jwks_base_url_matches(chatgpt_base_url, jwks_base_url)
+                })
+            }) =>
+        {
+            Ok(base_url)
+        }
+        (Err(error), _) => Err(std::io::Error::other(error)),
+    }
 }
 
 pub(super) fn require_agent_identity_authapi_base_url(
@@ -225,9 +256,20 @@ pub(super) async fn verified_record_from_jwt(
     auth_route_config: &AuthRouteConfig,
 ) -> std::io::Result<AgentIdentityAuthRecord> {
     AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
-    let jwks_url = agent_identity_jwks_url(chatgpt_base_url);
+    let jwks_base_url =
+        match agent_identity_endpoint_override(CODEX_AGENT_IDENTITY_JWKS_BASE_URL_ENV_VAR) {
+            Some(base_url) => {
+                if !agent_identity_jwks_base_url_matches(chatgpt_base_url, &base_url) {
+                    ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url)
+                        .map_err(std::io::Error::other)?;
+                }
+                base_url
+            }
+            None => chatgpt_base_url.to_string(),
+        };
+    let jwks_url = agent_identity_jwks_url(&jwks_base_url);
     let client = create_default_auth_client(&jwks_url, auth_route_config)?;
-    let jwks = fetch_agent_identity_jwks(&client, chatgpt_base_url)
+    let jwks = fetch_agent_identity_jwks(&client, &jwks_base_url)
         .await
         .map_err(std::io::Error::other)?;
     let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?;
@@ -351,6 +393,7 @@ mod tests {
     use codex_agent_identity::generate_agent_key_material;
     use pretty_assertions::assert_eq;
     use serde_json::json;
+    use serial_test::serial;
     use wiremock::Mock;
     use wiremock::MockServer;
     use wiremock::ResponseTemplate;
@@ -414,6 +457,7 @@ mod tests {
     }
 
     #[tokio::test]
+    #[serial(codex_auth_env)]
     async fn from_jwt_registers_task() -> anyhow::Result<()> {
         let server = MockServer::start().await;
         Mock::given(method("GET"))

codex-rs/login/src/auth/auth_tests.rs

@@ -99,6 +99,7 @@ fn login_with_api_key_overwrites_existing_auth_json() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn login_with_access_token_writes_agent_identity_jwt() {
     let dir = tempdir().unwrap();
     let auth_path = dir.path().join("auth.json");
@@ -168,6 +169,152 @@ async fn login_with_access_token_rejects_agent_identity_workspace_mismatch() {
     assert!(server.received_requests().await.unwrap().is_empty());
 }
 
+#[tokio::test]
+#[serial(codex_auth_env)]
+async fn agent_identity_jwt_uses_explicit_staging_endpoint_overrides() -> anyhow::Result<()> {
+    let jwks_server = MockServer::start().await;
+    let authapi_server = MockServer::start().await;
+    let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
+    let jwt = signed_agent_identity_jwt(&record, json!(record.plan_type))?;
+    Mock::given(method("GET"))
+        .and(path("/api/codex/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(1)
+        .mount(&jwks_server)
+        .await;
+    mock_agent_task_registration(
+        &authapi_server,
+        "/api/accounts",
+        &record.agent_runtime_id,
+        "task-id",
+    )
+    .await;
+    let authapi_base_url = format!("{}/api/accounts/", authapi_server.uri());
+    let _authapi_guard =
+        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url);
+    let jwks_base_url = format!("{}/api/codex/", jwks_server.uri());
+    let _jwks_guard = EnvVarGuard::set("CODEX_AGENT_IDENTITY_JWKS_BASE_URL", &jwks_base_url);
+
+    let auth = CodexAuth::from_agent_identity_jwt(
+        &jwt,
+        Some(ChatGptEnvironment::Staging.chatgpt_base_url()),
+        &crate::test_support::transport_default_auth_route_config(),
+    )
+    .await?;
+
+    let CodexAuth::AgentIdentity(agent_identity) = auth else {
+        panic!("JWT should load as agent identity auth");
+    };
+    assert_eq!(agent_identity.run_task_id(), "task-id");
+    jwks_server.verify().await;
+    authapi_server.verify().await;
+    Ok(())
+}
+
+#[tokio::test]
+#[serial(codex_auth_env)]
+async fn agent_identity_jwt_supports_existing_staging_launcher() -> anyhow::Result<()> {
+    let jwks_server = MockServer::start().await;
+    let authapi_server = MockServer::start().await;
+    let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
+    let jwt = signed_agent_identity_jwt(&record, json!(record.plan_type))?;
+    Mock::given(method("GET"))
+        .and(path("/api/codex/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(1)
+        .mount(&jwks_server)
+        .await;
+    mock_agent_task_registration(
+        &authapi_server,
+        "/api/accounts",
+        &record.agent_runtime_id,
+        "task-id",
+    )
+    .await;
+    let authapi_base_url = format!("{}/api/accounts", authapi_server.uri());
+    let _authapi_guard =
+        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url);
+    let jwks_base_url = format!("{}/api/codex", jwks_server.uri());
+    let _jwks_guard = EnvVarGuard::set("CODEX_AGENT_IDENTITY_JWKS_BASE_URL", &jwks_base_url);
+
+    let auth = CodexAuth::from_agent_identity_jwt(
+        &jwt,
+        Some(&jwks_base_url),
+        &crate::test_support::transport_default_auth_route_config(),
+    )
+    .await?;
+
+    let CodexAuth::AgentIdentity(agent_identity) = auth else {
+        panic!("JWT should load as agent identity auth");
+    };
+    assert_eq!(agent_identity.run_task_id(), "task-id");
+    jwks_server.verify().await;
+    authapi_server.verify().await;
+    Ok(())
+}
+
+#[test]
+#[serial(codex_auth_env)]
+fn agent_identity_authapi_override_preserves_chatgpt_environment_validation() {
+    let _authapi_guard = EnvVarGuard::set(
+        "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL",
+        "https://authapi.example/api/accounts",
+    );
+    let _jwks_guard = EnvVarGuard::set(
+        "CODEX_AGENT_IDENTITY_JWKS_BASE_URL",
+        "https://jwks.example/api/codex",
+    );
+
+    let error = agent_identity_authapi_base_url(Some("https://attacker.example/backend-api"))
+        .expect_err("AuthAPI overrides must not bypass ChatGPT environment validation");
+
+    assert_eq!(
+        error.to_string(),
+        "Agent Identity only supports production and staging ChatGPT environments"
+    );
+}
+
+#[test]
+#[serial(codex_auth_env)]
+fn agent_identity_custom_jwks_base_requires_explicit_authapi_override() {
+    let _authapi_guard = EnvVarGuard::remove("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL");
+    let jwks_base_url = "https://jwks.example/api/codex";
+    let _jwks_guard = EnvVarGuard::set("CODEX_AGENT_IDENTITY_JWKS_BASE_URL", jwks_base_url);
+
+    let error = agent_identity_authapi_base_url(Some(jwks_base_url))
+        .expect_err("custom JWKS bases must also explicitly configure AuthAPI");
+
+    assert_eq!(
+        error.to_string(),
+        "Agent Identity only supports production and staging ChatGPT environments"
+    );
+}
+
+#[tokio::test]
+#[serial(codex_auth_env)]
+async fn agent_identity_jwks_override_preserves_chatgpt_environment_validation() {
+    let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
+    let jwt =
+        signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity");
+    let _jwks_guard = EnvVarGuard::set(
+        "CODEX_AGENT_IDENTITY_JWKS_BASE_URL",
+        "https://jwks.example/api/codex",
+    );
+
+    let error = verified_record_from_jwt(
+        &jwt,
+        "https://attacker.example/backend-api",
+        &crate::test_support::transport_default_auth_route_config(),
+    )
+    .await
+    .expect_err("JWKS overrides must not bypass ChatGPT environment validation");
+
+    assert_eq!(
+        error.to_string(),
+        "Agent Identity only supports production and staging ChatGPT environments"
+    );
+}
+
 #[tokio::test]
 #[serial(codex_auth_env)]
 async fn stored_agent_identity_jwt_keeps_auth_json_unchanged() -> anyhow::Result<()> {
@@ -744,6 +891,7 @@ async fn chatgpt_auth_non_retryable_registration_error_is_hard_failure() -> anyh
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> anyhow::Result<()> {
     let record = agent_identity_record(WORKSPACE_ID_ALLOWED);
     let agent_identity =
@@ -781,6 +929,7 @@ async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> an
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn login_with_access_token_rejects_unsigned_jwt() {
     let dir = tempdir().unwrap();
     let record = agent_identity_record(WORKSPACE_ID_ALLOWED);