From 25ac0e4527b2f9c60bd86bb8534dbc0627563d3c Mon Sep 17 00:00:00 2001 From: Shijie Rao Date: Tue, 28 Apr 2026 15:35:00 -0400 Subject: [PATCH] Load cloud requirements for agent identity (#19708) ## Why Agent Identity sessions can represent Business and Enterprise ChatGPT workspaces, but cloud requirements were skipped before fetch. That meant workspace-managed requirements were not loaded for Agent Identity even when the JWT carried the same account identity and plan information that normal ChatGPT token auth exposes. This PR now sits on top of the Agent Identity stack through [#19764](https://github.com/openai/codex/pull/19764). Because [#19763](https://github.com/openai/codex/pull/19763) moved task registration into Agent Identity auth loading, cloud requirements no longer needs a separate runtime-initialization step before building the backend client. ## What changed - Stop skipping `CodexAuth::AgentIdentity` in the cloud requirements loader. - Share the cloud requirements eligibility check between startup load and background cache refresh. - Rely on eagerly loaded Agent Identity auth so backend requests can attach task-scoped `AgentAssertion` headers. - Decode Agent Identity JWT `plan_type` as the auth-layer plan type, then convert it through a shared `auth::PlanType` -> `account::PlanType` mapping. - Add the missing serde alias for the `education` plan string and add coverage for raw Agent Identity plan aliases such as `hc` and `education`. ## Testing - `cargo test -p codex-agent-identity -p codex-login -p codex-cloud-requirements -p codex-protocol` --- codex-rs/agent-identity/src/lib.rs | 33 +++++++-- codex-rs/cloud-requirements/src/lib.rs | 99 +++++++++++++++++++++----- codex-rs/login/src/auth/auth_tests.rs | 62 ++++++++++++++-- codex-rs/login/src/auth/manager.rs | 22 +----- codex-rs/login/src/auth/storage.rs | 2 +- codex-rs/protocol/src/account.rs | 48 +++++++++++++ codex-rs/protocol/src/auth.rs | 21 ++++++ 7 files changed, 239 insertions(+), 48 deletions(-) diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index cdfb258f4..7aad81a34 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -8,7 +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::auth::PlanType as AuthPlanType; use codex_protocol::protocol::SessionSource; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; @@ -73,7 +73,7 @@ pub struct AgentIdentityJwtClaims { pub account_id: String, pub chatgpt_user_id: String, pub email: String, - pub plan_type: AccountPlanType, + pub plan_type: AuthPlanType, pub chatgpt_account_is_fedramp: bool, } @@ -408,6 +408,8 @@ mod tests { use jsonwebtoken::Header; use pretty_assertions::assert_eq; + use codex_protocol::auth::KnownPlan; + use super::*; #[test] @@ -517,12 +519,33 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, } ); } + #[test] + fn decode_agent_identity_jwt_maps_raw_plan_aliases() { + 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", + "chatgpt_user_id": "user-id", + "email": "user@example.com", + "plan_type": "hc", + "chatgpt_account_is_fedramp": false, + })); + + let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode"); + + assert_eq!(claims.plan_type, AuthPlanType::Known(KnownPlan::Enterprise)); + } + #[test] fn decode_agent_identity_jwt_verifies_when_jwks_is_present() { let jwks = test_jwks("test-key"); @@ -536,7 +559,7 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, }; let jwt = jsonwebtoken::encode( @@ -568,7 +591,7 @@ mod tests { account_id: "account-id".to_string(), chatgpt_user_id: "user-id".to_string(), email: "user@example.com".to_string(), - plan_type: AccountPlanType::Pro, + plan_type: AuthPlanType::Known(KnownPlan::Pro), chatgpt_account_is_fedramp: false, }; assert_eq!( diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index d38fbb084..c54817964 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -179,6 +179,14 @@ fn auth_identity(auth: &CodexAuth) -> (Option, Option) { (auth.get_chatgpt_user_id(), auth.get_account_id()) } +fn cloud_requirements_eligible_auth(auth: &CodexAuth) -> bool { + let Some(plan_type) = auth.account_plan_type() else { + return false; + }; + auth.uses_codex_backend() + && (plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) +} + fn cache_payload_bytes(payload: &CloudRequirementsCacheSignedPayload) -> Option> { serde_json::to_vec(&payload).ok() } @@ -329,17 +337,7 @@ 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); - }; - if !auth.uses_codex_backend() - || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) - { + if !cloud_requirements_eligible_auth(&auth) { return Ok(None); } let (chatgpt_user_id, account_id) = auth_identity(&auth); @@ -554,12 +552,7 @@ impl CloudRequirementsService { let Some(auth) = self.auth_manager.auth().await else { return false; }; - let Some(plan_type) = auth.account_plan_type() else { - return false; - }; - if !auth.uses_codex_backend() - || !(plan_type.is_business_like() || matches!(plan_type, PlanType::Enterprise)) - { + if !cloud_requirements_eligible_auth(&auth) { return false; } @@ -837,18 +830,41 @@ mod tests { use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_config::types::AuthCredentialsStoreMode; + use codex_login::auth::AgentIdentityAuth; + use codex_login::auth::AgentIdentityAuthRecord; use codex_protocol::protocol::AskForApproval; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::collections::VecDeque; + use std::ffi::OsString; use std::future::pending; + use std::io::Read; + use std::io::Write; + use std::net::TcpListener; use std::path::Path; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::thread; use tempfile::TempDir; use tempfile::tempdir; + struct EnvVarGuard { + key: &'static str, + original: Option, + } + + impl Drop for EnvVarGuard { + fn drop(&mut self) { + unsafe { + match &self.original { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } + } + } + fn write_auth_json(codex_home: &Path, value: serde_json::Value) -> std::io::Result<()> { std::fs::write(codex_home.join("auth.json"), serde_json::to_string(&value)?)?; Ok(()) @@ -1200,6 +1216,55 @@ mod tests { ); } + #[tokio::test] + async fn cloud_requirements_eligible_auth_allows_agent_identity_business_plan() { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind task registration server"); + let addr = listener + .local_addr() + .expect("task registration server addr"); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().expect("accept task registration request"); + let mut request = [0; 4096]; + let _ = stream + .read(&mut request) + .expect("read task registration request"); + let body = r#"{"task_id":"task-123"}"#; + write!( + stream, + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + .expect("write task registration response"); + }); + let record = AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-123".to_string(), + agent_private_key: "MC4CAQAwBQYDK2VwBCIEIDQg14jybCLydjHQwXeBzsDM7oB6BSAenodx6oCovQ/D" + .to_string(), + account_id: "account-12345".to_string(), + chatgpt_user_id: "user-12345".to_string(), + email: "user@example.com".to_string(), + plan_type: PlanType::Business, + chatgpt_account_is_fedramp: false, + }; + let authapi_base_url = format!("http://{addr}/backend-api"); + let original_authapi_base_url = std::env::var_os("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL"); + unsafe { + std::env::set_var("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &authapi_base_url); + } + let _authapi_guard = EnvVarGuard { + key: "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", + original: original_authapi_base_url, + }; + let auth = AgentIdentityAuth::load(record) + .await + .map(CodexAuth::AgentIdentity) + .expect("agent identity auth"); + server.join().expect("task registration server joined"); + + assert!(cloud_requirements_eligible_auth(&auth)); + } + #[tokio::test] async fn fetch_cloud_requirements_allows_business_like_usage_based_plan() { let codex_home = tempdir().expect("tempdir"); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 90748f786..80ec9d07a 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -88,7 +88,8 @@ 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 = signed_agent_identity_jwt(&record).expect("signed agent identity"); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/wham/agent-identities/jwks")) @@ -709,7 +710,8 @@ 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"); + signed_agent_identity_jwt(&expected_record, json!(expected_record.plan_type)) + .expect("signed agent identity"); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/wham/agent-identities/jwks")) @@ -925,6 +927,13 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { } fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result { + fake_agent_identity_jwt_with_plan_type(record, serde_json::to_value(record.plan_type)?) +} + +fn fake_agent_identity_jwt_with_plan_type( + record: &AgentIdentityAuthRecord, + plan_type: serde_json::Value, +) -> 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!({ @@ -937,7 +946,7 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result< "account_id": record.account_id, "chatgpt_user_id": record.chatgpt_user_id, "email": record.email, - "plan_type": record.plan_type, + "plan_type": plan_type, "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp, }); let payload_b64 = encode(&serde_json::to_vec(&payload)?); @@ -947,6 +956,7 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result< fn signed_agent_identity_jwt( record: &AgentIdentityAuthRecord, + plan_type: serde_json::Value, ) -> jsonwebtoken::errors::Result { let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); header.kid = Some("test-key".to_string()); @@ -962,7 +972,7 @@ fn signed_agent_identity_jwt( "account_id": record.account_id, "chatgpt_user_id": record.chatgpt_user_id, "email": record.email, - "plan_type": record.plan_type, + "plan_type": plan_type, "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp, }), &jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?, @@ -1011,6 +1021,50 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN 8U4M2TSWCKUY/A6sT4W8+mT9 -----END PRIVATE KEY-----"#; +#[tokio::test] +#[serial(codex_auth_env)] +async fn agent_identity_plan_type_maps_raw_enterprise_alias() { + assert_agent_identity_plan_alias(json!("hc"), AccountPlanType::Enterprise).await; +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn agent_identity_plan_type_maps_raw_education_alias() { + assert_agent_identity_plan_alias(json!("education"), AccountPlanType::Edu).await; +} + +async fn assert_agent_identity_plan_alias( + plan_type: serde_json::Value, + expected_plan_type: AccountPlanType, +) { + let record = agent_identity_record("account-id"); + let jwt = signed_agent_identity_jwt(&record, plan_type).expect("agent identity jwt"); + 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 chatgpt_base_url = format!("{}/backend-api", server.uri()); + let _authapi_guard = + EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url); + let auth = CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url)) + .await + .expect("agent identity auth"); + + pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type)); + server.verify().await; +} + #[tokio::test] #[serial(codex_auth_env)] async fn plan_type_maps_known_plan() { diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9f140c276..29897db7b 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -39,7 +39,6 @@ use crate::token_data::parse_jwt_expiration; use codex_client::CodexHttpClient; use codex_config::types::AuthCredentialsStoreMode; use codex_protocol::account::PlanType as AccountPlanType; -use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; @@ -380,29 +379,10 @@ impl CodexAuth { return Some(auth.plan_type()); } - let map_known = |kp: &InternalKnownPlan| match kp { - InternalKnownPlan::Free => AccountPlanType::Free, - InternalKnownPlan::Go => AccountPlanType::Go, - InternalKnownPlan::Plus => AccountPlanType::Plus, - InternalKnownPlan::Pro => AccountPlanType::Pro, - InternalKnownPlan::ProLite => AccountPlanType::ProLite, - InternalKnownPlan::Team => AccountPlanType::Team, - InternalKnownPlan::SelfServeBusinessUsageBased => { - AccountPlanType::SelfServeBusinessUsageBased - } - InternalKnownPlan::Business => AccountPlanType::Business, - InternalKnownPlan::EnterpriseCbpUsageBased => AccountPlanType::EnterpriseCbpUsageBased, - InternalKnownPlan::Enterprise => AccountPlanType::Enterprise, - InternalKnownPlan::Edu => AccountPlanType::Edu, - }; - self.get_current_token_data().map(|t| { t.id_token .chatgpt_plan_type - .map(|pt| match pt { - InternalPlanType::Known(k) => map_known(&k), - InternalPlanType::Unknown(_) => AccountPlanType::Unknown, - }) + .map(AccountPlanType::from) .unwrap_or(AccountPlanType::Unknown) }) } diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index b2e8946fa..3a1c8ae6a 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -75,7 +75,7 @@ impl From for AgentIdentityAuthRecord { account_id: claims.account_id, chatgpt_user_id: claims.chatgpt_user_id, email: claims.email, - plan_type: claims.plan_type, + plan_type: claims.plan_type.into(), chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, } } diff --git a/codex-rs/protocol/src/account.rs b/codex-rs/protocol/src/account.rs index aea9ad843..bfbccbeb2 100644 --- a/codex-rs/protocol/src/account.rs +++ b/codex-rs/protocol/src/account.rs @@ -3,6 +3,9 @@ use serde::Deserialize; use serde::Serialize; use ts_rs::TS; +use crate::auth::KnownPlan; +use crate::auth::PlanType as AuthPlanType; + #[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] #[serde(rename_all = "lowercase")] #[ts(rename_all = "lowercase")] @@ -57,9 +60,38 @@ impl PlanType { } } +impl From for PlanType { + fn from(plan_type: AuthPlanType) -> Self { + match plan_type { + AuthPlanType::Known(plan) => plan.into(), + AuthPlanType::Unknown(_) => Self::Unknown, + } + } +} + +impl From for PlanType { + fn from(plan: KnownPlan) -> Self { + match plan { + KnownPlan::Free => Self::Free, + KnownPlan::Go => Self::Go, + KnownPlan::Plus => Self::Plus, + KnownPlan::Pro => Self::Pro, + KnownPlan::ProLite => Self::ProLite, + KnownPlan::Team => Self::Team, + KnownPlan::SelfServeBusinessUsageBased => Self::SelfServeBusinessUsageBased, + KnownPlan::Business => Self::Business, + KnownPlan::EnterpriseCbpUsageBased => Self::EnterpriseCbpUsageBased, + KnownPlan::Enterprise => Self::Enterprise, + KnownPlan::Edu => Self::Edu, + } + } +} + #[cfg(test)] mod tests { use super::PlanType; + use crate::auth::KnownPlan; + use crate::auth::PlanType as AuthPlanType; use pretty_assertions::assert_eq; #[test] @@ -121,4 +153,20 @@ mod tests { assert_eq!(PlanType::Edu.is_workspace_account(), true); assert_eq!(PlanType::Pro.is_workspace_account(), false); } + + #[test] + fn auth_plan_type_converts_to_account_plan_type() { + assert_eq!( + PlanType::from(AuthPlanType::Known(KnownPlan::EnterpriseCbpUsageBased)), + PlanType::EnterpriseCbpUsageBased + ); + assert_eq!( + PlanType::from(AuthPlanType::Known(KnownPlan::Enterprise)), + PlanType::Enterprise + ); + assert_eq!( + PlanType::from(AuthPlanType::Unknown("mystery-tier".to_string())), + PlanType::Unknown + ); + } } diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs index 99e067bf2..495f138f4 100644 --- a/codex-rs/protocol/src/auth.rs +++ b/codex-rs/protocol/src/auth.rs @@ -46,6 +46,7 @@ pub enum KnownPlan { EnterpriseCbpUsageBased, #[serde(alias = "hc")] Enterprise, + #[serde(alias = "education")] Edu, } @@ -118,3 +119,23 @@ pub enum RefreshTokenFailedReason { Revoked, Other, } + +#[cfg(test)] +mod tests { + use super::KnownPlan; + use super::PlanType; + use pretty_assertions::assert_eq; + + #[test] + fn plan_type_deserializes_raw_aliases() { + assert_eq!( + serde_json::from_str::("\"hc\"").expect("hc should deserialize"), + PlanType::Known(KnownPlan::Enterprise) + ); + assert_eq!( + serde_json::from_str::("\"education\"") + .expect("education should deserialize"), + PlanType::Known(KnownPlan::Edu) + ); + } +}