feat: verify agent identity JWTs with JWKS (#19764)

codex-rs/login/Cargo.toml

@@ -45,6 +45,7 @@ webbrowser = { workspace = true }
 [dev-dependencies]
 anyhow = { workspace = true }
 core_test_support = { workspace = true }
+jsonwebtoken = { workspace = true }
 keyring = { workspace = true }
 pretty_assertions = { workspace = true }
 regex-lite = { workspace = true }

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

@@ -82,7 +82,7 @@ mod tests {
     use serial_test::serial;
 
     #[test]
-    #[serial(agent_identity_authapi_base_url_env)]
+    #[serial(codex_auth_env)]
     fn agent_identity_authapi_base_url_prefers_env_value() {
         let _guard = EnvVarGuard::set(
             CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR,
@@ -95,7 +95,7 @@ mod tests {
     }
 
     #[test]
-    #[serial(agent_identity_authapi_base_url_env)]
+    #[serial(codex_auth_env)]
     fn agent_identity_authapi_base_url_uses_prod_authapi_by_default() {
         let _guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR);
         assert_eq!(

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

@@ -16,6 +16,11 @@ use serde_json::json;
 use std::sync::Arc;
 use tempfile::TempDir;
 use tempfile::tempdir;
+use wiremock::Mock;
+use wiremock::MockServer;
+use wiremock::ResponseTemplate;
+use wiremock::matchers::method;
+use wiremock::matchers::path;
 
 #[tokio::test]
 async fn refresh_without_id_token() {
@@ -78,15 +83,29 @@ fn login_with_api_key_overwrites_existing_auth_json() {
     assert!(auth.tokens.is_none(), "tokens should be cleared");
 }
 
-#[test]
-fn login_with_agent_identity_writes_only_token() {
+#[tokio::test]
+async fn login_with_agent_identity_writes_only_token() {
     let dir = tempdir().unwrap();
     let auth_path = dir.path().join("auth.json");
     let record = agent_identity_record("account-123");
-    let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
+    let agent_identity = signed_agent_identity_jwt(&record).expect("signed agent identity");
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .and(path("/backend-api/wham/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(1)
+        .mount(&server)
+        .await;
+    let chatgpt_base_url = format!("{}/backend-api", server.uri());
 
-    super::login_with_agent_identity(dir.path(), &agent_identity, AuthCredentialsStoreMode::File)
-        .expect("login_with_agent_identity should succeed");
+    super::login_with_agent_identity(
+        dir.path(),
+        &agent_identity,
+        AuthCredentialsStoreMode::File,
+        Some(&chatgpt_base_url),
+    )
+    .await
+    .expect("login_with_agent_identity should succeed");
 
     let storage = FileAuthStorage::new(dir.path().to_path_buf());
     let auth = storage
@@ -99,15 +118,21 @@ fn login_with_agent_identity_writes_only_token() {
     );
     assert!(auth.tokens.is_none(), "tokens should be cleared");
     assert!(auth.openai_api_key.is_none(), "API key should be cleared");
+    server.verify().await;
 }
 
-#[test]
-fn login_with_agent_identity_rejects_invalid_jwt() {
+#[tokio::test]
+async fn login_with_agent_identity_rejects_invalid_jwt() {
     let dir = tempdir().unwrap();
 
-    let err =
-        super::login_with_agent_identity(dir.path(), "not-a-jwt", AuthCredentialsStoreMode::File)
-            .expect_err("invalid Agent Identity token should fail");
+    let err = super::login_with_agent_identity(
+        dir.path(),
+        "not-a-jwt",
+        AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
+    )
+    .await
+    .expect_err("invalid Agent Identity token should fail");
 
     assert_eq!(err.kind(), std::io::ErrorKind::Other);
     assert!(
@@ -117,11 +142,47 @@ fn login_with_agent_identity_rejects_invalid_jwt() {
 }
 
 #[tokio::test]
+async fn login_with_agent_identity_rejects_unsigned_jwt() {
+    let dir = tempdir().unwrap();
+    let record = agent_identity_record("account-123");
+    let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .and(path("/backend-api/wham/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(1)
+        .mount(&server)
+        .await;
+    let chatgpt_base_url = format!("{}/backend-api", server.uri());
+
+    super::login_with_agent_identity(
+        dir.path(),
+        &agent_identity,
+        AuthCredentialsStoreMode::File,
+        Some(&chatgpt_base_url),
+    )
+    .await
+    .expect_err("unsigned Agent Identity token should fail");
+
+    assert!(
+        !get_auth_file(dir.path()).exists(),
+        "unsigned Agent Identity token should not write auth.json"
+    );
+    server.verify().await;
+}
+
+#[tokio::test]
+#[serial(codex_auth_env)]
 async fn missing_auth_json_returns_none() {
     let dir = tempdir().unwrap();
-    let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File)
-        .await
-        .expect("call should succeed");
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
+    let auth = CodexAuth::from_auth_storage(
+        dir.path(),
+        AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
+    )
+    .await
+    .expect("call should succeed");
     assert_eq!(auth, None);
 }
 
@@ -129,6 +190,7 @@ async fn missing_auth_json_returns_none() {
 #[serial(codex_auth_env)]
 async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let fake_jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -143,6 +205,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .unwrap()
@@ -186,6 +249,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
 #[serial(codex_auth_env)]
 async fn loads_api_key_from_auth_json() {
     let dir = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let auth_file = dir.path().join("auth.json");
     std::fs::write(
         auth_file,
@@ -197,6 +261,7 @@ async fn loads_api_key_from_auth_json() {
         dir.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .unwrap()
@@ -255,8 +320,10 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -271,6 +338,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -288,6 +356,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
         codex_home.path(),
         updated_auth_dot_json,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("updated auth should parse");
@@ -591,6 +660,7 @@ async fn build_config(
         auth_credentials_store_mode: AuthCredentialsStoreMode::File,
         forced_login_method,
         forced_chatgpt_workspace_id,
+        chatgpt_base_url: None,
     }
 }
 
@@ -611,6 +681,14 @@ impl EnvVarGuard {
         }
         Self { key, original }
     }
+
+    fn remove(key: &'static str) -> Self {
+        let original = env::var_os(key);
+        unsafe {
+            env::remove_var(key);
+        }
+        Self { key, original }
+    }
 }
 
 #[cfg(test)]
@@ -625,6 +703,55 @@ impl Drop for EnvVarGuard {
     }
 }
 
+#[tokio::test]
+#[serial(codex_auth_env)]
+async fn load_auth_reads_agent_identity_from_env() {
+    let codex_home = tempdir().unwrap();
+    let expected_record = agent_identity_record("account-123");
+    let agent_identity =
+        signed_agent_identity_jwt(&expected_record).expect("signed agent identity");
+    let server = MockServer::start().await;
+    Mock::given(method("GET"))
+        .and(path("/backend-api/wham/agent-identities/jwks"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
+        .expect(1)
+        .mount(&server)
+        .await;
+    Mock::given(method("POST"))
+        .and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "task_id": "task-123",
+        })))
+        .expect(1)
+        .mount(&server)
+        .await;
+    let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
+
+    let chatgpt_base_url = format!("{}/backend-api", server.uri());
+    let _authapi_guard =
+        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
+    let auth = super::load_auth(
+        codex_home.path(),
+        /*enable_codex_api_key_env*/ false,
+        AuthCredentialsStoreMode::File,
+        Some(&chatgpt_base_url),
+    )
+    .await
+    .expect("env auth should load")
+    .expect("env auth should be present");
+
+    let CodexAuth::AgentIdentity(agent_identity) = auth else {
+        panic!("env auth should load as agent identity");
+    };
+    assert_eq!(agent_identity.record(), &expected_record);
+    assert_eq!(agent_identity.process_task_id(), "task-123");
+    assert!(
+        !get_auth_file(codex_home.path()).exists(),
+        "env auth should not write auth.json"
+    );
+    server.verify().await;
+}
+
 #[tokio::test]
 #[serial(codex_auth_env)]
 async fn load_auth_keeps_codex_api_key_env_precedence() {
@@ -638,6 +765,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ true,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("env auth should load")
@@ -650,6 +778,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
 #[serial(codex_auth_env)]
 async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
         .expect("seed api key");
 
@@ -674,6 +803,7 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
 #[serial(codex_auth_env)]
 async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -705,6 +835,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
 #[serial(codex_auth_env)]
 async fn enforce_login_restrictions_allows_matching_workspace() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -736,6 +867,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() {
 async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_forced_chatgpt_workspace_id_is_set()
  {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
         .expect("seed api key");
 
@@ -759,6 +891,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f
 #[serial(codex_auth_env)]
 async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
     let _guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env");
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let codex_home = tempdir().unwrap();
 
     let config = build_config(
@@ -795,6 +928,10 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
     let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
     let header_b64 = encode(br#"{"alg":"EdDSA","typ":"JWT"}"#);
     let payload = json!({
+        "iss": "https://chatgpt.com/codex-backend/agent-identity",
+        "aud": "codex-app-server",
+        "iat": 1_700_000_000usize,
+        "exp": 4_000_000_000usize,
         "agent_runtime_id": record.agent_runtime_id,
         "agent_private_key": record.agent_private_key,
         "account_id": record.account_id,
@@ -808,9 +945,77 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
     Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
 }
 
+fn signed_agent_identity_jwt(
+    record: &AgentIdentityAuthRecord,
+) -> jsonwebtoken::errors::Result<String> {
+    let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
+    header.kid = Some("test-key".to_string());
+    jsonwebtoken::encode(
+        &header,
+        &json!({
+            "iss": "https://chatgpt.com/codex-backend/agent-identity",
+            "aud": "codex-app-server",
+            "iat": 1_700_000_000usize,
+            "exp": 4_000_000_000usize,
+            "agent_runtime_id": record.agent_runtime_id,
+            "agent_private_key": record.agent_private_key,
+            "account_id": record.account_id,
+            "chatgpt_user_id": record.chatgpt_user_id,
+            "email": record.email,
+            "plan_type": record.plan_type,
+            "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp,
+        }),
+        &jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?,
+    )
+}
+
+fn test_jwks_body() -> serde_json::Value {
+    json!({
+        "keys": [{
+            "kty": "RSA",
+            "kid": "test-key",
+            "use": "sig",
+            "alg": "RS256",
+            "n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
+            "e": "AQAB",
+        }]
+    })
+}
+
+const TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM: &[u8] = br#"-----BEGIN PRIVATE KEY-----
+MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
+bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
+FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
+mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
+0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
+zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
+uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
+1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
+RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
+NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
+rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
+Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
+aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
+iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
+YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
+3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
+sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
+gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
+sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
+BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
+KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
+r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
+29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
+J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
+5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
+8U4M2TSWCKUY/A6sT4W8+mT9
+-----END PRIVATE KEY-----"#;
+
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn plan_type_maps_known_plan() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -825,6 +1030,7 @@ async fn plan_type_maps_known_plan() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -834,8 +1040,10 @@ async fn plan_type_maps_known_plan() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn plan_type_maps_self_serve_business_usage_based_plan() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -850,6 +1058,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -862,8 +1071,10 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -878,6 +1089,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -890,8 +1102,10 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn plan_type_maps_unknown_to_unknown() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -906,6 +1120,7 @@ async fn plan_type_maps_unknown_to_unknown() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -915,8 +1130,10 @@ async fn plan_type_maps_unknown_to_unknown() {
 }
 
 #[tokio::test]
+#[serial(codex_auth_env)]
 async fn missing_plan_type_maps_to_unknown() {
     let codex_home = tempdir().unwrap();
+    let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
     let _jwt = write_auth_file(
         AuthFileParams {
             openai_api_key: None,
@@ -931,6 +1148,7 @@ async fn missing_plan_type_maps_to_unknown() {
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
         AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
     )
     .await
     .expect("load auth")

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

@@ -16,6 +16,8 @@ use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 use tokio::sync::Semaphore;
 
+use codex_agent_identity::decode_agent_identity_jwt;
+use codex_agent_identity::fetch_agent_identity_jwks;
 use codex_app_server_protocol::AuthMode;
 use codex_app_server_protocol::AuthMode as ApiAuthMode;
 use codex_protocol::config_types::ForcedLoginMethod;
@@ -29,6 +31,7 @@ pub use crate::auth::storage::AuthDotJson;
 use crate::auth::storage::AuthStorageBackend;
 use crate::auth::storage::create_auth_storage;
 use crate::auth::util::try_parse_error_message;
+use crate::default_client::build_reqwest_client;
 use crate::default_client::create_client;
 use crate::token_data::TokenData;
 use crate::token_data::parse_chatgpt_jwt_claims;
@@ -88,6 +91,7 @@ const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be
 const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str =
     "Your access token could not be refreshed. Please log out and sign in again.";
 const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again.";
+const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
 const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
 pub(super) const REVOKE_TOKEN_URL: &str = "https://auth.openai.com/oauth/revoke";
 pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE";
@@ -197,6 +201,7 @@ impl CodexAuth {
         codex_home: &Path,
         auth_dot_json: AuthDotJson,
         auth_credentials_store_mode: AuthCredentialsStoreMode,
+        chatgpt_base_url: Option<&str>,
     ) -> std::io::Result<Self> {
         let auth_mode = auth_dot_json.resolved_mode();
         let client = create_client();
@@ -212,7 +217,7 @@ impl CodexAuth {
                     "agent identity auth is missing an agent identity token.",
                 ));
             };
-            return Self::from_agent_identity_jwt(&agent_identity).await;
+            return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await;
         }
 
         let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
@@ -237,17 +242,26 @@ impl CodexAuth {
     pub async fn from_auth_storage(
         codex_home: &Path,
         auth_credentials_store_mode: AuthCredentialsStoreMode,
+        chatgpt_base_url: Option<&str>,
     ) -> std::io::Result<Option<Self>> {
         load_auth(
             codex_home,
             /*enable_codex_api_key_env*/ false,
             auth_credentials_store_mode,
+            chatgpt_base_url,
         )
         .await
     }
 
-    pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
-        let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
+    pub async fn from_agent_identity_jwt(
+        jwt: &str,
+        chatgpt_base_url: Option<&str>,
+    ) -> std::io::Result<Self> {
+        let base_url = chatgpt_base_url
+            .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
+            .trim_end_matches('/')
+            .to_string();
+        let record = verified_agent_identity_record(jwt, &base_url).await?;
         Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
     }
 
@@ -493,6 +507,18 @@ pub fn read_codex_agent_identity_from_env() -> Option<String> {
         .filter(|value| !value.is_empty())
 }
 
+async fn verified_agent_identity_record(
+    jwt: &str,
+    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 claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?;
+    Ok(claims.into())
+}
+
 /// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)`
 /// if a file was removed, `Ok(false)` if no auth file was present.
 pub fn logout(
@@ -535,12 +561,17 @@ pub fn login_with_api_key(
 }
 
 /// Writes an `auth.json` that contains only the Agent Identity token.
-pub fn login_with_agent_identity(
+pub async fn login_with_agent_identity(
     codex_home: &Path,
     agent_identity: &str,
     auth_credentials_store_mode: AuthCredentialsStoreMode,
+    chatgpt_base_url: Option<&str>,
 ) -> std::io::Result<()> {
-    AgentIdentityAuthRecord::from_agent_identity_jwt(agent_identity)?;
+    let base_url = chatgpt_base_url
+        .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
+        .trim_end_matches('/')
+        .to_string();
+    verified_agent_identity_record(agent_identity, &base_url).await?;
     let auth_dot_json = AuthDotJson {
         auth_mode: Some(ApiAuthMode::AgentIdentity),
         openai_api_key: None,
@@ -599,6 +630,7 @@ pub struct AuthConfig {
     pub auth_credentials_store_mode: AuthCredentialsStoreMode,
     pub forced_login_method: Option<ForcedLoginMethod>,
     pub forced_chatgpt_workspace_id: Option<String>,
+    pub chatgpt_base_url: Option<String>,
 }
 
 pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> {
@@ -606,6 +638,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<
         &config.codex_home,
         /*enable_codex_api_key_env*/ true,
         config.auth_credentials_store_mode,
+        config.chatgpt_base_url.as_deref(),
     )
     .await?
     else {
@@ -711,6 +744,7 @@ async fn load_auth(
     codex_home: &Path,
     enable_codex_api_key_env: bool,
     auth_credentials_store_mode: AuthCredentialsStoreMode,
+    chatgpt_base_url: Option<&str>,
 ) -> std::io::Result<Option<CodexAuth>> {
     // API key via env var takes precedence over any other auth method.
     if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() {
@@ -728,6 +762,7 @@ async fn load_auth(
             codex_home,
             auth_dot_json,
             AuthCredentialsStoreMode::Ephemeral,
+            chatgpt_base_url,
         )
         .await?;
         return Ok(Some(auth));
@@ -739,7 +774,7 @@ async fn load_auth(
     }
 
     if let Some(agent_identity) = read_codex_agent_identity_from_env() {
-        return CodexAuth::from_agent_identity_jwt(&agent_identity)
+        return CodexAuth::from_agent_identity_jwt(&agent_identity, chatgpt_base_url)
             .await
             .map(Some);
     }
@@ -751,9 +786,13 @@ async fn load_auth(
         None => return Ok(None),
     };
 
-    let auth =
-        CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, auth_credentials_store_mode)
-            .await?;
+    let auth = CodexAuth::from_auth_dot_json(
+        codex_home,
+        auth_dot_json,
+        auth_credentials_store_mode,
+        chatgpt_base_url,
+    )
+    .await?;
     Ok(Some(auth))
 }
 
@@ -1288,6 +1327,7 @@ impl AuthManager {
             &codex_home,
             enable_codex_api_key_env,
             auth_credentials_store_mode,
+            chatgpt_base_url.as_deref(),
         )
         .await
         .ok()
@@ -1491,6 +1531,7 @@ impl AuthManager {
             &self.codex_home,
             self.enable_codex_api_key_env,
             self.auth_credentials_store_mode,
+            self.chatgpt_base_url.as_deref(),
         )
         .await
         .ok()

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

@@ -19,6 +19,7 @@ use std::sync::Mutex;
 use tracing::warn;
 
 use crate::token_data::TokenData;
+use codex_agent_identity::AgentIdentityJwtClaims;
 use codex_agent_identity::decode_agent_identity_jwt;
 use codex_app_server_protocol::AuthMode;
 use codex_config::types::AuthCredentialsStoreMode;
@@ -59,10 +60,16 @@ pub struct AgentIdentityAuthRecord {
 
 impl AgentIdentityAuthRecord {
     pub(crate) fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
-        let claims = decode_agent_identity_jwt(jwt, /*public_key_base64*/ None)
-            .map_err(std::io::Error::other)?;
+        let claims =
+            decode_agent_identity_jwt(jwt, /*jwks*/ None).map_err(std::io::Error::other)?;
 
-        Ok(Self {
+        Ok(claims.into())
+    }
+}
+
+impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {
+    fn from(claims: AgentIdentityJwtClaims) -> Self {
+        Self {
             agent_runtime_id: claims.agent_runtime_id,
             agent_private_key: claims.agent_private_key,
             account_id: claims.account_id,
@@ -70,7 +77,7 @@ impl AgentIdentityAuthRecord {
             email: claims.email,
             plan_type: claims.plan_type,
             chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,
-        })
+        }
     }
 }