From 4f0091f99c600ec7db188e182ec36fd07b8bc3aa Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Sun, 26 Apr 2026 21:19:48 -0700 Subject: [PATCH 1/2] feat: verify agent identity JWTs with JWKS --- codex-rs/Cargo.lock | 1 + codex-rs/agent-identity/src/lib.rs | 237 +++++++++++++++---- codex-rs/cli/src/login.rs | 6 +- codex-rs/cli/tests/login.rs | 16 +- codex-rs/core/tests/suite/client.rs | 1 + codex-rs/exec/src/lib.rs | 1 + codex-rs/login/Cargo.toml | 1 + codex-rs/login/src/auth/auth_tests.rs | 236 +++++++++++++++++- codex-rs/login/src/auth/manager.rs | 50 +++- codex-rs/login/src/auth/storage.rs | 15 +- codex-rs/models-manager/src/manager_tests.rs | 1 + codex-rs/tui/src/lib.rs | 1 + 12 files changed, 487 insertions(+), 79 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 464b7d72a21c..313fe3dd6bf4 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2836,6 +2836,7 @@ dependencies = [ "codex-terminal-detection", "codex-utils-template", "core_test_support", + "jsonwebtoken", "keyring", "once_cell", "os_info", diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index bf139f787027..12c7fd4f2807 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -19,6 +19,9 @@ use ed25519_dalek::pkcs8::EncodePrivateKey; use jsonwebtoken::Algorithm; use jsonwebtoken::DecodingKey; use jsonwebtoken::Validation; +use jsonwebtoken::decode; +use jsonwebtoken::decode_header; +use jsonwebtoken::jwk::JwkSet; use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; @@ -28,6 +31,9 @@ use sha2::Digest as _; use sha2::Sha512; const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30); +const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10); +const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server"; +const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity"; /// Stored key material for a registered agent identity. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -58,6 +64,10 @@ pub struct GeneratedAgentKeyMaterial { /// Claims carried by an Agent Identity JWT. #[derive(Clone, Debug, Deserialize, PartialEq, Eq)] pub struct AgentIdentityJwtClaims { + pub iss: String, + pub aud: String, + pub iat: usize, + pub exp: usize, pub agent_runtime_id: String, pub agent_private_key: String, pub account_id: String, @@ -115,27 +125,47 @@ pub fn authorization_header_for_agent_task( Ok(format!("AgentAssertion {serialized_assertion}")) } +pub async fn fetch_agent_identity_jwks( + client: &reqwest::Client, + chatgpt_base_url: &str, +) -> Result { + client + .get(agent_identity_jwks_url(chatgpt_base_url)) + .timeout(AGENT_IDENTITY_JWKS_TIMEOUT) + .send() + .await + .context("failed to fetch agent identity JWKS")? + .error_for_status() + .context("failed to fetch agent identity JWKS")? + .json() + .await + .context("failed to decode agent identity JWKS") +} + pub fn decode_agent_identity_jwt( jwt: &str, - public_key_base64: Option<&str>, + jwks: Option<&JwkSet>, ) -> Result { - let Some(public_key_base64) = public_key_base64 else { + let Some(jwks) = jwks else { return decode_agent_identity_jwt_payload(jwt); }; - let mut validation = Validation::new(Algorithm::EdDSA); - validation.required_spec_claims.clear(); - validation.validate_exp = false; - validation.validate_aud = false; - - let public_key = BASE64_STANDARD - .decode(public_key_base64) - .context("agent identity JWT public key is not valid base64")?; - let decoding_key = DecodingKey::from_ed_der(&public_key); - - jsonwebtoken::decode::(jwt, &decoding_key, &validation) + let header = decode_header(jwt).context("failed to decode agent identity JWT header")?; + let kid = header + .kid + .context("agent identity JWT header does not include a kid")?; + let jwk = jwks + .find(&kid) + .with_context(|| format!("agent identity JWT kid {kid} is not trusted"))?; + let decoding_key = DecodingKey::from_jwk(jwk).context("failed to build JWT decoding key")?; + let mut validation = Validation::new(Algorithm::RS256); + validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]); + validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]); + validation.required_spec_claims.insert("iss".to_string()); + validation.required_spec_claims.insert("aud".to_string()); + decode::(jwt, &decoding_key, &validation) .map(|data| data.claims) - .context("failed to decode agent identity JWT") + .context("failed to verify agent identity JWT") } fn decode_agent_identity_jwt_payload(jwt: &str) -> Result { @@ -279,6 +309,17 @@ pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { format!("{trimmed}/authenticate_app_v2") } +pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + if trimmed.contains("/backend-api") { + format!("{trimmed}/wham/agent-identities/jwks") + } else if trimmed.contains("/api/codex") { + format!("{trimmed}/agent-identities/jwks") + } else { + format!("{trimmed}/api/codex/agent-identities/jwks") + } +} + pub fn agent_identity_request_id() -> Result { let mut request_id_bytes = [0u8; 16]; OsRng @@ -301,12 +342,16 @@ pub fn normalize_chatgpt_base_url(chatgpt_base_url: &str) -> String { break; } } - if let Some(stripped) = base_url.strip_suffix("/codex") { + if base_url.ends_with("/codex") + && !base_url.ends_with("/api/codex") + && let Some(stripped) = base_url.strip_suffix("/codex") + { base_url = stripped.to_string(); } if (base_url.starts_with("https://chatgpt.com") || base_url.starts_with("https://chat.openai.com")) && !base_url.contains("/backend-api") + && !base_url.contains("/api/codex") { base_url = format!("{base_url}/backend-api"); } @@ -471,6 +516,10 @@ mod tests { #[test] fn decode_agent_identity_jwt_reads_claims() { let jwt = jwt_with_payload(serde_json::json!({ + "iss": AGENT_IDENTITY_JWT_ISSUER, + "aud": AGENT_IDENTITY_JWT_AUDIENCE, + "iat": 1_700_000_000usize, + "exp": 4_000_000_000usize, "agent_runtime_id": "agent-runtime-id", "agent_private_key": "private-key", "account_id": "account-id", @@ -480,12 +529,15 @@ mod tests { "chatgpt_account_is_fedramp": false, })); - let claims = - decode_agent_identity_jwt(&jwt, /*public_key_base64*/ None).expect("JWT should decode"); + let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode"); assert_eq!( claims, AgentIdentityJwtClaims { + iss: AGENT_IDENTITY_JWT_ISSUER.to_string(), + aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(), + iat: 1_700_000_000, + exp: 4_000_000_000, agent_runtime_id: "agent-runtime-id".to_string(), agent_private_key: "private-key".to_string(), account_id: "account-id".to_string(), @@ -498,15 +550,13 @@ mod tests { } #[test] - fn decode_agent_identity_jwt_verifies_when_public_key_is_present() { - let mut secret_key_bytes = [0u8; 32]; - secret_key_bytes[0] = 1; - let signing_key = SigningKey::from_bytes(&secret_key_bytes); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .expect("private key should encode"); - let public_key_base64 = BASE64_STANDARD.encode(signing_key.verifying_key().as_bytes()); + fn decode_agent_identity_jwt_verifies_when_jwks_is_present() { + let jwks = test_jwks("test-key"); let claims = AgentIdentityJwtClaims { + iss: AGENT_IDENTITY_JWT_ISSUER.to_string(), + aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(), + iat: 1_700_000_000, + exp: 4_000_000_000, agent_runtime_id: "agent-runtime-id".to_string(), agent_private_key: "private-key".to_string(), account_id: "account-id".to_string(), @@ -516,8 +566,12 @@ mod tests { chatgpt_account_is_fedramp: false, }; let jwt = jsonwebtoken::encode( - &Header::new(Algorithm::EdDSA), + &test_jwt_header("test-key"), &serde_json::json!({ + "iss": claims.iss, + "aud": claims.aud, + "iat": claims.iat, + "exp": claims.exp, "agent_runtime_id": claims.agent_runtime_id, "agent_private_key": claims.agent_private_key, "account_id": claims.account_id, @@ -526,11 +580,15 @@ mod tests { "plan_type": "pro", "chatgpt_account_is_fedramp": claims.chatgpt_account_is_fedramp, }), - &EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()), + &test_rsa_encoding_key(), ) .expect("JWT should encode"); let expected_claims = AgentIdentityJwtClaims { + iss: AGENT_IDENTITY_JWT_ISSUER.to_string(), + aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(), + iat: 1_700_000_000, + exp: 4_000_000_000, agent_runtime_id: "agent-runtime-id".to_string(), agent_private_key: "private-key".to_string(), account_id: "account-id".to_string(), @@ -540,31 +598,45 @@ mod tests { chatgpt_account_is_fedramp: false, }; assert_eq!( - decode_agent_identity_jwt(&jwt, Some(&public_key_base64)).expect("JWT should verify"), + decode_agent_identity_jwt(&jwt, Some(&jwks)).expect("JWT should verify"), expected_claims ); } #[test] - fn decode_agent_identity_jwt_rejects_wrong_public_key() { - let mut signing_secret_key_bytes = [0u8; 32]; - signing_secret_key_bytes[0] = 1; - let signing_key = SigningKey::from_bytes(&signing_secret_key_bytes); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .expect("private key should encode"); + fn decode_agent_identity_jwt_rejects_untrusted_kid() { + let jwks = test_jwks("other-key"); - let mut other_secret_key_bytes = [0u8; 32]; - other_secret_key_bytes[0] = 2; - let other_public_key_base64 = BASE64_STANDARD.encode( - SigningKey::from_bytes(&other_secret_key_bytes) - .verifying_key() - .as_bytes(), - ); + let jwt = jsonwebtoken::encode( + &test_jwt_header("test-key"), + &serde_json::json!({ + "iss": AGENT_IDENTITY_JWT_ISSUER, + "aud": AGENT_IDENTITY_JWT_AUDIENCE, + "iat": 1_700_000_000, + "exp": 4_000_000_000usize, + "agent_runtime_id": "agent-runtime-id", + "agent_private_key": "private-key", + "account_id": "account-id", + "chatgpt_user_id": "user-id", + "email": "user@example.com", + "plan_type": "pro", + "chatgpt_account_is_fedramp": false, + }), + &test_rsa_encoding_key(), + ) + .expect("JWT should encode"); + + decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify"); + } + #[test] + fn decode_agent_identity_jwt_requires_issuer_and_audience() { + let jwks = test_jwks("test-key"); let jwt = jsonwebtoken::encode( - &Header::new(Algorithm::EdDSA), + &test_jwt_header("test-key"), &serde_json::json!({ + "iat": 1_700_000_000, + "exp": 4_000_000_000usize, "agent_runtime_id": "agent-runtime-id", "agent_private_key": "private-key", "account_id": "account-id", @@ -573,12 +645,65 @@ mod tests { "plan_type": "pro", "chatgpt_account_is_fedramp": false, }), - &EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()), + &test_rsa_encoding_key(), ) .expect("JWT should encode"); - decode_agent_identity_jwt(&jwt, Some(&other_public_key_base64)) - .expect_err("JWT should not verify"); + decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify"); + } + + fn test_jwt_header(kid: &str) -> Header { + let mut header = Header::new(Algorithm::RS256); + header.kid = Some(kid.to_string()); + header + } + + fn test_rsa_encoding_key() -> EncodingKey { + EncodingKey::from_rsa_pem( + 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-----"#, + ) + .expect("test RSA key should parse") + } + + fn test_jwks(kid: &str) -> jsonwebtoken::jwk::JwkSet { + serde_json::from_value(serde_json::json!({ + "keys": [{ + "kty": "RSA", + "kid": kid, + "use": "sig", + "alg": "RS256", + "n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q", + "e": "AQAB", + }] + })) + .expect("test JWKS should parse") } #[test] @@ -587,6 +712,26 @@ mod tests { normalize_chatgpt_base_url("https://chatgpt.com/codex"), "https://chatgpt.com/backend-api" ); + assert_eq!( + normalize_chatgpt_base_url("http://localhost:8080/api/codex"), + "http://localhost:8080/api/codex" + ); + assert_eq!( + normalize_chatgpt_base_url("https://chatgpt.com/api/codex"), + "https://chatgpt.com/api/codex" + ); + } + + #[test] + fn agent_identity_jwks_url_matches_base_url_style() { + assert_eq!( + agent_identity_jwks_url("https://chatgpt.com/backend-api"), + "https://chatgpt.com/backend-api/wham/agent-identities/jwks" + ); + assert_eq!( + agent_identity_jwks_url("http://localhost:8080/api/codex"), + "http://localhost:8080/api/codex/agent-identities/jwks" + ); } fn jwt_with_payload(payload: serde_json::Value) -> String { diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 6f62624b2898..7bb79e90c712 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -207,7 +207,10 @@ pub async fn run_login_with_agent_identity( &config.codex_home, &agent_identity, config.cli_auth_credentials_store_mode, - ) { + Some(&config.chatgpt_base_url), + ) + .await + { Ok(_) => { eprintln!("{LOGIN_SUCCESS_MESSAGE}"); std::process::exit(0); @@ -366,6 +369,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { &config.codex_home, config.cli_auth_credentials_store_mode, config.agent_identity_authapi_base_url.as_deref(), + Some(&config.chatgpt_base_url), ) .await { diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs index 8f26cd51d419..7fd9f7af2771 100644 --- a/codex-rs/cli/tests/login.rs +++ b/codex-rs/cli/tests/login.rs @@ -6,8 +6,6 @@ use pretty_assertions::assert_eq; use serde_json::Value; use tempfile::TempDir; -const FAKE_AGENT_IDENTITY_JWT: &str = "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJhZ2VudF9ydW50aW1lX2lkIjoiYWdlbnQtcnVudGltZS1pZCIsImFnZW50X3ByaXZhdGVfa2V5IjoicHJpdmF0ZS1rZXkiLCJhY2NvdW50X2lkIjoiYWNjb3VudC0xMjMiLCJjaGF0Z3B0X3VzZXJfaWQiOiJ1c2VyLWlkIiwiZW1haWwiOiJ1c2VyQGV4YW1wbGUuY29tIiwicGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lzX2ZlZHJhbXAiOmZhbHNlfQ.c2ln"; - fn codex_command(codex_home: &Path) -> Result { let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); cmd.env("CODEX_HOME", codex_home); @@ -53,22 +51,16 @@ fn login_with_api_key_reads_stdin_and_writes_auth_json() -> Result<()> { } #[test] -fn login_with_agent_identity_reads_stdin_and_writes_auth_json() -> Result<()> { +fn login_with_agent_identity_rejects_invalid_jwt() -> Result<()> { let codex_home = TempDir::new()?; write_file_auth_config(codex_home.path())?; let mut cmd = codex_command(codex_home.path())?; cmd.args(["login", "--with-agent-identity"]) - .write_stdin(format!("{FAKE_AGENT_IDENTITY_JWT}\n")) + .write_stdin("not-a-jwt\n") .assert() - .success() - .stderr(contains("Successfully logged in")); - - let auth = read_auth_json(codex_home.path())?; - assert_eq!(auth["auth_mode"], "agentIdentity"); - assert_eq!(auth["agent_identity"], FAKE_AGENT_IDENTITY_JWT); - assert!(auth["OPENAI_API_KEY"].is_null()); - assert!(auth.get("tokens").is_none()); + .failure() + .stderr(contains("Error logging in with Agent Identity")); Ok(()) } diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c8dffd69b555..15f933ed3ff1 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1094,6 +1094,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { codex_home.path(), AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await { diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 1c273a4b153a..ae8907d2f4b8 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -439,6 +439,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result auth_credentials_store_mode: config.cli_auth_credentials_store_mode, forced_login_method: config.forced_login_method, forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(), + chatgpt_base_url: Some(config.chatgpt_base_url.clone()), agent_identity_authapi_base_url: config.agent_identity_authapi_base_url.clone(), }) .await diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index 026a3edf0845..161d1b862cb7 100644 --- a/codex-rs/login/Cargo.toml +++ b/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 } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index a72f31974eb3..4203926aadb3 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/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"); - - super::login_with_agent_identity(dir.path(), &agent_identity, AuthCredentialsStoreMode::File) - .expect("login_with_agent_identity should succeed"); + 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, + 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,22 @@ 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, + /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, + ) + .await + .expect_err("invalid Agent Identity token should fail"); assert_eq!(err.kind(), std::io::ErrorKind::Other); assert!( @@ -117,12 +143,45 @@ 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 _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR); let auth = CodexAuth::from_auth_storage( dir.path(), AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("call should succeed"); @@ -133,6 +192,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, @@ -148,6 +208,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .unwrap() @@ -191,6 +252,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, @@ -203,6 +265,7 @@ async fn loads_api_key_from_auth_json() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .unwrap() @@ -262,8 +325,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, @@ -279,6 +344,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") @@ -297,6 +363,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { updated_auth_dot_json, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("updated auth should parse"); @@ -600,6 +667,7 @@ async fn build_config( auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_login_method, forced_chatgpt_workspace_id, + chatgpt_base_url: None, agent_identity_authapi_base_url: None, } } @@ -621,6 +689,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)] @@ -635,6 +711,53 @@ 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 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() { @@ -649,6 +772,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() { /*enable_codex_api_key_env*/ true, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("env auth should load") @@ -661,6 +785,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"); @@ -685,6 +810,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, @@ -716,6 +842,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, @@ -747,6 +874,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"); @@ -770,6 +898,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( @@ -806,6 +935,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, @@ -819,9 +952,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 { + 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, @@ -837,6 +1038,7 @@ async fn plan_type_maps_known_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") @@ -846,8 +1048,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, @@ -863,6 +1067,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") @@ -875,8 +1080,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, @@ -892,6 +1099,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") @@ -904,8 +1112,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, @@ -921,6 +1131,7 @@ async fn plan_type_maps_unknown_to_unknown() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") @@ -930,8 +1141,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, @@ -947,6 +1160,7 @@ async fn missing_plan_type_maps_to_unknown() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("load auth") diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index e4abb3ab7d54..55762553c7d5 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -16,6 +16,9 @@ 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_agent_identity::normalize_chatgpt_base_url; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -29,6 +32,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 +92,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"; @@ -198,6 +203,7 @@ impl CodexAuth { auth_dot_json: AuthDotJson, auth_credentials_store_mode: AuthCredentialsStoreMode, agent_identity_authapi_base_url: Option<&str>, + chatgpt_base_url: Option<&str>, ) -> std::io::Result { let auth_mode = auth_dot_json.resolved_mode(); let client = create_client(); @@ -213,8 +219,12 @@ impl CodexAuth { "agent identity auth is missing an agent identity token.", )); }; - return Self::from_agent_identity_jwt(&agent_identity, agent_identity_authapi_base_url) - .await; + return Self::from_agent_identity_jwt( + &agent_identity, + agent_identity_authapi_base_url, + chatgpt_base_url, + ) + .await; } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); @@ -240,12 +250,14 @@ impl CodexAuth { codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, agent_identity_authapi_base_url: Option<&str>, + chatgpt_base_url: Option<&str>, ) -> std::io::Result> { load_auth( codex_home, /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, agent_identity_authapi_base_url, + chatgpt_base_url, ) .await } @@ -253,8 +265,10 @@ impl CodexAuth { pub async fn from_agent_identity_jwt( jwt: &str, agent_identity_authapi_base_url: Option<&str>, + chatgpt_base_url: Option<&str>, ) -> std::io::Result { - let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; + let normalized_base_url = normalized_agent_identity_base_url(chatgpt_base_url); + let record = verified_agent_identity_record(jwt, &normalized_base_url).await?; Ok(Self::AgentIdentity( AgentIdentityAuth::load(record, agent_identity_authapi_base_url).await?, )) @@ -502,6 +516,22 @@ pub fn read_codex_agent_identity_from_env() -> Option { .filter(|value| !value.is_empty()) } +fn normalized_agent_identity_base_url(chatgpt_base_url: Option<&str>) -> String { + normalize_chatgpt_base_url(chatgpt_base_url.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)) +} + +async fn verified_agent_identity_record( + jwt: &str, + 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 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( @@ -545,12 +575,14 @@ 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 normalized_base_url = normalized_agent_identity_base_url(chatgpt_base_url); + verified_agent_identity_record(agent_identity, &normalized_base_url).await?; let auth_dot_json = AuthDotJson { auth_mode: Some(ApiAuthMode::AgentIdentity), openai_api_key: None, @@ -609,6 +641,7 @@ pub struct AuthConfig { pub auth_credentials_store_mode: AuthCredentialsStoreMode, pub forced_login_method: Option, pub forced_chatgpt_workspace_id: Option, + pub chatgpt_base_url: Option, pub agent_identity_authapi_base_url: Option, } @@ -618,6 +651,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result< /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, config.agent_identity_authapi_base_url.as_deref(), + config.chatgpt_base_url.as_deref(), ) .await? else { @@ -724,6 +758,7 @@ async fn load_auth( enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, agent_identity_authapi_base_url: Option<&str>, + chatgpt_base_url: Option<&str>, ) -> std::io::Result> { // 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() { @@ -742,6 +777,7 @@ async fn load_auth( auth_dot_json, AuthCredentialsStoreMode::Ephemeral, agent_identity_authapi_base_url, + chatgpt_base_url, ) .await?; return Ok(Some(auth)); @@ -756,6 +792,7 @@ async fn load_auth( return CodexAuth::from_agent_identity_jwt( &agent_identity, agent_identity_authapi_base_url, + chatgpt_base_url, ) .await .map(Some); @@ -773,6 +810,7 @@ async fn load_auth( auth_dot_json, auth_credentials_store_mode, agent_identity_authapi_base_url, + chatgpt_base_url, ) .await?; Ok(Some(auth)) @@ -1321,6 +1359,7 @@ impl AuthManager { enable_codex_api_key_env, auth_credentials_store_mode, agent_identity_authapi_base_url.as_deref(), + chatgpt_base_url.as_deref(), ) .await .ok() @@ -1529,6 +1568,7 @@ impl AuthManager { self.enable_codex_api_key_env, self.auth_credentials_store_mode, self.agent_identity_authapi_base_url.as_deref(), + self.chatgpt_base_url.as_deref(), ) .await .ok() diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index b61ce081067d..b2e8946fa979 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/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 { - 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 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, - }) + } } } diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index d431f0525229..9acf8f2a6b42 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -232,6 +232,7 @@ c2ln", codex_home, AuthCredentialsStoreMode::File, /*agent_identity_authapi_base_url*/ None, + /*chatgpt_base_url*/ None, ) .await .expect("auth should load") diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index fbb5fad4026a..d98d3dc2762c 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -896,6 +896,7 @@ pub async fn run_main( auth_credentials_store_mode: config.cli_auth_credentials_store_mode, forced_login_method: config.forced_login_method, forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(), + chatgpt_base_url: Some(config.chatgpt_base_url.clone()), agent_identity_authapi_base_url: config.agent_identity_authapi_base_url.clone(), }) .await From b1c80703678ce9c2711528189b8f21ada7f55536 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Mon, 27 Apr 2026 14:15:39 -0700 Subject: [PATCH 2/2] fix: simplify agent identity backend URL handling --- codex-rs/agent-identity/src/lib.rs | 65 ++++----------------------- codex-rs/login/src/auth/auth_tests.rs | 2 +- codex-rs/login/src/auth/manager.rs | 19 ++++---- 3 files changed, 20 insertions(+), 66 deletions(-) diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index 12c7fd4f2807..172686f72fb0 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -129,14 +129,16 @@ pub async fn fetch_agent_identity_jwks( client: &reqwest::Client, chatgpt_base_url: &str, ) -> Result { - client + let response = client .get(agent_identity_jwks_url(chatgpt_base_url)) .timeout(AGENT_IDENTITY_JWKS_TIMEOUT) .send() .await - .context("failed to fetch agent identity JWKS")? + .context("failed to request agent identity JWKS")? .error_for_status() - .context("failed to fetch agent identity JWKS")? + .context("agent identity JWKS endpoint returned an error")?; + + response .json() .await .context("failed to decode agent identity JWKS") @@ -311,13 +313,7 @@ pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String { let trimmed = chatgpt_base_url.trim_end_matches('/'); - if trimmed.contains("/backend-api") { - format!("{trimmed}/wham/agent-identities/jwks") - } else if trimmed.contains("/api/codex") { - format!("{trimmed}/agent-identities/jwks") - } else { - format!("{trimmed}/api/codex/agent-identities/jwks") - } + format!("{trimmed}/wham/agent-identities/jwks") } pub fn agent_identity_request_id() -> Result { @@ -331,33 +327,6 @@ pub fn agent_identity_request_id() -> Result { )) } -pub fn normalize_chatgpt_base_url(chatgpt_base_url: &str) -> String { - let mut base_url = chatgpt_base_url.trim_end_matches('/').to_string(); - for suffix in [ - "/wham/remote/control/server/enroll", - "/wham/remote/control/server", - ] { - if let Some(stripped) = base_url.strip_suffix(suffix) { - base_url = stripped.to_string(); - break; - } - } - if base_url.ends_with("/codex") - && !base_url.ends_with("/api/codex") - && let Some(stripped) = base_url.strip_suffix("/codex") - { - base_url = stripped.to_string(); - } - if (base_url.starts_with("https://chatgpt.com") - || base_url.starts_with("https://chat.openai.com")) - && !base_url.contains("/backend-api") - && !base_url.contains("/api/codex") - { - base_url = format!("{base_url}/backend-api"); - } - base_url -} - pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { AgentBillOfMaterials { agent_version: env!("CARGO_PKG_VERSION").to_string(), @@ -707,30 +676,14 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN } #[test] - fn normalize_chatgpt_base_url_strips_codex_before_backend_api() { - assert_eq!( - normalize_chatgpt_base_url("https://chatgpt.com/codex"), - "https://chatgpt.com/backend-api" - ); - assert_eq!( - normalize_chatgpt_base_url("http://localhost:8080/api/codex"), - "http://localhost:8080/api/codex" - ); - assert_eq!( - normalize_chatgpt_base_url("https://chatgpt.com/api/codex"), - "https://chatgpt.com/api/codex" - ); - } - - #[test] - fn agent_identity_jwks_url_matches_base_url_style() { + fn agent_identity_jwks_url_uses_backend_api_base_url() { assert_eq!( agent_identity_jwks_url("https://chatgpt.com/backend-api"), "https://chatgpt.com/backend-api/wham/agent-identities/jwks" ); assert_eq!( - agent_identity_jwks_url("http://localhost:8080/api/codex"), - "http://localhost:8080/api/codex/agent-identities/jwks" + agent_identity_jwks_url("https://chatgpt.com/backend-api/"), + "https://chatgpt.com/backend-api/wham/agent-identities/jwks" ); } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 4203926aadb3..ebefb214409f 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -129,7 +129,6 @@ async fn login_with_agent_identity_rejects_invalid_jwt() { dir.path(), "not-a-jwt", AuthCredentialsStoreMode::File, - /*agent_identity_authapi_base_url*/ None, /*chatgpt_base_url*/ None, ) .await @@ -740,6 +739,7 @@ async fn load_auth_reads_agent_identity_from_env() { codex_home.path(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*agent_identity_authapi_base_url*/ Some(&chatgpt_base_url), Some(&chatgpt_base_url), ) .await diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 55762553c7d5..e96c5c8b7f9f 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -18,7 +18,6 @@ use tokio::sync::Semaphore; use codex_agent_identity::decode_agent_identity_jwt; use codex_agent_identity::fetch_agent_identity_jwks; -use codex_agent_identity::normalize_chatgpt_base_url; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -267,8 +266,11 @@ impl CodexAuth { agent_identity_authapi_base_url: Option<&str>, chatgpt_base_url: Option<&str>, ) -> std::io::Result { - let normalized_base_url = normalized_agent_identity_base_url(chatgpt_base_url); - let record = verified_agent_identity_record(jwt, &normalized_base_url).await?; + 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, agent_identity_authapi_base_url).await?, )) @@ -516,10 +518,6 @@ pub fn read_codex_agent_identity_from_env() -> Option { .filter(|value| !value.is_empty()) } -fn normalized_agent_identity_base_url(chatgpt_base_url: Option<&str>) -> String { - normalize_chatgpt_base_url(chatgpt_base_url.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)) -} - async fn verified_agent_identity_record( jwt: &str, chatgpt_base_url: &str, @@ -581,8 +579,11 @@ pub async fn login_with_agent_identity( auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option<&str>, ) -> std::io::Result<()> { - let normalized_base_url = normalized_agent_identity_base_url(chatgpt_base_url); - verified_agent_identity_record(agent_identity, &normalized_base_url).await?; + 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,