From d82eda652b96dd151209ea59946f71752a617557 Mon Sep 17 00:00:00 2001 From: Richard Lee Date: Wed, 20 May 2026 12:46:08 -0700 Subject: [PATCH] Add agent identity JWKS env override --- codex-rs/exec-server/README.md | 3 ++ codex-rs/login/src/auth/auth_tests.rs | 59 +++++++++++++++++++++++++++ codex-rs/login/src/auth/manager.rs | 12 ++++-- 3 files changed, 71 insertions(+), 3 deletions(-) diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 468ac205f7..762dedc60f 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -31,6 +31,9 @@ remote registration needs authentication. Containerized callers that receive an Agent Identity JWT in `CODEX_ACCESS_TOKEN` can opt into that auth path with `--use-agent-identity-auth`; Codex then registers an Agent task and sends the derived AgentAssertion headers on the registry request. +When the container cannot fetch the Agent Identity JWKS endpoint, set +`CODEX_AGENT_IDENTITY_JWKS_JSON` to the trusted JWKS endpoint JSON so Codex can +verify the JWT without that network request. Wire framing: diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 63ab3e0c42..bcf5d09598 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -126,6 +126,65 @@ async fn login_with_access_token_writes_only_token() { server.verify().await; } +#[tokio::test] +#[serial(codex_auth_env)] +async fn login_with_access_token_reads_agent_identity_jwks_from_env() { + let dir = tempdir().unwrap(); + let auth_path = dir.path().join("auth.json"); + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); + let _jwks_guard = EnvVarGuard::set( + CODEX_AGENT_IDENTITY_JWKS_JSON_ENV_VAR, + &test_jwks_body().to_string(), + ); + let server = MockServer::start().await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + + super::login_with_access_token( + dir.path(), + &agent_identity, + AuthCredentialsStoreMode::File, + Some(&chatgpt_base_url), + ) + .await + .expect("login_with_access_token should use JWKS from env"); + + let storage = FileAuthStorage::new(dir.path().to_path_buf()); + let auth = storage + .try_read_auth_json(&auth_path) + .expect("auth.json should parse"); + assert_eq!(auth.auth_mode, Some(AuthMode::AgentIdentity)); + assert_eq!( + auth.agent_identity.as_deref(), + Some(agent_identity.as_str()) + ); +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn login_with_access_token_rejects_invalid_agent_identity_jwks_env() { + let dir = tempdir().unwrap(); + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); + let _jwks_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_JWKS_JSON_ENV_VAR, "not-jwks-json"); + + super::login_with_access_token( + dir.path(), + &agent_identity, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect_err("invalid JWKS env should fail"); + + assert!( + !get_auth_file(dir.path()).exists(), + "invalid JWKS env should not write auth.json" + ); +} + #[tokio::test] async fn login_with_access_token_rejects_invalid_jwt() { let dir = tempdir().unwrap(); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index a2e4e8e0d8..c24f0ef2d1 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -465,6 +465,7 @@ impl ChatgptAuth { pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; pub const CODEX_ACCESS_TOKEN_ENV_VAR: &str = "CODEX_ACCESS_TOKEN"; +const CODEX_AGENT_IDENTITY_JWKS_JSON_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_JWKS_JSON"; pub fn read_openai_api_key_from_env() -> Option { env::var(OPENAI_API_KEY_ENV_VAR) @@ -493,9 +494,14 @@ async fn verified_agent_identity_record( chatgpt_base_url: &str, ) -> std::io::Result { AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; - let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url) - .await - .map_err(std::io::Error::other)?; + let jwks = + if let Some(jwks_json) = read_non_empty_env_var(CODEX_AGENT_IDENTITY_JWKS_JSON_ENV_VAR) { + serde_json::from_str(&jwks_json).map_err(std::io::Error::other)? + } else { + fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url) + .await + .map_err(std::io::Error::other)? + }; let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?; Ok(claims.into()) }