Add agent identity JWKS env override

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:
 

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();

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<String> {
     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> {
     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())
 }