Load cloud requirements for agent identity

shijie-oai · shijie-oai · commit b9a0661d1e09 · 2026-04-28T10:47:29.000-07:00
diff --git 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
@@ -179,6 +179,14 @@ fn auth_identity(auth: &CodexAuth) -> (Option<String>, Option<String>) {
     (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<Vec<u8>> {
     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<OsString>,
+    }
+
+    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
@@ -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<String> {
+    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<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 = 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<String> {
     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,48 @@ 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 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
@@ -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
@@ -75,7 +75,7 @@ impl From<AgentIdentityJwtClaims> 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
diff --git a/codex-rs/protocol/src/auth.rs b/codex-rs/protocol/src/auth.rs

Original file line number	Diff line number	Diff line change
`@@ -75,7 +75,7 @@ impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {`
`75`	`75`	`account_id: claims.account_id,`
`76`	`76`	`chatgpt_user_id: claims.chatgpt_user_id,`
`77`	`77`	`email: claims.email,`
`78`		`- plan_type: claims.plan_type,`
	`78`	`+ plan_type: claims.plan_type.into(),`
`79`	`79`	`chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,`
`80`	`80`	`}`
`81`	`81`	`}`