codex_agent_identity/
lib.rs

1use std::collections::BTreeMap;
2use std::error::Error as StdError;
3use std::fmt;
4use std::time::Duration;
5
6use anyhow::Context;
7use anyhow::Result;
8use base64::Engine as _;
9use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
10use base64::engine::general_purpose::URL_SAFE_NO_PAD;
11use chrono::SecondsFormat;
12use chrono::Utc;
13use codex_http_client::HttpClient;
14use codex_http_client::HttpError;
15use codex_protocol::auth::PlanType as AuthPlanType;
16use codex_protocol::protocol::SessionSource;
17use crypto_box::SecretKey as Curve25519SecretKey;
18use ed25519_dalek::Signer as _;
19use ed25519_dalek::SigningKey;
20use ed25519_dalek::VerifyingKey;
21use ed25519_dalek::pkcs8::DecodePrivateKey;
22use ed25519_dalek::pkcs8::EncodePrivateKey;
23use http::StatusCode;
24use jsonwebtoken::Algorithm;
25use jsonwebtoken::DecodingKey;
26use jsonwebtoken::Validation;
27use jsonwebtoken::decode;
28use jsonwebtoken::decode_header;
29use jsonwebtoken::jwk::JwkSet;
30use rand::TryRngCore;
31use rand::rngs::OsRng;
32use serde::Deserialize;
33use serde::Serialize;
34use serde::de::DeserializeOwned;
35use sha2::Digest as _;
36use sha2::Sha512;
37
38const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30);
39const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10);
40const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server";
41const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity";
42const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
43const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
44const STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.api.openai.org/api/accounts";
45const AGENT_IDENTITY_KEY_SEED_BYTES: usize = 64;
46const AGENT_IDENTITY_KEY_DERIVATION_CONTEXT: &[u8] = b"codex-agent-identity-ed25519-v1";
47
48#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
49pub enum ChatGptEnvironment {
50    #[default]
51    Production,
52    Staging,
53}
54
55impl ChatGptEnvironment {
56    pub fn from_chatgpt_base_url(chatgpt_base_url: &str) -> Result<Self> {
57        match chatgpt_base_url.trim_end_matches('/') {
58            "https://chatgpt.com"
59            | "https://chatgpt.com/backend-api"
60            | "https://chatgpt.com/codex"
61            | "https://chatgpt.com/backend-api/codex"
62            | "https://chat.openai.com"
63            | "https://chat.openai.com/backend-api"
64            | "https://chat.openai.com/codex"
65            | "https://chat.openai.com/backend-api/codex" => Ok(Self::Production),
66            "https://chatgpt-staging.com"
67            | "https://chatgpt-staging.com/backend-api"
68            | "https://chatgpt-staging.com/codex"
69            | "https://chatgpt-staging.com/backend-api/codex" => Ok(Self::Staging),
70            _ => anyhow::bail!(
71                "Agent Identity only supports production and staging ChatGPT environments"
72            ),
73        }
74    }
75
76    pub fn chatgpt_base_url(self) -> &'static str {
77        match self {
78            Self::Production => "https://chatgpt.com/backend-api",
79            Self::Staging => "https://chatgpt-staging.com/backend-api",
80        }
81    }
82
83    pub fn agent_identity_authapi_base_url(self) -> &'static str {
84        match self {
85            Self::Production => PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL,
86            Self::Staging => STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL,
87        }
88    }
89}
90
91/// Borrowed durable signing material for a registered agent identity.
92///
93/// This intentionally does not include a task id. Task ids are scoped to a
94/// single Codex run, while the agent runtime id and private key are the
95/// reusable identity material used to register and sign that run task.
96#[derive(Clone, Copy, Debug, PartialEq, Eq)]
97pub struct AgentIdentityKey<'a> {
98    pub agent_runtime_id: &'a str,
99    pub private_key_pkcs8_base64: &'a str,
100}
101
102#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
103pub struct AgentBillOfMaterials {
104    pub agent_version: String,
105    pub agent_harness_id: String,
106    pub running_location: String,
107}
108
109pub struct GeneratedAgentKeyMaterial {
110    pub private_key_pkcs8_base64: String,
111    pub public_key_ssh: String,
112}
113
114/// Claims carried by an Agent Identity JWT.
115#[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
116pub struct AgentIdentityJwtClaims {
117    pub iss: String,
118    pub aud: String,
119    pub iat: usize,
120    pub exp: usize,
121    pub agent_runtime_id: String,
122    pub agent_private_key: String,
123    pub account_id: String,
124    pub chatgpt_user_id: String,
125    pub email: Option<String>,
126    pub plan_type: AuthPlanType,
127    pub chatgpt_account_is_fedramp: bool,
128}
129
130#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
131struct AgentAssertionEnvelope {
132    agent_runtime_id: String,
133    task_id: String,
134    timestamp: String,
135    signature: String,
136}
137
138#[derive(Serialize)]
139struct RegisterTaskRequest {
140    timestamp: String,
141    signature: String,
142}
143
144#[derive(Deserialize)]
145struct RegisterTaskResponse {
146    #[serde(default)]
147    task_id: Option<String>,
148    #[serde(default, rename = "taskId")]
149    task_id_camel: Option<String>,
150    #[serde(default)]
151    encrypted_task_id: Option<String>,
152    #[serde(default, rename = "encryptedTaskId")]
153    encrypted_task_id_camel: Option<String>,
154}
155
156#[derive(Debug, Serialize)]
157struct RegisterAgentRequest {
158    abom: AgentBillOfMaterials,
159    agent_public_key: String,
160    capabilities: Vec<String>,
161    ttl: Option<u64>,
162}
163
164#[derive(Debug, Deserialize)]
165struct RegisterAgentResponse {
166    agent_runtime_id: String,
167}
168
169/// HTTP status failure returned by Agent Identity registration endpoints.
170#[derive(Debug)]
171pub struct AgentIdentityRegistrationHttpError {
172    operation: &'static str,
173    status: StatusCode,
174    body: String,
175}
176
177impl AgentIdentityRegistrationHttpError {
178    fn new(operation: &'static str, status: StatusCode, body: String) -> Self {
179        Self {
180            operation,
181            status,
182            body,
183        }
184    }
185
186    /// HTTP status returned by the registration endpoint.
187    pub fn status(&self) -> StatusCode {
188        self.status
189    }
190}
191
192impl fmt::Display for AgentIdentityRegistrationHttpError {
193    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194        if self.body.is_empty() {
195            write!(f, "{} failed with status {}", self.operation, self.status)
196        } else {
197            write!(
198                f,
199                "{} failed with status {}: {}",
200                self.operation, self.status, self.body
201            )
202        }
203    }
204}
205
206impl StdError for AgentIdentityRegistrationHttpError {}
207
208/// Returns whether an Agent Identity registration error is safe to retry.
209pub fn is_retryable_registration_error(error: &anyhow::Error) -> bool {
210    error.chain().any(is_retryable_registration_cause)
211}
212
213fn is_retryable_registration_cause(cause: &(dyn StdError + 'static)) -> bool {
214    if let Some(error) = cause.downcast_ref::<AgentIdentityRegistrationHttpError>() {
215        return is_retryable_registration_status(error.status());
216    }
217
218    if let Some(error) = cause.downcast_ref::<HttpError>() {
219        if let Some(status) = error.status() {
220            return is_retryable_registration_status(status);
221        }
222        return error.is_timeout() || error.is_connect() || error.is_request();
223    }
224
225    false
226}
227
228fn is_retryable_registration_status(status: StatusCode) -> bool {
229    status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error()
230}
231
232pub fn authorization_header_for_agent_task(
233    key: AgentIdentityKey<'_>,
234    task_id: &str,
235) -> Result<String> {
236    let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
237    let envelope = AgentAssertionEnvelope {
238        agent_runtime_id: key.agent_runtime_id.to_string(),
239        task_id: task_id.to_string(),
240        timestamp: timestamp.clone(),
241        signature: sign_agent_assertion_payload(key, task_id, &timestamp)?,
242    };
243    let serialized_assertion = serialize_agent_assertion(&envelope)?;
244    Ok(format!("AgentAssertion {serialized_assertion}"))
245}
246
247pub async fn fetch_agent_identity_jwks(
248    client: &HttpClient,
249    agent_identity_jwt_base_url: &str,
250) -> Result<JwkSet> {
251    let response = client
252        .get(agent_identity_jwks_url(agent_identity_jwt_base_url))
253        .timeout(AGENT_IDENTITY_JWKS_TIMEOUT)
254        .send()
255        .await
256        .context("failed to request agent identity JWKS")?
257        .error_for_status()
258        .context("agent identity JWKS endpoint returned an error")?;
259
260    response
261        .json()
262        .await
263        .context("failed to decode agent identity JWKS")
264}
265
266pub fn decode_agent_identity_jwt(
267    jwt: &str,
268    jwks: Option<&JwkSet>,
269) -> Result<AgentIdentityJwtClaims> {
270    let Some(jwks) = jwks else {
271        return decode_agent_identity_jwt_payload(jwt);
272    };
273
274    let header = decode_header(jwt).context("failed to decode agent identity JWT header")?;
275    let kid = header
276        .kid
277        .context("agent identity JWT header does not include a kid")?;
278    let jwk = jwks
279        .find(&kid)
280        .with_context(|| format!("agent identity JWT kid {kid} is not trusted"))?;
281    let decoding_key = DecodingKey::from_jwk(jwk).context("failed to build JWT decoding key")?;
282    let mut validation = Validation::new(Algorithm::RS256);
283    validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]);
284    validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]);
285    validation.required_spec_claims.insert("iss".to_string());
286    validation.required_spec_claims.insert("aud".to_string());
287    decode::<AgentIdentityJwtClaims>(jwt, &decoding_key, &validation)
288        .map(|data| data.claims)
289        .context("failed to verify agent identity JWT")
290}
291
292fn decode_agent_identity_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T> {
293    let mut parts = jwt.split('.');
294    let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) {
295        (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s),
296        _ => anyhow::bail!("invalid agent identity JWT format"),
297    };
298    anyhow::ensure!(parts.next().is_none(), "invalid agent identity JWT format");
299
300    let payload_bytes = URL_SAFE_NO_PAD
301        .decode(payload_b64)
302        .context("agent identity JWT payload is not valid base64url")?;
303    serde_json::from_slice(&payload_bytes).context("agent identity JWT payload is not valid JSON")
304}
305
306pub fn sign_task_registration_payload(
307    key: AgentIdentityKey<'_>,
308    timestamp: &str,
309) -> Result<String> {
310    let signing_key = signing_key_from_private_key_pkcs8_base64(key.private_key_pkcs8_base64)?;
311    let payload = format!("{}:{timestamp}", key.agent_runtime_id);
312    Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes()))
313}
314
315pub async fn register_agent_task(
316    client: &HttpClient,
317    agent_identity_authapi_base_url: &str,
318    key: AgentIdentityKey<'_>,
319) -> Result<String> {
320    let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
321    let request = RegisterTaskRequest {
322        signature: sign_task_registration_payload(key, &timestamp)?,
323        timestamp,
324    };
325    let url = agent_task_registration_url(agent_identity_authapi_base_url, key.agent_runtime_id);
326
327    let response = client
328        .post(url)
329        .timeout(AGENT_TASK_REGISTRATION_TIMEOUT)
330        .json(&request)
331        .send()
332        .await
333        .context("failed to register agent task")?;
334    if !response.status().is_success() {
335        let status = response.status();
336        let body = response.text().await.unwrap_or_default();
337        let body = if body.len() > 512 {
338            format!("{}...", body.chars().take(512).collect::<String>())
339        } else {
340            body
341        };
342        return Err(AgentIdentityRegistrationHttpError::new(
343            "agent task registration",
344            status,
345            body,
346        )
347        .into());
348    }
349
350    let response = response
351        .json()
352        .await
353        .context("failed to decode agent task registration response")?;
354
355    task_id_from_register_task_response(key, response)
356}
357
358pub async fn register_agent_identity(
359    client: &HttpClient,
360    agent_identity_authapi_base_url: &str,
361    access_token: &str,
362    is_fedramp_account: bool,
363    key_material: &GeneratedAgentKeyMaterial,
364    abom: AgentBillOfMaterials,
365    capabilities: Vec<String>,
366) -> Result<String> {
367    let url = agent_registration_url(agent_identity_authapi_base_url);
368    let request = RegisterAgentRequest {
369        abom,
370        agent_public_key: key_material.public_key_ssh.clone(),
371        capabilities,
372        ttl: None,
373    };
374
375    let mut request_builder = client
376        .post(&url)
377        .bearer_auth(access_token)
378        .json(&request)
379        .timeout(AGENT_REGISTRATION_TIMEOUT);
380    if is_fedramp_account {
381        request_builder = request_builder.header("X-OpenAI-Fedramp", "true");
382    }
383
384    let response = request_builder
385        .send()
386        .await
387        .with_context(|| format!("failed to send agent identity registration request to {url}"))?
388        .error_for_status()
389        .with_context(|| format!("agent identity registration failed for {url}"))?
390        .json::<RegisterAgentResponse>()
391        .await
392        .with_context(|| format!("failed to parse agent identity response from {url}"))?;
393
394    Ok(response.agent_runtime_id)
395}
396
397fn task_id_from_register_task_response(
398    key: AgentIdentityKey<'_>,
399    response: RegisterTaskResponse,
400) -> Result<String> {
401    if let Some(task_id) = response.task_id.or(response.task_id_camel) {
402        return Ok(task_id);
403    }
404    let encrypted_task_id = response
405        .encrypted_task_id
406        .or(response.encrypted_task_id_camel)
407        .context("agent task registration response omitted task id")?;
408    decrypt_task_id_response(key, &encrypted_task_id)
409}
410
411pub fn decrypt_task_id_response(
412    key: AgentIdentityKey<'_>,
413    encrypted_task_id: &str,
414) -> Result<String> {
415    let signing_key = signing_key_from_private_key_pkcs8_base64(key.private_key_pkcs8_base64)?;
416    let ciphertext = BASE64_STANDARD
417        .decode(encrypted_task_id)
418        .context("encrypted task id is not valid base64")?;
419    let plaintext = curve25519_secret_key_from_signing_key(&signing_key)
420        .unseal(&ciphertext)
421        .map_err(|_| anyhow::anyhow!("failed to decrypt encrypted task id"))?;
422    String::from_utf8(plaintext).context("decrypted task id is not valid UTF-8")
423}
424
425pub fn generate_agent_key_material() -> Result<GeneratedAgentKeyMaterial> {
426    let mut seed_material = [0u8; AGENT_IDENTITY_KEY_SEED_BYTES];
427    OsRng
428        .try_fill_bytes(&mut seed_material)
429        .context("failed to generate agent identity private key seed material")?;
430    // Ed25519 stores a 32-byte seed, so derive it from all sampled seed material.
431    let mut digest = Sha512::new();
432    digest.update(AGENT_IDENTITY_KEY_DERIVATION_CONTEXT);
433    digest.update(seed_material);
434    let digest = digest.finalize();
435    let mut secret_key_bytes = [0u8; 32];
436    secret_key_bytes.copy_from_slice(&digest[..32]);
437    let signing_key = SigningKey::from_bytes(&secret_key_bytes);
438    let private_key_pkcs8 = signing_key
439        .to_pkcs8_der()
440        .context("failed to encode agent identity private key as PKCS#8")?;
441
442    Ok(GeneratedAgentKeyMaterial {
443        private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()),
444        public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()),
445    })
446}
447
448pub fn public_key_ssh_from_private_key_pkcs8_base64(
449    private_key_pkcs8_base64: &str,
450) -> Result<String> {
451    let signing_key = signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64)?;
452    Ok(encode_ssh_ed25519_public_key(&signing_key.verifying_key()))
453}
454
455pub fn verifying_key_from_private_key_pkcs8_base64(
456    private_key_pkcs8_base64: &str,
457) -> Result<VerifyingKey> {
458    let signing_key = signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64)?;
459    Ok(signing_key.verifying_key())
460}
461
462pub fn curve25519_secret_key_from_private_key_pkcs8_base64(
463    private_key_pkcs8_base64: &str,
464) -> Result<Curve25519SecretKey> {
465    let signing_key = signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64)?;
466    Ok(curve25519_secret_key_from_signing_key(&signing_key))
467}
468
469pub fn agent_registration_url(agent_identity_authapi_base_url: &str) -> String {
470    agent_identity_authapi_url(agent_identity_authapi_base_url, "/v1/agent/register")
471}
472
473pub fn agent_task_registration_url(
474    agent_identity_authapi_base_url: &str,
475    agent_runtime_id: &str,
476) -> String {
477    agent_identity_authapi_url(
478        agent_identity_authapi_base_url,
479        &format!("/v1/agent/{agent_runtime_id}/task/register"),
480    )
481}
482
483pub fn agent_identity_jwks_url(agent_identity_jwt_base_url: &str) -> String {
484    let trimmed = agent_identity_jwt_base_url.trim_end_matches('/');
485    if trimmed.contains("/backend-api") {
486        format!("{trimmed}/wham/agent-identities/jwks")
487    } else {
488        format!("{trimmed}/agent-identities/jwks")
489    }
490}
491
492fn agent_identity_authapi_url(agent_identity_authapi_base_url: &str, api_path: &str) -> String {
493    let base_url = agent_identity_authapi_base_url.trim_end_matches('/');
494    format!("{base_url}{api_path}")
495}
496
497pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials {
498    AgentBillOfMaterials {
499        agent_version: env!("CARGO_PKG_VERSION").to_string(),
500        agent_harness_id: match &session_source {
501            SessionSource::VSCode => "codex-app".to_string(),
502            SessionSource::Cli
503            | SessionSource::Exec
504            | SessionSource::Mcp
505            | SessionSource::Custom(_)
506            | SessionSource::Internal(_)
507            | SessionSource::SubAgent(_)
508            | SessionSource::Unknown => "codex-cli".to_string(),
509        },
510        running_location: format!("{}-{}", session_source, std::env::consts::OS),
511    }
512}
513
514pub fn encode_ssh_ed25519_public_key(verifying_key: &VerifyingKey) -> String {
515    let mut blob = Vec::with_capacity(4 + 11 + 4 + 32);
516    append_ssh_string(&mut blob, b"ssh-ed25519");
517    append_ssh_string(&mut blob, verifying_key.as_bytes());
518    format!("ssh-ed25519 {}", BASE64_STANDARD.encode(blob))
519}
520
521fn sign_agent_assertion_payload(
522    key: AgentIdentityKey<'_>,
523    task_id: &str,
524    timestamp: &str,
525) -> Result<String> {
526    let signing_key = signing_key_from_private_key_pkcs8_base64(key.private_key_pkcs8_base64)?;
527    let payload = format!("{}:{task_id}:{timestamp}", key.agent_runtime_id);
528    Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes()))
529}
530
531fn serialize_agent_assertion(envelope: &AgentAssertionEnvelope) -> Result<String> {
532    let payload = serde_json::to_vec(&BTreeMap::from([
533        ("agent_runtime_id", envelope.agent_runtime_id.as_str()),
534        ("signature", envelope.signature.as_str()),
535        ("task_id", envelope.task_id.as_str()),
536        ("timestamp", envelope.timestamp.as_str()),
537    ]))
538    .context("failed to serialize agent assertion envelope")?;
539    Ok(URL_SAFE_NO_PAD.encode(payload))
540}
541
542fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey {
543    let digest = Sha512::digest(signing_key.to_bytes());
544    let mut secret_key = [0u8; 32];
545    secret_key.copy_from_slice(&digest[..32]);
546    secret_key[0] &= 248;
547    secret_key[31] &= 127;
548    secret_key[31] |= 64;
549    Curve25519SecretKey::from(secret_key)
550}
551
552fn append_ssh_string(buf: &mut Vec<u8>, value: &[u8]) {
553    buf.extend_from_slice(&(value.len() as u32).to_be_bytes());
554    buf.extend_from_slice(value);
555}
556
557fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result<SigningKey> {
558    let private_key = BASE64_STANDARD
559        .decode(private_key_pkcs8_base64)
560        .context("stored agent identity private key is not valid base64")?;
561    SigningKey::from_pkcs8_der(&private_key)
562        .context("stored agent identity private key is not valid PKCS#8")
563}
564
565#[cfg(test)]
566mod tests {
567    use base64::Engine as _;
568    use ed25519_dalek::Signature;
569    use ed25519_dalek::Verifier as _;
570    use jsonwebtoken::EncodingKey;
571    use jsonwebtoken::Header;
572    use pretty_assertions::assert_eq;
573
574    use codex_protocol::auth::KnownPlan;
575
576    use super::*;
577
578    #[test]
579    fn register_task_request_uses_single_run_task_shape() {
580        let request = RegisterTaskRequest {
581            timestamp: "2026-04-23T00:00:00Z".to_string(),
582            signature: "signature".to_string(),
583        };
584
585        let serialized = serde_json::to_value(request).expect("serialize request");
586
587        assert_eq!(
588            serialized,
589            serde_json::json!({
590                "timestamp": "2026-04-23T00:00:00Z",
591                "signature": "signature",
592            })
593        );
594    }
595
596    #[test]
597    fn authorization_header_for_agent_task_serializes_signed_agent_assertion() {
598        let signing_key = SigningKey::from_bytes(&[7u8; 32]);
599        let private_key = signing_key
600            .to_pkcs8_der()
601            .expect("encode test key material");
602        let key = AgentIdentityKey {
603            agent_runtime_id: "agent-123",
604            private_key_pkcs8_base64: &BASE64_STANDARD.encode(private_key.as_bytes()),
605        };
606
607        let header = authorization_header_for_agent_task(key, "task-123")
608            .expect("build agent assertion header");
609        let token = header
610            .strip_prefix("AgentAssertion ")
611            .expect("agent assertion scheme");
612        let payload = URL_SAFE_NO_PAD
613            .decode(token)
614            .expect("valid base64url payload");
615        let envelope: AgentAssertionEnvelope =
616            serde_json::from_slice(&payload).expect("valid assertion envelope");
617
618        assert_eq!(
619            envelope,
620            AgentAssertionEnvelope {
621                agent_runtime_id: "agent-123".to_string(),
622                task_id: "task-123".to_string(),
623                timestamp: envelope.timestamp.clone(),
624                signature: envelope.signature.clone(),
625            }
626        );
627        let signature_bytes = BASE64_STANDARD
628            .decode(&envelope.signature)
629            .expect("valid base64 signature");
630        let signature = Signature::from_slice(&signature_bytes).expect("valid signature bytes");
631        signing_key
632            .verifying_key()
633            .verify(
634                format!(
635                    "{}:{}:{}",
636                    envelope.agent_runtime_id, envelope.task_id, envelope.timestamp
637                )
638                .as_bytes(),
639                &signature,
640            )
641            .expect("signature should verify");
642    }
643
644    #[test]
645    fn decode_agent_identity_jwt_reads_claims() {
646        let jwt = jwt_with_payload(serde_json::json!({
647            "iss": AGENT_IDENTITY_JWT_ISSUER,
648            "aud": AGENT_IDENTITY_JWT_AUDIENCE,
649            "iat": 1_700_000_000usize,
650            "exp": 4_000_000_000usize,
651            "agent_runtime_id": "agent-runtime-id",
652            "agent_private_key": "private-key",
653            "account_id": "account-id",
654            "chatgpt_user_id": "user-id",
655            "email": "user@example.com",
656            "plan_type": "pro",
657            "chatgpt_account_is_fedramp": false,
658        }));
659
660        let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode");
661
662        assert_eq!(
663            claims,
664            AgentIdentityJwtClaims {
665                iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
666                aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
667                iat: 1_700_000_000,
668                exp: 4_000_000_000,
669                agent_runtime_id: "agent-runtime-id".to_string(),
670                agent_private_key: "private-key".to_string(),
671                account_id: "account-id".to_string(),
672                chatgpt_user_id: "user-id".to_string(),
673                email: Some("user@example.com".to_string()),
674                plan_type: AuthPlanType::Known(KnownPlan::Pro),
675                chatgpt_account_is_fedramp: false,
676            }
677        );
678    }
679
680    #[test]
681    fn decode_agent_identity_jwt_accepts_missing_email() {
682        let jwt = jwt_with_payload(serde_json::json!({
683            "iss": AGENT_IDENTITY_JWT_ISSUER,
684            "aud": AGENT_IDENTITY_JWT_AUDIENCE,
685            "iat": 1_700_000_000usize,
686            "exp": 4_000_000_000usize,
687            "agent_runtime_id": "agent-runtime-id",
688            "agent_private_key": "private-key",
689            "account_id": "account-id",
690            "chatgpt_user_id": "user-id",
691            "plan_type": "pro",
692            "chatgpt_account_is_fedramp": false,
693        }));
694
695        let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode");
696
697        assert_eq!(claims.email, None);
698    }
699
700    #[test]
701    fn decode_agent_identity_jwt_maps_raw_plan_aliases() {
702        let jwt = jwt_with_payload(serde_json::json!({
703            "iss": AGENT_IDENTITY_JWT_ISSUER,
704            "aud": AGENT_IDENTITY_JWT_AUDIENCE,
705            "iat": 1_700_000_000usize,
706            "exp": 4_000_000_000usize,
707            "agent_runtime_id": "agent-runtime-id",
708            "agent_private_key": "private-key",
709            "account_id": "account-id",
710            "chatgpt_user_id": "user-id",
711            "email": "user@example.com",
712            "plan_type": "hc",
713            "chatgpt_account_is_fedramp": false,
714        }));
715
716        let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode");
717
718        assert_eq!(claims.plan_type, AuthPlanType::Known(KnownPlan::Enterprise));
719    }
720
721    #[test]
722    fn decode_agent_identity_jwt_verifies_when_jwks_is_present() {
723        let jwks = test_jwks("test-key");
724        let claims = AgentIdentityJwtClaims {
725            iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
726            aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
727            iat: 1_700_000_000,
728            exp: 4_000_000_000,
729            agent_runtime_id: "agent-runtime-id".to_string(),
730            agent_private_key: "private-key".to_string(),
731            account_id: "account-id".to_string(),
732            chatgpt_user_id: "user-id".to_string(),
733            email: Some("user@example.com".to_string()),
734            plan_type: AuthPlanType::Known(KnownPlan::Pro),
735            chatgpt_account_is_fedramp: false,
736        };
737        let jwt = jsonwebtoken::encode(
738            &test_jwt_header("test-key"),
739            &serde_json::json!({
740                "iss": claims.iss,
741                "aud": claims.aud,
742                "iat": claims.iat,
743                "exp": claims.exp,
744                "agent_runtime_id": claims.agent_runtime_id,
745                "agent_private_key": claims.agent_private_key,
746                "account_id": claims.account_id,
747                "chatgpt_user_id": claims.chatgpt_user_id,
748                "email": claims.email,
749                "plan_type": "pro",
750                "chatgpt_account_is_fedramp": claims.chatgpt_account_is_fedramp,
751            }),
752            &test_rsa_encoding_key(),
753        )
754        .expect("JWT should encode");
755
756        let expected_claims = AgentIdentityJwtClaims {
757            iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
758            aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
759            iat: 1_700_000_000,
760            exp: 4_000_000_000,
761            agent_runtime_id: "agent-runtime-id".to_string(),
762            agent_private_key: "private-key".to_string(),
763            account_id: "account-id".to_string(),
764            chatgpt_user_id: "user-id".to_string(),
765            email: Some("user@example.com".to_string()),
766            plan_type: AuthPlanType::Known(KnownPlan::Pro),
767            chatgpt_account_is_fedramp: false,
768        };
769        assert_eq!(
770            decode_agent_identity_jwt(&jwt, Some(&jwks)).expect("JWT should verify"),
771            expected_claims
772        );
773    }
774
775    #[test]
776    fn decode_agent_identity_jwt_rejects_untrusted_kid() {
777        let jwks = test_jwks("other-key");
778
779        let jwt = jsonwebtoken::encode(
780            &test_jwt_header("test-key"),
781            &serde_json::json!({
782                "iss": AGENT_IDENTITY_JWT_ISSUER,
783                "aud": AGENT_IDENTITY_JWT_AUDIENCE,
784                "iat": 1_700_000_000,
785                "exp": 4_000_000_000usize,
786                "agent_runtime_id": "agent-runtime-id",
787                "agent_private_key": "private-key",
788                "account_id": "account-id",
789                "chatgpt_user_id": "user-id",
790                "email": "user@example.com",
791                "plan_type": "pro",
792                "chatgpt_account_is_fedramp": false,
793            }),
794            &test_rsa_encoding_key(),
795        )
796        .expect("JWT should encode");
797
798        decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
799    }
800
801    #[test]
802    fn decode_agent_identity_jwt_requires_issuer_and_audience() {
803        let jwks = test_jwks("test-key");
804        let jwt = jsonwebtoken::encode(
805            &test_jwt_header("test-key"),
806            &serde_json::json!({
807                "iat": 1_700_000_000,
808                "exp": 4_000_000_000usize,
809                "agent_runtime_id": "agent-runtime-id",
810                "agent_private_key": "private-key",
811                "account_id": "account-id",
812                "chatgpt_user_id": "user-id",
813                "email": "user@example.com",
814                "plan_type": "pro",
815                "chatgpt_account_is_fedramp": false,
816            }),
817            &test_rsa_encoding_key(),
818        )
819        .expect("JWT should encode");
820
821        decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
822    }
823
824    fn test_jwt_header(kid: &str) -> Header {
825        let mut header = Header::new(Algorithm::RS256);
826        header.kid = Some(kid.to_string());
827        header
828    }
829
830    fn test_rsa_encoding_key() -> EncodingKey {
831        EncodingKey::from_rsa_pem(
832            br#"-----BEGIN PRIVATE KEY-----
833MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
834bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
835FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
836mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
8370eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
838zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
839uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
8401ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
841RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
842NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
843rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
844Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
845aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
846iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
847YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
8483fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
849sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
850gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
851sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
852BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
853KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
854r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
85529HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
856J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
8575da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
8588U4M2TSWCKUY/A6sT4W8+mT9
859-----END PRIVATE KEY-----"#,
860        )
861        .expect("test RSA key should parse")
862    }
863
864    fn test_jwks(kid: &str) -> jsonwebtoken::jwk::JwkSet {
865        serde_json::from_value(serde_json::json!({
866            "keys": [{
867                "kty": "RSA",
868                "kid": kid,
869                "use": "sig",
870                "alg": "RS256",
871                "n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
872                "e": "AQAB",
873            }]
874        }))
875        .expect("test JWKS should parse")
876    }
877
878    #[test]
879    fn chatgpt_environment_maps_known_urls_to_authapi() -> anyhow::Result<()> {
880        assert_eq!(
881            ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt.com/backend-api/codex")?,
882            ChatGptEnvironment::Production
883        );
884        assert_eq!(
885            ChatGptEnvironment::Production.agent_identity_authapi_base_url(),
886            "https://auth.openai.com/api/accounts"
887        );
888        assert_eq!(
889            ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt-staging.com/backend-api")?,
890            ChatGptEnvironment::Staging
891        );
892        assert_eq!(
893            ChatGptEnvironment::Staging.agent_identity_authapi_base_url(),
894            "https://auth.api.openai.org/api/accounts"
895        );
896        Ok(())
897    }
898
899    #[test]
900    fn chatgpt_environment_rejects_custom_urls() {
901        assert!(ChatGptEnvironment::from_chatgpt_base_url("http://localhost:8080").is_err(),);
902    }
903
904    #[test]
905    fn agent_registration_url_appends_to_authapi_base_url() {
906        assert_eq!(
907            agent_registration_url("https://auth.openai.com/api/accounts"),
908            "https://auth.openai.com/api/accounts/v1/agent/register"
909        );
910        assert_eq!(
911            agent_registration_url("http://localhost:8080"),
912            "http://localhost:8080/v1/agent/register"
913        );
914        assert_eq!(
915            agent_registration_url("http://localhost:8080/backend-api"),
916            "http://localhost:8080/backend-api/v1/agent/register"
917        );
918    }
919
920    #[test]
921    fn agent_task_registration_url_appends_to_authapi_base_url() {
922        assert_eq!(
923            agent_task_registration_url("https://auth.openai.com/api/accounts", "agent-runtime-id"),
924            "https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
925        );
926        assert_eq!(
927            agent_task_registration_url(
928                "https://auth.openai.com/api/accounts/",
929                "agent-runtime-id"
930            ),
931            "https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
932        );
933        assert_eq!(
934            agent_task_registration_url("http://localhost:8080", "agent-runtime-id"),
935            "http://localhost:8080/v1/agent/agent-runtime-id/task/register"
936        );
937    }
938
939    #[test]
940    fn retryable_registration_error_accepts_429_and_5xx() {
941        let too_many_requests = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
942            "agent registration",
943            StatusCode::TOO_MANY_REQUESTS,
944            "rate limited".to_string(),
945        ));
946        let unavailable = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
947            "agent registration",
948            StatusCode::SERVICE_UNAVAILABLE,
949            "try later".to_string(),
950        ));
951
952        assert!(is_retryable_registration_error(&too_many_requests));
953        assert!(is_retryable_registration_error(&unavailable));
954    }
955
956    #[test]
957    fn retryable_registration_error_rejects_hard_failures() {
958        let forbidden = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
959            "agent registration",
960            StatusCode::FORBIDDEN,
961            "not allowed".to_string(),
962        ));
963        let malformed = anyhow::anyhow!("failed to sign registration request");
964
965        assert!(!is_retryable_registration_error(&forbidden));
966        assert!(!is_retryable_registration_error(&malformed));
967    }
968
969    #[test]
970    fn agent_identity_jwks_url_uses_agent_identity_jwt_route() {
971        assert_eq!(
972            agent_identity_jwks_url("https://chatgpt.com/backend-api"),
973            "https://chatgpt.com/backend-api/wham/agent-identities/jwks"
974        );
975        assert_eq!(
976            agent_identity_jwks_url("https://chatgpt.com/backend-api/"),
977            "https://chatgpt.com/backend-api/wham/agent-identities/jwks"
978        );
979    }
980
981    #[test]
982    fn agent_identity_jwks_url_uses_jwt_issuer_base_url() {
983        assert_eq!(
984            agent_identity_jwks_url("http://localhost:8080/api/codex"),
985            "http://localhost:8080/api/codex/agent-identities/jwks"
986        );
987        assert_eq!(
988            agent_identity_jwks_url("http://localhost:8080/api/codex/"),
989            "http://localhost:8080/api/codex/agent-identities/jwks"
990        );
991    }
992
993    fn jwt_with_payload(payload: serde_json::Value) -> String {
994        let encode = |bytes: &[u8]| URL_SAFE_NO_PAD.encode(bytes);
995        let header_b64 = encode(br#"{"alg":"none","typ":"JWT"}"#);
996        let payload_b64 = encode(&serde_json::to_vec(&payload).expect("payload should serialize"));
997        let signature_b64 = encode(b"sig");
998        format!("{header_b64}.{payload_b64}.{signature_b64}")
999    }
1000}
codex_agent_identity/lib.rs

codex_agent_identity/
lib.rs
