From b94343ab9f3f3e77f945d88eb10743bbd15eb8d3 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Fri, 7 Aug 2026 02:25:41 +0000 Subject: [PATCH] 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 | 56 +++++++- codex-rs/login/src/auth/auth_tests.rs | 149 ++++++++++++++++++++++ 2 files changed, 199 insertions(+), 6 deletions(-) diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 9894951f04..81185bc1bf 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/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 { + 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 { 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::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")) diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 66b1b1f784..78eb803dba 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/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);