codex_login/auth/
agent_identity.rs

1use std::future::Future;
2use std::sync::Arc;
3
4use codex_agent_identity::AgentIdentityKey;
5use codex_agent_identity::ChatGptEnvironment;
6use codex_agent_identity::agent_identity_jwks_url;
7use codex_agent_identity::agent_registration_url;
8use codex_agent_identity::agent_task_registration_url;
9use codex_agent_identity::build_abom;
10use codex_agent_identity::decode_agent_identity_jwt;
11use codex_agent_identity::fetch_agent_identity_jwks;
12use codex_agent_identity::generate_agent_key_material;
13use codex_agent_identity::is_retryable_registration_error;
14use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64;
15use codex_agent_identity::register_agent_identity;
16use codex_agent_identity::register_agent_task;
17use codex_http_client::HttpClient;
18use codex_protocol::account::PlanType as AccountPlanType;
19use codex_protocol::protocol::SessionSource;
20use thiserror::Error;
21
22use crate::default_client::create_default_auth_client;
23use crate::outbound_proxy::AuthRouteConfig;
24
25use super::storage::AgentIdentityAuthRecord;
26
27pub(super) const MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS: usize = 3;
28
29pub(super) fn agent_identity_authapi_base_url(
30    chatgpt_base_url: Option<&str>,
31) -> std::io::Result<String> {
32    let environment = match chatgpt_base_url {
33        Some(chatgpt_base_url) => ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url)
34            .map_err(std::io::Error::other)?,
35        None => ChatGptEnvironment::default(),
36    };
37    Ok(environment.agent_identity_authapi_base_url().to_string())
38}
39
40pub(super) fn require_agent_identity_authapi_base_url(
41    agent_identity_authapi_base_url: Option<&str>,
42) -> std::io::Result<&str> {
43    agent_identity_authapi_base_url.ok_or_else(|| {
44        std::io::Error::other(
45            "Agent Identity only supports production and staging ChatGPT environments",
46        )
47    })
48}
49
50#[derive(Clone, Debug, Error)]
51pub enum AgentIdentityAuthError {
52    #[error(
53        "agent identity bootstrap unavailable after {attempts} attempts during {operation}: {message}"
54    )]
55    BootstrapUnavailable {
56        operation: &'static str,
57        attempts: usize,
58        message: String,
59    },
60}
61
62impl AgentIdentityAuthError {
63    pub(super) fn bootstrap_unavailable(error: &std::io::Error) -> Option<&Self> {
64        match error
65            .get_ref()
66            .and_then(|source| source.downcast_ref::<Self>())
67        {
68            Some(error @ Self::BootstrapUnavailable { .. }) => Some(error),
69            None => None,
70        }
71    }
72}
73
74#[derive(Debug, Error)]
75#[error("retryable agent identity registration failure: {message}")]
76pub(super) struct RetryableAgentIdentityRegistrationError {
77    message: String,
78}
79
80impl RetryableAgentIdentityRegistrationError {
81    pub(super) fn new(message: String) -> Self {
82        Self { message }
83    }
84}
85
86#[derive(Clone, Debug)]
87pub struct AgentIdentityAuth {
88    record: Arc<AgentIdentityAuthRecord>,
89}
90
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub(super) struct ManagedChatGptAgentIdentityBinding {
93    pub(super) account_id: String,
94    pub(super) chatgpt_user_id: String,
95    pub(super) email: Option<String>,
96    pub(super) plan_type: AccountPlanType,
97    pub(super) chatgpt_account_is_fedramp: bool,
98    pub(super) access_token: String,
99}
100
101impl AgentIdentityAuth {
102    pub async fn from_record(
103        mut record: AgentIdentityAuthRecord,
104        agent_identity_authapi_base_url: &str,
105        auth_route_config: Option<&AuthRouteConfig>,
106    ) -> std::io::Result<Self> {
107        public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key)
108            .map_err(std::io::Error::other)?;
109        if record_needs_task_registration(&record) {
110            record.task_id = Some(
111                register_task_for_record_with_retries(
112                    &record,
113                    agent_identity_authapi_base_url,
114                    auth_route_config,
115                )
116                .await?,
117            );
118        }
119        Ok(Self {
120            record: Arc::new(record),
121        })
122    }
123
124    pub async fn from_jwt(
125        jwt: &str,
126        chatgpt_base_url: &str,
127        agent_identity_authapi_base_url: &str,
128        auth_route_config: Option<&AuthRouteConfig>,
129    ) -> std::io::Result<Self> {
130        let record = verified_record_from_jwt(jwt, chatgpt_base_url, auth_route_config).await?;
131        Self::from_record(record, agent_identity_authapi_base_url, auth_route_config).await
132    }
133
134    #[cfg(test)]
135    fn from_initialized_record(mut record: AgentIdentityAuthRecord, run_task_id: String) -> Self {
136        record.task_id = Some(run_task_id);
137        Self {
138            record: Arc::new(record),
139        }
140    }
141
142    pub fn record(&self) -> &AgentIdentityAuthRecord {
143        self.record.as_ref()
144    }
145
146    pub fn run_task_id(&self) -> &str {
147        match self.record.task_id.as_deref() {
148            Some(task_id) => task_id,
149            None => unreachable!("AgentIdentityAuth should only be constructed with a task_id"),
150        }
151    }
152
153    pub fn account_id(&self) -> &str {
154        &self.record.account_id
155    }
156
157    pub fn chatgpt_user_id(&self) -> &str {
158        &self.record.chatgpt_user_id
159    }
160
161    pub fn email(&self) -> Option<&str> {
162        self.record.email.as_deref()
163    }
164
165    pub fn plan_type(&self) -> AccountPlanType {
166        self.record.plan_type
167    }
168
169    pub fn is_fedramp_account(&self) -> bool {
170        self.record.chatgpt_account_is_fedramp
171    }
172}
173
174pub(super) async fn register_managed_chatgpt_agent_identity(
175    binding: ManagedChatGptAgentIdentityBinding,
176    agent_identity_authapi_base_url: &str,
177    session_source: SessionSource,
178    auth_route_config: Option<&AuthRouteConfig>,
179) -> std::io::Result<AgentIdentityAuth> {
180    let key_material = generate_agent_key_material().map_err(std::io::Error::other)?;
181    let registration_url = agent_registration_url(agent_identity_authapi_base_url);
182    let client = create_default_auth_client(&registration_url, auth_route_config)?;
183    let runtime_id = retry_registration(|| async {
184        register_agent_identity(
185            &client,
186            agent_identity_authapi_base_url,
187            &binding.access_token,
188            binding.chatgpt_account_is_fedramp,
189            &key_material,
190            build_abom(session_source.clone()),
191            vec!["responsesapi".to_string()],
192        )
193        .await
194        .map_err(|err| {
195            if is_retryable_registration_error(&err) {
196                std::io::Error::other(RetryableAgentIdentityRegistrationError::new(
197                    err.to_string(),
198                ))
199            } else {
200                std::io::Error::other(err)
201            }
202        })
203    })
204    .await
205    .map_err(|err| classify_bootstrap_error("agent identity registration", err))?;
206
207    let record = AgentIdentityAuthRecord {
208        agent_runtime_id: runtime_id,
209        agent_private_key: key_material.private_key_pkcs8_base64,
210        account_id: binding.account_id,
211        chatgpt_user_id: binding.chatgpt_user_id,
212        email: binding.email,
213        plan_type: binding.plan_type,
214        chatgpt_account_is_fedramp: binding.chatgpt_account_is_fedramp,
215        task_id: None,
216    };
217    AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url, auth_route_config)
218        .await
219        .map_err(|err| classify_bootstrap_error("agent task registration", err))
220}
221
222pub(super) async fn verified_record_from_jwt(
223    jwt: &str,
224    chatgpt_base_url: &str,
225    auth_route_config: Option<&AuthRouteConfig>,
226) -> std::io::Result<AgentIdentityAuthRecord> {
227    AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
228    let jwks_url = agent_identity_jwks_url(chatgpt_base_url);
229    let client = create_default_auth_client(&jwks_url, auth_route_config)?;
230    let jwks = fetch_agent_identity_jwks(&client, chatgpt_base_url)
231        .await
232        .map_err(std::io::Error::other)?;
233    let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?;
234    Ok(claims.into())
235}
236
237pub(super) fn record_needs_task_registration(record: &AgentIdentityAuthRecord) -> bool {
238    record
239        .task_id
240        .as_deref()
241        .is_none_or(|task_id| task_id.trim().is_empty())
242}
243
244pub(super) fn record_matches_managed_chatgpt_binding(
245    record: &AgentIdentityAuthRecord,
246    binding: &ManagedChatGptAgentIdentityBinding,
247) -> bool {
248    record.account_id == binding.account_id
249        && record.chatgpt_user_id == binding.chatgpt_user_id
250        && public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key).is_ok()
251}
252
253pub(super) fn classify_bootstrap_error(
254    operation: &'static str,
255    err: std::io::Error,
256) -> std::io::Error {
257    if is_retryable_io_registration_error(&err) {
258        std::io::Error::other(AgentIdentityAuthError::BootstrapUnavailable {
259            operation,
260            attempts: MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS,
261            message: err.to_string(),
262        })
263    } else {
264        err
265    }
266}
267
268pub(super) fn is_retryable_io_registration_error(err: &std::io::Error) -> bool {
269    err.get_ref().is_some_and(
270        <dyn std::error::Error + std::marker::Send + std::marker::Sync + 'static>::is::<
271            RetryableAgentIdentityRegistrationError,
272        >,
273    )
274}
275
276pub(super) async fn retry_registration<T, F, Fut>(mut operation: F) -> std::io::Result<T>
277where
278    F: FnMut() -> Fut,
279    Fut: Future<Output = std::io::Result<T>>,
280{
281    let mut attempt = 1;
282    loop {
283        match operation().await {
284            Ok(value) => return Ok(value),
285            Err(err)
286                if attempt < MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS
287                    && is_retryable_io_registration_error(&err) =>
288            {
289                tracing::warn!(
290                    attempt,
291                    max_attempts = MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS,
292                    error = %err,
293                    "agent identity registration attempt failed; retrying"
294                );
295                attempt += 1;
296            }
297            Err(err) => return Err(err),
298        }
299    }
300}
301
302async fn register_task_for_record_with_retries(
303    record: &AgentIdentityAuthRecord,
304    agent_identity_authapi_base_url: &str,
305    auth_route_config: Option<&AuthRouteConfig>,
306) -> std::io::Result<String> {
307    let task_registration_url =
308        agent_task_registration_url(agent_identity_authapi_base_url, &record.agent_runtime_id);
309    let client = create_default_auth_client(&task_registration_url, auth_route_config)?;
310    retry_registration(|| async {
311        register_task_for_record(&client, record, agent_identity_authapi_base_url).await
312    })
313    .await
314}
315
316async fn register_task_for_record(
317    client: &HttpClient,
318    record: &AgentIdentityAuthRecord,
319    agent_identity_authapi_base_url: &str,
320) -> std::io::Result<String> {
321    register_agent_task(
322        client,
323        agent_identity_authapi_base_url,
324        key_for_record(record),
325    )
326    .await
327    .map_err(|err| {
328        if is_retryable_registration_error(&err) {
329            std::io::Error::other(RetryableAgentIdentityRegistrationError::new(
330                err.to_string(),
331            ))
332        } else {
333            std::io::Error::other(err)
334        }
335    })
336}
337
338fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
339    AgentIdentityKey {
340        agent_runtime_id: &record.agent_runtime_id,
341        private_key_pkcs8_base64: &record.agent_private_key,
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use std::sync::Arc;
348    use std::sync::atomic::AtomicUsize;
349    use std::sync::atomic::Ordering;
350
351    use codex_agent_identity::generate_agent_key_material;
352    use pretty_assertions::assert_eq;
353    use serde_json::json;
354    use wiremock::Mock;
355    use wiremock::MockServer;
356    use wiremock::ResponseTemplate;
357    use wiremock::matchers::method;
358    use wiremock::matchers::path;
359
360    use super::*;
361
362    fn agent_identity_record(private_key: String) -> AgentIdentityAuthRecord {
363        AgentIdentityAuthRecord {
364            agent_runtime_id: "agent-runtime-1".to_string(),
365            agent_private_key: private_key,
366            account_id: "account-1".to_string(),
367            chatgpt_user_id: "user-1".to_string(),
368            email: Some("agent@example.com".to_string()),
369            plan_type: AccountPlanType::Plus,
370            chatgpt_account_is_fedramp: false,
371            task_id: None,
372        }
373    }
374
375    fn agent_identity_record_with_generated_key() -> AgentIdentityAuthRecord {
376        let key_material = generate_agent_key_material().expect("generate key material");
377        agent_identity_record(key_material.private_key_pkcs8_base64)
378    }
379
380    #[tokio::test]
381    async fn from_record_registers_task() -> anyhow::Result<()> {
382        let server = MockServer::start().await;
383        Mock::given(method("POST"))
384            .and(path("/v1/agent/agent-runtime-1/task/register"))
385            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
386                "task_id": "task-run-1",
387            })))
388            .expect(1)
389            .mount(&server)
390            .await;
391
392        let auth = AgentIdentityAuth::from_record(
393            agent_identity_record_with_generated_key(),
394            &server.uri(),
395            /*auth_route_config*/ None,
396        )
397        .await?;
398
399        assert_eq!(auth.run_task_id(), "task-run-1");
400        let requests = server
401            .received_requests()
402            .await
403            .expect("failed to fetch task registration request");
404        let request_body = requests[0]
405            .body_json::<serde_json::Value>()
406            .expect("task registration request should be JSON");
407        let request_body = request_body
408            .as_object()
409            .expect("request body should be object");
410        assert!(request_body.get("timestamp").is_some());
411        assert!(request_body.get("signature").is_some());
412        assert_eq!(request_body.len(), 2);
413        Ok(())
414    }
415
416    #[tokio::test]
417    async fn from_jwt_registers_task() -> anyhow::Result<()> {
418        let server = MockServer::start().await;
419        Mock::given(method("GET"))
420            .and(path("/backend-api/wham/agent-identities/jwks"))
421            .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
422            .expect(1)
423            .mount(&server)
424            .await;
425        Mock::given(method("POST"))
426            .and(path("/v1/agent/agent-runtime-1/task/register"))
427            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
428                "task_id": "task-run-1",
429            })))
430            .expect(1)
431            .mount(&server)
432            .await;
433
434        let record = agent_identity_record_with_generated_key();
435        let jwt = signed_agent_identity_jwt(&record)?;
436        let auth = AgentIdentityAuth::from_jwt(
437            &jwt,
438            &format!("{}/backend-api", server.uri()),
439            &server.uri(),
440            /*auth_route_config*/ None,
441        )
442        .await?;
443
444        assert_eq!(auth.record().agent_runtime_id, "agent-runtime-1");
445        assert_eq!(auth.run_task_id(), "task-run-1");
446        Ok(())
447    }
448
449    #[test]
450    fn run_task_is_shared_across_clones() {
451        let auth = AgentIdentityAuth::from_initialized_record(
452            agent_identity_record_with_generated_key(),
453            "task-run-1".to_string(),
454        );
455        let cloned = auth.clone();
456
457        assert!(Arc::ptr_eq(&auth.record, &cloned.record));
458        assert_eq!(cloned.run_task_id(), "task-run-1");
459    }
460
461    #[tokio::test]
462    async fn from_record_retries_transient_registration() -> anyhow::Result<()> {
463        let server = MockServer::start().await;
464        let request_count = Arc::new(AtomicUsize::new(0));
465        let response_count = Arc::clone(&request_count);
466        Mock::given(method("POST"))
467            .and(path("/v1/agent/agent-runtime-1/task/register"))
468            .respond_with(move |_request: &wiremock::Request| {
469                if response_count.fetch_add(1, Ordering::SeqCst) == 0 {
470                    ResponseTemplate::new(500)
471                } else {
472                    ResponseTemplate::new(200).set_body_json(json!({
473                        "task_id": "task-run-1",
474                    }))
475                }
476            })
477            .expect(2)
478            .mount(&server)
479            .await;
480        let auth = AgentIdentityAuth::from_record(
481            agent_identity_record_with_generated_key(),
482            &server.uri(),
483            /*auth_route_config*/ None,
484        )
485        .await?;
486
487        assert_eq!(request_count.load(Ordering::SeqCst), 2);
488        assert_eq!(auth.run_task_id(), "task-run-1");
489        Ok(())
490    }
491
492    fn signed_agent_identity_jwt(
493        record: &AgentIdentityAuthRecord,
494    ) -> jsonwebtoken::errors::Result<String> {
495        let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
496        header.kid = Some("test-key".to_string());
497        jsonwebtoken::encode(
498            &header,
499            &json!({
500                "iss": "https://chatgpt.com/codex-backend/agent-identity",
501                "aud": "codex-app-server",
502                "iat": 1_700_000_000usize,
503                "exp": 4_000_000_000usize,
504                "agent_runtime_id": record.agent_runtime_id,
505                "agent_private_key": record.agent_private_key,
506                "account_id": record.account_id,
507                "chatgpt_user_id": record.chatgpt_user_id,
508                "email": record.email,
509                "plan_type": record.plan_type,
510                "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp,
511            }),
512            &jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?,
513        )
514    }
515
516    fn test_jwks_body() -> serde_json::Value {
517        json!({
518            "keys": [{
519                "kty": "RSA",
520                "kid": "test-key",
521                "use": "sig",
522                "alg": "RS256",
523                "n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
524                "e": "AQAB",
525            }]
526        })
527    }
528
529    const TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM: &[u8] = br#"-----BEGIN PRIVATE KEY-----
530MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
531bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
532FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
533mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
5340eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
535zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
536uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
5371ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
538RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
539NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
540rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
541Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
542aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
543iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
544YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
5453fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
546sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
547gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
548sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
549BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
550KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
551r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
55229HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
553J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
5545da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
5558U4M2TSWCKUY/A6sT4W8+mT9
556-----END PRIVATE KEY-----"#;
557}
codex_login/auth/agent_identity.rs

codex_login/auth/
agent_identity.rs
