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) + ); + } +}