From fed0a8f4faa58db3138488cca77628c1d54a2cd8 Mon Sep 17 00:00:00 2001 From: efrazer-oai Date: Sun, 26 Apr 2026 12:49:54 -0700 Subject: [PATCH] feat: load AgentIdentity from JWT login/env (#18904) ## Summary This PR lets programmatic AgentIdentity users provide one token through either stdin login or environment auth. `codex login --with-agent-identity` reads an Agent Identity JWT from stdin, validates that it has the required claims, and stores that token as the `agent_identity` value in `auth.json`. The file format is token-only; the decoded account and key fields are runtime state, not hand-authored auth.json fields. The Agent Identity JWT claim shape and decoder live in `codex-agent-identity`; `codex-login` only owns env/storage precedence and conversion into `CodexAuth::AgentIdentity`. When env auth is enabled, `CODEX_AGENT_IDENTITY` can provide the same JWT without writing auth state to disk. `CODEX_API_KEY` still wins if both env vars are set. Reference old stack: https://github.com/openai/codex/pull/17387/changes Reference JWT/env stack: https://github.com/openai/codex/pull/18176 ## Stack 1. https://github.com/openai/codex/pull/18757: full revert 2. https://github.com/openai/codex/pull/18871: isolated Agent Identity crate 3. https://github.com/openai/codex/pull/18785: explicit AgentIdentity auth mode and startup task allocation 4. https://github.com/openai/codex/pull/18811: migrate Codex backend auth callsites through AuthProvider 5. This PR: accept AgentIdentity JWTs through login/env ## Testing Tests: targeted login and Agent Identity crate tests, CLI checks, scoped formatter/linter cleanup, and CI. --------- Co-authored-by: Shijie Rao --- codex-rs/Cargo.lock | 1 + codex-rs/agent-identity/Cargo.toml | 1 + codex-rs/agent-identity/src/lib.rs | 199 +++++++++++++++++++++- codex-rs/cli/src/lib.rs | 2 + codex-rs/cli/src/login.rs | 64 ++++++- codex-rs/cli/src/main.rs | 19 ++- codex-rs/cli/tests/login.rs | 74 ++++++++ codex-rs/cloud-requirements/src/lib.rs | 5 + codex-rs/login/src/auth/agent_identity.rs | 20 +-- codex-rs/login/src/auth/auth_tests.rs | 125 +++++++++++++- codex-rs/login/src/auth/manager.rs | 44 ++++- codex-rs/login/src/auth/storage.rs | 20 ++- codex-rs/login/src/auth/storage_tests.rs | 59 +++++-- codex-rs/login/src/lib.rs | 3 + 14 files changed, 587 insertions(+), 49 deletions(-) create mode 100644 codex-rs/cli/tests/login.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5ff3f462f..fd4ed6d8d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1758,6 +1758,7 @@ dependencies = [ "codex-protocol", "crypto_box", "ed25519-dalek", + "jsonwebtoken", "pretty_assertions", "rand 0.9.3", "reqwest", diff --git a/codex-rs/agent-identity/Cargo.toml b/codex-rs/agent-identity/Cargo.toml index 7976c3354..4610d6ec9 100644 --- a/codex-rs/agent-identity/Cargo.toml +++ b/codex-rs/agent-identity/Cargo.toml @@ -19,6 +19,7 @@ chrono = { workspace = true } codex-protocol = { workspace = true } crypto_box = { workspace = true } ed25519-dalek = { workspace = true } +jsonwebtoken = { workspace = true } rand = { workspace = true } reqwest = { workspace = true, features = ["json"] } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index a6d7e25df..bf139f787 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -8,6 +8,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::SecondsFormat; use chrono::Utc; +use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::protocol::SessionSource; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; @@ -15,10 +16,14 @@ use ed25519_dalek::SigningKey; use ed25519_dalek::VerifyingKey; use ed25519_dalek::pkcs8::DecodePrivateKey; use ed25519_dalek::pkcs8::EncodePrivateKey; +use jsonwebtoken::Algorithm; +use jsonwebtoken::DecodingKey; +use jsonwebtoken::Validation; use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; use serde::Serialize; +use serde::de::DeserializeOwned; use sha2::Digest as _; use sha2::Sha512; @@ -50,6 +55,18 @@ pub struct GeneratedAgentKeyMaterial { pub public_key_ssh: String, } +/// Claims carried by an Agent Identity JWT. +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub struct AgentIdentityJwtClaims { + pub agent_runtime_id: String, + pub agent_private_key: String, + pub account_id: String, + pub chatgpt_user_id: String, + pub email: String, + pub plan_type: AccountPlanType, + pub chatgpt_account_is_fedramp: bool, +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] struct AgentAssertionEnvelope { agent_runtime_id: String, @@ -98,6 +115,43 @@ pub fn authorization_header_for_agent_task( Ok(format!("AgentAssertion {serialized_assertion}")) } +pub fn decode_agent_identity_jwt( + jwt: &str, + public_key_base64: Option<&str>, +) -> Result { + let Some(public_key_base64) = public_key_base64 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) + .map(|data| data.claims) + .context("failed to decode agent identity JWT") +} + +fn decode_agent_identity_jwt_payload(jwt: &str) -> Result { + let mut parts = jwt.split('.'); + let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) { + (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s), + _ => anyhow::bail!("invalid agent identity JWT format"), + }; + anyhow::ensure!(parts.next().is_none(), "invalid agent identity JWT format"); + + let payload_bytes = URL_SAFE_NO_PAD + .decode(payload_b64) + .context("agent identity JWT payload is not valid base64url")?; + serde_json::from_slice(&payload_bytes).context("agent identity JWT payload is not valid JSON") +} + pub fn sign_task_registration_payload( key: AgentIdentityKey<'_>, timestamp: &str, @@ -117,19 +171,27 @@ pub async fn register_agent_task( signature: sign_task_registration_payload(key, ×tamp)?, timestamp, }; + let url = agent_task_registration_url(chatgpt_base_url, key.agent_runtime_id); let response = client - .post(agent_task_registration_url( - chatgpt_base_url, - key.agent_runtime_id, - )) + .post(url) .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) .json(&request) .send() .await - .context("failed to register agent task")? - .error_for_status() - .context("failed to register agent task")? + .context("failed to register agent task")?; + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + let body = if body.len() > 512 { + format!("{}...", body.chars().take(512).collect::()) + } else { + body + }; + anyhow::bail!("failed to register agent task with status {status}: {body}"); + } + + let response = response .json() .await .context("failed to decode agent task registration response")?; @@ -323,6 +385,8 @@ mod tests { use base64::Engine as _; use ed25519_dalek::Signature; use ed25519_dalek::Verifier as _; + use jsonwebtoken::EncodingKey; + use jsonwebtoken::Header; use pretty_assertions::assert_eq; use super::*; @@ -404,6 +468,119 @@ mod tests { ); } + #[test] + fn decode_agent_identity_jwt_reads_claims() { + let jwt = jwt_with_payload(serde_json::json!({ + "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, + })); + + let claims = + decode_agent_identity_jwt(&jwt, /*public_key_base64*/ None).expect("JWT should decode"); + + assert_eq!( + claims, + AgentIdentityJwtClaims { + agent_runtime_id: "agent-runtime-id".to_string(), + agent_private_key: "private-key".to_string(), + account_id: "account-id".to_string(), + chatgpt_user_id: "user-id".to_string(), + email: "user@example.com".to_string(), + plan_type: AccountPlanType::Pro, + chatgpt_account_is_fedramp: false, + } + ); + } + + #[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()); + let claims = AgentIdentityJwtClaims { + agent_runtime_id: "agent-runtime-id".to_string(), + agent_private_key: "private-key".to_string(), + account_id: "account-id".to_string(), + chatgpt_user_id: "user-id".to_string(), + email: "user@example.com".to_string(), + plan_type: AccountPlanType::Pro, + chatgpt_account_is_fedramp: false, + }; + let jwt = jsonwebtoken::encode( + &Header::new(Algorithm::EdDSA), + &serde_json::json!({ + "agent_runtime_id": claims.agent_runtime_id, + "agent_private_key": claims.agent_private_key, + "account_id": claims.account_id, + "chatgpt_user_id": claims.chatgpt_user_id, + "email": claims.email, + "plan_type": "pro", + "chatgpt_account_is_fedramp": claims.chatgpt_account_is_fedramp, + }), + &EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()), + ) + .expect("JWT should encode"); + + let expected_claims = AgentIdentityJwtClaims { + agent_runtime_id: "agent-runtime-id".to_string(), + agent_private_key: "private-key".to_string(), + account_id: "account-id".to_string(), + chatgpt_user_id: "user-id".to_string(), + email: "user@example.com".to_string(), + plan_type: AccountPlanType::Pro, + chatgpt_account_is_fedramp: false, + }; + assert_eq!( + decode_agent_identity_jwt(&jwt, Some(&public_key_base64)).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(), + ); + + let jwt = jsonwebtoken::encode( + &Header::new(Algorithm::EdDSA), + &serde_json::json!({ + "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, + }), + &EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()), + ) + .expect("JWT should encode"); + + decode_agent_identity_jwt(&jwt, Some(&other_public_key_base64)) + .expect_err("JWT should not verify"); + } + #[test] fn normalize_chatgpt_base_url_strips_codex_before_backend_api() { assert_eq!( @@ -411,4 +588,12 @@ mod tests { "https://chatgpt.com/backend-api" ); } + + fn jwt_with_payload(payload: serde_json::Value) -> String { + let encode = |bytes: &[u8]| URL_SAFE_NO_PAD.encode(bytes); + let header_b64 = encode(br#"{"alg":"none","typ":"JWT"}"#); + let payload_b64 = encode(&serde_json::to_vec(&payload).expect("payload should serialize")); + let signature_b64 = encode(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") + } } diff --git a/codex-rs/cli/src/lib.rs b/codex-rs/cli/src/lib.rs index cac34b3b6..3f3448c64 100644 --- a/codex-rs/cli/src/lib.rs +++ b/codex-rs/cli/src/lib.rs @@ -9,8 +9,10 @@ use codex_utils_cli::CliConfigOverrides; pub use debug_sandbox::run_command_under_landlock; pub use debug_sandbox::run_command_under_seatbelt; pub use debug_sandbox::run_command_under_windows; +pub use login::read_agent_identity_from_stdin; pub use login::read_api_key_from_stdin; pub use login::run_login_status; +pub use login::run_login_with_agent_identity; pub use login::run_login_with_api_key; pub use login::run_login_with_chatgpt; pub use login::run_login_with_device_code; diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 42241aa93..4fa7272ae 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -13,6 +13,7 @@ use codex_core::config::Config; use codex_login::CLIENT_ID; use codex_login::CodexAuth; use codex_login::ServerOptions; +use codex_login::login_with_agent_identity; use codex_login::login_with_api_key; use codex_login::logout_with_revoke; use codex_login::run_device_code_login; @@ -34,6 +35,8 @@ const CHATGPT_LOGIN_DISABLED_MESSAGE: &str = "ChatGPT login is disabled. Use API key login instead."; const API_KEY_LOGIN_DISABLED_MESSAGE: &str = "API key login is disabled. Use ChatGPT login instead."; +const AGENT_IDENTITY_LOGIN_DISABLED_MESSAGE: &str = + "Agent Identity login is disabled. Use API key login instead."; const LOGIN_SUCCESS_MESSAGE: &str = "Successfully logged in"; /// Installs a small file-backed tracing layer for direct `codex login` flows. @@ -187,31 +190,74 @@ pub async fn run_login_with_api_key( } } +pub async fn run_login_with_agent_identity( + cli_config_overrides: CliConfigOverrides, + agent_identity: String, +) -> ! { + let config = load_config_or_exit(cli_config_overrides).await; + let _login_log_guard = init_login_file_logging(&config); + tracing::info!("starting agent identity login flow"); + + if matches!(config.forced_login_method, Some(ForcedLoginMethod::Api)) { + eprintln!("{AGENT_IDENTITY_LOGIN_DISABLED_MESSAGE}"); + std::process::exit(1); + } + + match login_with_agent_identity( + &config.codex_home, + &agent_identity, + config.cli_auth_credentials_store_mode, + ) { + Ok(_) => { + eprintln!("{LOGIN_SUCCESS_MESSAGE}"); + std::process::exit(0); + } + Err(e) => { + eprintln!("Error logging in with Agent Identity: {e}"); + std::process::exit(1); + } + } +} + pub fn read_api_key_from_stdin() -> String { + read_stdin_secret( + "--with-api-key expects the API key on stdin. Try piping it, e.g. `printenv OPENAI_API_KEY | codex login --with-api-key`.", + "Reading API key from stdin...", + "No API key provided via stdin.", + ) +} + +pub fn read_agent_identity_from_stdin() -> String { + read_stdin_secret( + "--with-agent-identity expects the Agent Identity token on stdin. Try piping it, e.g. `printenv CODEX_AGENT_IDENTITY | codex login --with-agent-identity`.", + "Reading Agent Identity token from stdin...", + "No Agent Identity token provided via stdin.", + ) +} + +fn read_stdin_secret(terminal_message: &str, reading_message: &str, empty_message: &str) -> String { let mut stdin = std::io::stdin(); if stdin.is_terminal() { - eprintln!( - "--with-api-key expects the API key on stdin. Try piping it, e.g. `printenv OPENAI_API_KEY | codex login --with-api-key`." - ); + eprintln!("{terminal_message}"); std::process::exit(1); } - eprintln!("Reading API key from stdin..."); + eprintln!("{reading_message}"); let mut buffer = String::new(); if let Err(err) = stdin.read_to_string(&mut buffer) { - eprintln!("Failed to read API key from stdin: {err}"); + eprintln!("Failed to read stdin: {err}"); std::process::exit(1); } - let api_key = buffer.trim().to_string(); - if api_key.is_empty() { - eprintln!("No API key provided via stdin."); + let secret = buffer.trim().to_string(); + if secret.is_empty() { + eprintln!("{empty_message}"); std::process::exit(1); } - api_key + secret } /// Login using the OAuth device code flow. diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 35005f5be..415769e36 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -10,8 +10,10 @@ use codex_chatgpt::apply_command::run_apply_command; use codex_cli::LandlockCommand; use codex_cli::SeatbeltCommand; use codex_cli::WindowsCommand; +use codex_cli::read_agent_identity_from_stdin; use codex_cli::read_api_key_from_stdin; use codex_cli::run_login_status; +use codex_cli::run_login_with_agent_identity; use codex_cli::run_login_with_api_key; use codex_cli::run_login_with_chatgpt; use codex_cli::run_login_with_device_code; @@ -359,6 +361,12 @@ struct LoginCommand { )] with_api_key: bool, + #[arg( + long = "with-agent-identity", + help = "Read the experimental Agent Identity token from stdin (e.g. `printenv CODEX_AGENT_IDENTITY | codex login --with-agent-identity`)" + )] + with_agent_identity: bool, + #[arg( long = "api-key", num_args = 0..=1, @@ -940,7 +948,12 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { run_login_status(login_cli.config_overrides).await; } None => { - if login_cli.use_device_code { + if login_cli.with_api_key && login_cli.with_agent_identity { + eprintln!( + "Choose one login credential source: --with-api-key or --with-agent-identity." + ); + std::process::exit(1); + } else if login_cli.use_device_code { run_login_with_device_code( login_cli.config_overrides, login_cli.issuer_base_url, @@ -955,6 +968,10 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { } else if login_cli.with_api_key { let api_key = read_api_key_from_stdin(); run_login_with_api_key(login_cli.config_overrides, api_key).await; + } else if login_cli.with_agent_identity { + let agent_identity = read_agent_identity_from_stdin(); + run_login_with_agent_identity(login_cli.config_overrides, agent_identity) + .await; } else { run_login_with_chatgpt(login_cli.config_overrides).await; } diff --git a/codex-rs/cli/tests/login.rs b/codex-rs/cli/tests/login.rs new file mode 100644 index 000000000..8f26cd51d --- /dev/null +++ b/codex-rs/cli/tests/login.rs @@ -0,0 +1,74 @@ +use std::path::Path; + +use anyhow::Result; +use predicates::str::contains; +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); + Ok(cmd) +} + +fn write_file_auth_config(codex_home: &Path) -> Result<()> { + std::fs::write( + codex_home.join("config.toml"), + "cli_auth_credentials_store = \"file\"\n", + )?; + Ok(()) +} + +fn read_auth_json(codex_home: &Path) -> Result { + let auth_json = std::fs::read_to_string(codex_home.join("auth.json"))?; + Ok(serde_json::from_str(&auth_json)?) +} + +#[test] +fn login_with_api_key_reads_stdin_and_writes_auth_json() -> Result<()> { + let codex_home = TempDir::new()?; + write_file_auth_config(codex_home.path())?; + + let mut cmd = codex_command(codex_home.path())?; + cmd.args([ + "-c", + "forced_login_method=\"api\"", + "login", + "--with-api-key", + ]) + .write_stdin("sk-test\n") + .assert() + .success() + .stderr(contains("Successfully logged in")); + + let auth = read_auth_json(codex_home.path())?; + assert_eq!(auth["OPENAI_API_KEY"], "sk-test"); + assert!(auth.get("tokens").is_none()); + assert!(auth.get("agent_identity").is_none()); + + Ok(()) +} + +#[test] +fn login_with_agent_identity_reads_stdin_and_writes_auth_json() -> 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")) + .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()); + + Ok(()) +} diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 8c51888a1..1d9975f12 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -329,6 +329,11 @@ impl CloudRequirementsService { let Some(auth) = self.auth_manager.auth().await else { return Ok(None); }; + if matches!(auth, CodexAuth::AgentIdentity(_)) { + // AgentIdentity does not carry a human bearer token, and identity-edge + // only allowlists task-scoped AgentAssertion calls for the Codex runtime. + return Ok(None); + } let Some(plan_type) = auth.account_plan_type() else { return Ok(None); }; diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 5f2dc9cfc..23bbdb504 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use codex_agent_identity::AgentIdentityKey; -use codex_agent_identity::normalize_chatgpt_base_url; use codex_agent_identity::register_agent_task; use codex_protocol::account::PlanType as AccountPlanType; use tokio::sync::OnceCell; @@ -10,7 +9,7 @@ use crate::default_client::build_reqwest_client; use super::storage::AgentIdentityAuthRecord; -const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api"; +const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts"; #[derive(Debug)] pub struct AgentIdentityAuth { @@ -43,17 +42,16 @@ impl AgentIdentityAuth { self.process_task_id.get().map(String::as_str) } - pub async fn ensure_runtime(&self, chatgpt_base_url: Option) -> std::io::Result<()> { + pub async fn ensure_runtime(&self) -> std::io::Result<()> { self.process_task_id .get_or_try_init(|| async { - let base_url = normalize_chatgpt_base_url( - chatgpt_base_url - .as_deref() - .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL), - ); - register_agent_task(&build_reqwest_client(), &base_url, self.key()) - .await - .map_err(std::io::Error::other) + register_agent_task( + &build_reqwest_client(), + AGENT_IDENTITY_AUTHAPI_BASE_URL, + self.key(), + ) + .await + .map_err(std::io::Error::other) }) .await .map(|_| ()) diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 6f17822e7..b38a40b6f 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -78,6 +78,44 @@ 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() { + 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 storage = FileAuthStorage::new(dir.path().to_path_buf()); + let auth = storage + .try_read_auth_json(&auth_path) + .expect("auth.json should parse"); + assert_eq!(auth.auth_mode, Some(AuthMode::AgentIdentity)); + assert_eq!( + auth.agent_identity.as_deref(), + Some(agent_identity.as_str()) + ); + assert!(auth.tokens.is_none(), "tokens should be cleared"); + assert!(auth.openai_api_key.is_none(), "API key should be cleared"); +} + +#[test] +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"); + + assert_eq!(err.kind(), std::io::ErrorKind::Other); + assert!( + !get_auth_file(dir.path()).exists(), + "invalid Agent Identity token should not write auth.json" + ); +} + #[test] fn missing_auth_json_returns_none() { let dir = tempdir().unwrap(); @@ -87,7 +125,7 @@ fn missing_auth_json_returns_none() { } #[tokio::test] -#[serial(codex_api_key)] +#[serial(codex_auth_env)] async fn pro_account_with_no_api_key_uses_chatgpt_auth() { let codex_home = tempdir().unwrap(); let fake_jwt = write_auth_file( @@ -143,7 +181,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { } #[tokio::test] -#[serial(codex_api_key)] +#[serial(codex_auth_env)] async fn loads_api_key_from_auth_json() { let dir = tempdir().unwrap(); let auth_file = dir.path().join("auth.json"); @@ -581,7 +619,54 @@ impl Drop for EnvVarGuard { } } +#[test] +#[serial(codex_auth_env)] +fn load_auth_reads_agent_identity_from_env() { + let codex_home = tempdir().unwrap(); + let expected_record = agent_identity_record("account-123"); + let agent_identity = fake_agent_identity_jwt(&expected_record).expect("fake agent identity"); + let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity); + + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ) + .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!( + !get_auth_file(codex_home.path()).exists(), + "env auth should not write auth.json" + ); +} + +#[test] +#[serial(codex_auth_env)] +fn load_auth_keeps_codex_api_key_env_precedence() { + let codex_home = tempdir().unwrap(); + let record = agent_identity_record("account-123"); + let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity"); + let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity); + let _api_key_guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env"); + + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ true, + AuthCredentialsStoreMode::File, + ) + .expect("env auth should load") + .expect("env auth should be present"); + + assert_eq!(auth.api_key(), Some("sk-env")); +} + #[tokio::test] +#[serial(codex_auth_env)] async fn enforce_login_restrictions_logs_out_for_method_mismatch() { let codex_home = tempdir().unwrap(); login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File) @@ -604,7 +689,7 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() { } #[tokio::test] -#[serial(codex_api_key)] +#[serial(codex_auth_env)] async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( @@ -634,7 +719,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { } #[tokio::test] -#[serial(codex_api_key)] +#[serial(codex_auth_env)] async fn enforce_login_restrictions_allows_matching_workspace() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( @@ -662,6 +747,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() { } #[tokio::test] +#[serial(codex_auth_env)] 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(); @@ -683,7 +769,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f } #[tokio::test] -#[serial(codex_api_key)] +#[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 codex_home = tempdir().unwrap(); @@ -703,6 +789,35 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() { ); } +fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { + AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-id".to_string(), + agent_private_key: "private-key".to_string(), + account_id: account_id.to_string(), + chatgpt_user_id: "user-id".to_string(), + email: "user@example.com".to_string(), + plan_type: AccountPlanType::Pro, + chatgpt_account_is_fedramp: false, + } +} + +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!({ + "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, + }); + let payload_b64 = encode(&serde_json::to_vec(&payload)?); + let signature_b64 = encode(b"sig"); + Ok(format!("{header_b64}.{payload_b64}.{signature_b64}")) +} + #[test] fn plan_type_maps_known_plan() { let codex_home = tempdir().unwrap(); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 419c6a4ba..85c9b6f02 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -207,12 +207,12 @@ impl CodexAuth { return Ok(Self::from_api_key(api_key)); } if auth_mode == ApiAuthMode::AgentIdentity { - let Some(record) = auth_dot_json.agent_identity else { + let Some(agent_identity) = auth_dot_json.agent_identity else { return Err(std::io::Error::other( - "agent identity auth is missing an agent identity record.", + "agent identity auth is missing an agent identity token.", )); }; - return Ok(Self::AgentIdentity(AgentIdentityAuth::new(record))); + return Self::from_agent_identity_jwt(&agent_identity); } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); @@ -245,6 +245,11 @@ impl CodexAuth { ) } + pub fn from_agent_identity_jwt(jwt: &str) -> std::io::Result { + let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; + Ok(Self::AgentIdentity(AgentIdentityAuth::new(record))) + } + pub fn auth_mode(&self) -> AuthMode { match self { Self::ApiKey(_) => AuthMode::ApiKey, @@ -318,10 +323,10 @@ impl CodexAuth { pub async fn initialize_runtime( &self, - chatgpt_base_url: Option, + _chatgpt_base_url: Option, ) -> std::io::Result<()> { match self { - Self::AgentIdentity(auth) => auth.ensure_runtime(chatgpt_base_url).await, + Self::AgentIdentity(auth) => auth.ensure_runtime().await, Self::ApiKey(_) | Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => Ok(()), } } @@ -474,6 +479,7 @@ impl ChatgptAuth { pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; +pub const CODEX_AGENT_IDENTITY_ENV_VAR: &str = "CODEX_AGENT_IDENTITY"; pub fn read_openai_api_key_from_env() -> Option { env::var(OPENAI_API_KEY_ENV_VAR) @@ -489,6 +495,13 @@ pub fn read_codex_api_key_from_env() -> Option { .filter(|value| !value.is_empty()) } +pub fn read_codex_agent_identity_from_env() -> Option { + env::var(CODEX_AGENT_IDENTITY_ENV_VAR) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + /// 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( @@ -529,6 +542,23 @@ pub fn login_with_api_key( save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } +/// Writes an `auth.json` that contains only the Agent Identity token. +pub fn login_with_agent_identity( + codex_home: &Path, + agent_identity: &str, + auth_credentials_store_mode: AuthCredentialsStoreMode, +) -> std::io::Result<()> { + AgentIdentityAuthRecord::from_agent_identity_jwt(agent_identity)?; + let auth_dot_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::AgentIdentity), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: Some(agent_identity.to_string()), + }; + save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) +} + /// Writes an in-memory auth payload for externally managed ChatGPT tokens. pub fn login_with_chatgpt_auth_tokens( codex_home: &Path, @@ -714,6 +744,10 @@ fn load_auth( return Ok(None); } + if let Some(agent_identity) = read_codex_agent_identity_from_env() { + return CodexAuth::from_agent_identity_jwt(&agent_identity).map(Some); + } + // Fall back to the configured persistent store (file/keyring/auto) for managed auth. let storage = create_auth_storage(codex_home.to_path_buf(), auth_credentials_store_mode); let auth_dot_json = match storage.load()? { diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index e2e801169..b61ce0810 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::decode_agent_identity_jwt; use codex_app_server_protocol::AuthMode; use codex_config::types::AuthCredentialsStoreMode; use codex_keyring_store::DefaultKeyringStore; @@ -42,7 +43,7 @@ pub struct AuthDotJson { pub last_refresh: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_identity: Option, + pub agent_identity: Option, } #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] @@ -56,6 +57,23 @@ pub struct AgentIdentityAuthRecord { pub chatgpt_account_is_fedramp: bool, } +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)?; + + Ok(Self { + agent_runtime_id: claims.agent_runtime_id, + agent_private_key: claims.agent_private_key, + account_id: claims.account_id, + chatgpt_user_id: claims.chatgpt_user_id, + email: claims.email, + plan_type: claims.plan_type, + chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, + }) + } +} + pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { codex_home.join("auth.json") } diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index c06a8cfde..b5646ef53 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -7,7 +7,6 @@ use serde_json::json; use tempfile::tempdir; use codex_keyring_store::tests::MockKeyringStore; -use codex_protocol::account::PlanType as AccountPlanType; use keyring::Error as KeyringError; #[tokio::test] @@ -59,20 +58,21 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let agent_identity = jwt_with_payload(json!({ + "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, + })); let auth_dot_json = AuthDotJson { auth_mode: Some(AuthMode::AgentIdentity), openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(AgentIdentityAuthRecord { - agent_runtime_id: "agent-runtime-id".to_string(), - agent_private_key: "private-key".to_string(), - account_id: "account-id".to_string(), - chatgpt_user_id: "user-id".to_string(), - email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, - chatgpt_account_is_fedramp: false, - }), + agent_identity: Some(agent_identity), }; storage.save(&auth_dot_json)?; @@ -82,6 +82,37 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn file_storage_loads_agent_identity_as_jwt() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let agent_identity_jwt = jwt_with_payload(json!({ + "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, + })); + let auth_file = get_auth_file(codex_home.path()); + std::fs::write( + &auth_file, + serde_json::to_string_pretty(&json!({ + "auth_mode": "agentIdentity", + "agent_identity": agent_identity_jwt, + }))?, + )?; + + let loaded = storage.load()?; + + assert_eq!( + loaded.expect("auth should load").agent_identity.as_deref(), + Some(agent_identity_jwt.as_str()) + ); + Ok(()) +} + #[test] fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; @@ -217,6 +248,14 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { } } +fn jwt_with_payload(payload: serde_json::Value) -> String { + 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_b64 = encode(&serde_json::to_vec(&payload).expect("payload should serialize")); + let signature_b64 = encode(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") +} + #[test] fn keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Result<()> { let codex_home = tempdir()?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index d69a77a97..3049b6f6b 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -22,6 +22,7 @@ pub use auth::AuthDotJson; pub use auth::AuthManager; pub use auth::AuthManagerConfig; pub use auth::CLIENT_ID; +pub use auth::CODEX_AGENT_IDENTITY_ENV_VAR; pub use auth::CODEX_API_KEY_ENV_VAR; pub use auth::CodexAuth; pub use auth::ExternalAuth; @@ -37,9 +38,11 @@ pub use auth::UnauthorizedRecovery; pub use auth::default_client; pub use auth::enforce_login_restrictions; pub use auth::load_auth_dot_json; +pub use auth::login_with_agent_identity; pub use auth::login_with_api_key; pub use auth::logout; pub use auth::logout_with_revoke; +pub use auth::read_codex_agent_identity_from_env; pub use auth::read_openai_api_key_from_env; pub use auth::save_auth; pub use auth_env_telemetry::AuthEnvTelemetry;