From f6797c3ac603ff558829c5cd5a27f95e1f4801b1 Mon Sep 17 00:00:00 2001 From: efrazer-oai Date: Tue, 28 Apr 2026 09:56:20 -0700 Subject: [PATCH] feat: verify agent identity JWTs with JWKS (#19764) --- codex-rs/Cargo.lock | 1 + codex-rs/agent-identity/src/lib.rs | 258 +++++++++++++------ codex-rs/cli/src/login.rs | 13 +- codex-rs/cli/tests/login.rs | 16 +- codex-rs/core/tests/suite/client.rs | 18 +- codex-rs/exec/src/lib.rs | 1 + codex-rs/login/Cargo.toml | 1 + codex-rs/login/src/auth/agent_identity.rs | 4 +- codex-rs/login/src/auth/auth_tests.rs | 244 +++++++++++++++++- codex-rs/login/src/auth/manager.rs | 59 ++++- codex-rs/login/src/auth/storage.rs | 15 +- codex-rs/models-manager/src/manager_tests.rs | 12 +- codex-rs/tui/src/lib.rs | 1 + 13 files changed, 517 insertions(+), 126 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e3f2ac1a1..51667c601 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2850,6 +2850,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 69801f6fc..cdfb258f4 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,49 @@ 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 { + let response = client + .get(agent_identity_jwks_url(chatgpt_base_url)) + .timeout(AGENT_IDENTITY_JWKS_TIMEOUT) + .send() + .await + .context("failed to request agent identity JWKS")? + .error_for_status() + .context("agent identity JWKS endpoint returned an error")?; + + response + .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 +311,15 @@ 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 { + format!("{trimmed}/agent-identities/jwks") + } +} + pub fn agent_identity_request_id() -> Result { let mut request_id_bytes = [0u8; 16]; OsRng @@ -290,29 +331,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 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 = format!("{base_url}/backend-api"); - } - base_url -} - pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { AgentBillOfMaterials { agent_version: env!("CARGO_PKG_VERSION").to_string(), @@ -472,6 +490,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", @@ -481,12 +503,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(), @@ -499,15 +524,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(), @@ -517,8 +540,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, @@ -527,11 +554,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(), @@ -541,31 +572,22 @@ 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"); - - 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(), - ); + fn decode_agent_identity_jwt_rejects_untrusted_kid() { + let jwks = test_jwks("other-key"); let jwt = jsonwebtoken::encode( - &Header::new(Algorithm::EdDSA), + &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", @@ -574,19 +596,111 @@ 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"); } #[test] - fn normalize_chatgpt_base_url_strips_codex_before_backend_api() { + fn decode_agent_identity_jwt_requires_issuer_and_audience() { + let jwks = test_jwks("test-key"); + let jwt = jsonwebtoken::encode( + &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", + "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"); + } + + 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] + fn agent_identity_jwks_url_uses_backend_api_base_url() { assert_eq!( - normalize_chatgpt_base_url("https://chatgpt.com/codex"), - "https://chatgpt.com/backend-api" + agent_identity_jwks_url("https://chatgpt.com/backend-api"), + "https://chatgpt.com/backend-api/wham/agent-identities/jwks" + ); + assert_eq!( + agent_identity_jwks_url("https://chatgpt.com/backend-api/"), + "https://chatgpt.com/backend-api/wham/agent-identities/jwks" + ); + } + + #[test] + fn agent_identity_jwks_url_uses_codex_api_base_url() { + assert_eq!( + agent_identity_jwks_url("http://localhost:8080/api/codex"), + "http://localhost:8080/api/codex/agent-identities/jwks" + ); + assert_eq!( + agent_identity_jwks_url("http://localhost:8080/api/codex/"), + "http://localhost:8080/api/codex/agent-identities/jwks" ); } diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 5c60c20cb..1baa344b8 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); @@ -362,8 +365,12 @@ pub async fn run_login_with_device_code_fallback_to_browser( pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; - match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode) - .await + match CodexAuth::from_auth_storage( + &config.codex_home, + config.cli_auth_credentials_store_mode, + Some(&config.chatgpt_base_url), + ) + .await { Ok(Some(auth)) => match auth.auth_mode() { AuthMode::ApiKey => match auth.get_token() { diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs index 8f26cd51d..7fd9f7af2 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 7bd5cee68..67340f316 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1090,13 +1090,17 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { let mut config = load_default_config_for_test(&codex_home).await; config.model_provider = model_provider; - let auth_manager = - match CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File).await - { - Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth), - Ok(None) => panic!("No CodexAuth found in codex_home"), - Err(e) => panic!("Failed to load CodexAuth: {e}"), - }; + let auth_manager = match CodexAuth::from_auth_storage( + codex_home.path(), + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + { + Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth), + Ok(None) => panic!("No CodexAuth found in codex_home"), + Err(e) => panic!("Failed to load CodexAuth: {e}"), + }; let thread_manager = ThreadManager::new( &config, auth_manager, diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index fc3392db2..d4044a728 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()), }) .await { diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index 026a3edf0..161d1b862 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/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index fd5417428..364471332 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/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!( diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index b3fe02125..90748f786 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"); + 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 { + 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") diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index ed4c0f01d..9f140c276 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/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 { 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> { 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 { - 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 { + 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 { .filter(|value| !value.is_empty()) } +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( @@ -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, pub forced_chatgpt_workspace_id: Option, + pub chatgpt_base_url: Option, } 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> { // 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() diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index b61ce0810..b2e8946fa 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 0f0998034..9519f9713 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/codex-rs/models-manager/src/manager_tests.rs @@ -228,10 +228,14 @@ c2ln", ) .expect("auth.json should be written"); - CodexAuth::from_auth_storage(codex_home, AuthCredentialsStoreMode::File) - .await - .expect("auth should load") - .expect("auth should be present") + CodexAuth::from_auth_storage( + codex_home, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await + .expect("auth should load") + .expect("auth should be present") } #[tokio::test] diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index b0bfe7c30..d09876739 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -900,6 +900,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()), }) .await {