From 656649cc9f018d287553a9c37ed40cc7f9fcfb03 Mon Sep 17 00:00:00 2001 From: shijie-openai Date: Sun, 26 Apr 2026 10:04:30 -0700 Subject: [PATCH] Fix agent identity runtime auth flow --- codex-rs/agent-identity/src/lib.rs | 56 ++++++++++++++++------- codex-rs/cloud-requirements/src/lib.rs | 5 ++ codex-rs/login/src/auth/agent_identity.rs | 21 ++++----- codex-rs/login/src/auth/manager.rs | 4 +- 4 files changed, 56 insertions(+), 30 deletions(-) diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index 6aad710d60..bf139f7870 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -23,6 +23,7 @@ use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; use serde::Serialize; +use serde::de::DeserializeOwned; use sha2::Digest as _; use sha2::Sha512; @@ -118,26 +119,39 @@ 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 decoding_key = if let Some(public_key_base64) = public_key_base64 { - let public_key = BASE64_STANDARD - .decode(public_key_base64) - .context("agent identity JWT public key is not valid base64")?; - DecodingKey::from_ed_der(&public_key) - } else { - validation.insecure_disable_signature_validation(); - DecodingKey::from_secret(&[]) - }; + 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, @@ -157,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")?; @@ -569,7 +591,7 @@ mod tests { 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":"EdDSA","typ":"JWT"}"#); + 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/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 8c51888a16..1d9975f128 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 5f2dc9cfc8..c11664a5bc 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,8 @@ 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 +43,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/manager.rs b/codex-rs/login/src/auth/manager.rs index b74f4e466f..85c9b6f024 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -323,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(()), } }