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`
This commit is contained in:
Shijie Rao authored and GitHub committed 2026-04-28 12:35:00 -07:00
1 parent 0700f979ba
commit 25ac0e4527
7 files changed
+239 -48

No files matched your search

+58 -4
View File
@@ -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,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() {