feat: add run task identity primitives (#19047)

## Stack

This is PR 1 of the simplified HAI single-run-task stack:

- [#19047](https://github.com/openai/codex/pull/19047) Agent Identity
assertion and task-registration primitives, including the shared
run-task helper used by existing Agent Identity JWT auth.
- [#19049](https://github.com/openai/codex/pull/19049)
Disabled-by-default ChatGPT auth opt-in that provisions/reuses persisted
Agent Identity runtime auth and its single run task.
- [#19051](https://github.com/openai/codex/pull/19051) Run-scoped
provider auth that uses one backend-owned task id for first-party
inference and compaction requests.

[#19054](https://github.com/openai/codex/pull/19054) collapsed out of
the active stack because the simplified design no longer needs a
separate background/control-plane task helper.

## Summary

The simplified POC shape is one backend-owned task per Agent Identity
run. This PR makes the first layer match that final shape directly
instead of introducing task targets, caller-owned external task refs, or
intermediate wrappers that later PRs would need to undo.

What changed:

- keeps the `AgentAssertion` wire payload as `agent_runtime_id`,
`task_id`, `timestamp`, and `signature`
- exposes `register_agent_task` as the single task-registration helper
for both existing Agent Identity JWT auth and the ChatGPT-registration
path added later in the stack
- makes task registration send only the signed registration timestamp;
the backend owns the returned opaque task id
- removes the unused target/task-kind/external-task-ref surfaces from
`codex-agent-identity`
- keeps Agent Identity JWT JWKS lookup separate from agent/task
registration URL derivation
- updates Agent Identity JWT auth to register one run task during auth
construction and share that task across cloned auth handles

This PR intentionally does not enable ChatGPT-derived Agent Identity.
That opt-in and config gate are added in the next PR.

## Testing

- `just test -p codex-agent-identity`

codex-rs/login/src/auth/agent_identity.rs

@@ -1,75 +1,84 @@
+use std::sync::Arc;
+
 use codex_agent_identity::AgentIdentityKey;
 use codex_agent_identity::register_agent_task;
 use codex_protocol::account::PlanType as AccountPlanType;
-use std::env;
 
 use crate::default_client::build_reqwest_client;
 
 use super::storage::AgentIdentityAuthRecord;
 
-const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
-const CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR: &str = "CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL";
-
 #[derive(Clone, Debug)]
 pub struct AgentIdentityAuth {
+    inner: Arc<AgentIdentityAuthInner>,
+}
+
+#[derive(Debug)]
+struct AgentIdentityAuthInner {
     record: AgentIdentityAuthRecord,
-    process_task_id: String,
+    run_task_id: String,
 }
 
 impl AgentIdentityAuth {
-    pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
-        let agent_identity_authapi_base_url = agent_identity_authapi_base_url();
-        let process_task_id = register_agent_task(
+    pub async fn load(
+        record: AgentIdentityAuthRecord,
+        agent_identity_authapi_base_url: &str,
+    ) -> std::io::Result<Self> {
+        let run_task_id = register_agent_task(
             &build_reqwest_client(),
-            &agent_identity_authapi_base_url,
-            key(&record),
+            agent_identity_authapi_base_url,
+            key_for_record(&record),
         )
         .await
         .map_err(std::io::Error::other)?;
         Ok(Self {
-            record,
-            process_task_id,
+            inner: Arc::new(AgentIdentityAuthInner {
+                record,
+                run_task_id,
+            }),
         })
     }
 
-    pub fn record(&self) -> &AgentIdentityAuthRecord {
-        &self.record
+    #[cfg(test)]
+    fn from_initialized_record(record: AgentIdentityAuthRecord, run_task_id: String) -> Self {
+        Self {
+            inner: Arc::new(AgentIdentityAuthInner {
+                record,
+                run_task_id,
+            }),
+        }
     }
 
-    pub fn process_task_id(&self) -> &str {
-        &self.process_task_id
+    pub fn record(&self) -> &AgentIdentityAuthRecord {
+        &self.inner.record
+    }
+
+    pub fn run_task_id(&self) -> &str {
+        &self.inner.run_task_id
     }
 
     pub fn account_id(&self) -> &str {
-        &self.record.account_id
+        &self.inner.record.account_id
     }
 
     pub fn chatgpt_user_id(&self) -> &str {
-        &self.record.chatgpt_user_id
+        &self.inner.record.chatgpt_user_id
     }
 
     pub fn email(&self) -> &str {
-        &self.record.email
+        &self.inner.record.email
     }
 
     pub fn plan_type(&self) -> AccountPlanType {
-        self.record.plan_type
+        self.inner.record.plan_type
     }
 
     pub fn is_fedramp_account(&self) -> bool {
-        self.record.chatgpt_account_is_fedramp
+        self.inner.record.chatgpt_account_is_fedramp
     }
 }
 
-fn agent_identity_authapi_base_url() -> String {
-    env::var(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR)
-        .ok()
-        .map(|base_url| base_url.trim().trim_end_matches('/').to_string())
-        .filter(|base_url| !base_url.is_empty())
-        .unwrap_or_else(|| PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL.to_string())
-}
-
-fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
+fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
     AgentIdentityKey {
         agent_runtime_id: &record.agent_runtime_id,
         private_key_pkcs8_base64: &record.agent_private_key,
@@ -78,63 +87,110 @@ fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
+
+    use codex_agent_identity::generate_agent_key_material;
+    use pretty_assertions::assert_eq;
+    use serde_json::json;
+    use wiremock::Mock;
+    use wiremock::MockServer;
+    use wiremock::ResponseTemplate;
+    use wiremock::matchers::method;
+    use wiremock::matchers::path;
+
     use super::*;
-    use serial_test::serial;
 
-    #[test]
-    #[serial(codex_auth_env)]
-    fn agent_identity_authapi_base_url_prefers_env_value() {
-        let _guard = EnvVarGuard::set(
-            CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR,
-            "https://authapi.example.test/api/accounts/",
-        );
-        assert_eq!(
-            agent_identity_authapi_base_url(),
-            "https://authapi.example.test/api/accounts"
-        );
-    }
-
-    #[test]
-    #[serial(codex_auth_env)]
-    fn agent_identity_authapi_base_url_uses_prod_authapi_by_default() {
-        let _guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR);
-        assert_eq!(
-            agent_identity_authapi_base_url(),
-            PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL
-        );
-    }
-
-    struct EnvVarGuard {
-        key: &'static str,
-        original: Option<std::ffi::OsString>,
-    }
-
-    impl EnvVarGuard {
-        fn set(key: &'static str, value: &str) -> Self {
-            let original = env::var_os(key);
-            unsafe {
-                env::set_var(key, value);
-            }
-            Self { key, original }
-        }
-
-        fn remove(key: &'static str) -> Self {
-            let original = env::var_os(key);
-            unsafe {
-                env::remove_var(key);
-            }
-            Self { key, original }
+    fn agent_identity_record(private_key: String) -> AgentIdentityAuthRecord {
+        AgentIdentityAuthRecord {
+            agent_runtime_id: "agent-runtime-1".to_string(),
+            agent_private_key: private_key,
+            account_id: "account-1".to_string(),
+            chatgpt_user_id: "user-1".to_string(),
+            email: "agent@example.com".to_string(),
+            plan_type: AccountPlanType::Plus,
+            chatgpt_account_is_fedramp: false,
         }
     }
 
-    impl Drop for EnvVarGuard {
-        fn drop(&mut self) {
-            unsafe {
-                match &self.original {
-                    Some(value) => env::set_var(self.key, value),
-                    None => env::remove_var(self.key),
+    fn agent_identity_record_with_generated_key() -> AgentIdentityAuthRecord {
+        let key_material = generate_agent_key_material().expect("generate key material");
+        agent_identity_record(key_material.private_key_pkcs8_base64)
+    }
+
+    #[tokio::test]
+    async fn load_registers_run_task() -> anyhow::Result<()> {
+        let server = MockServer::start().await;
+        Mock::given(method("POST"))
+            .and(path("/v1/agent/agent-runtime-1/task/register"))
+            .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+                "task_id": "task-run-1",
+            })))
+            .expect(1)
+            .mount(&server)
+            .await;
+
+        let auth =
+            AgentIdentityAuth::load(agent_identity_record_with_generated_key(), &server.uri())
+                .await?;
+
+        assert_eq!(auth.run_task_id(), "task-run-1");
+        let requests = server
+            .received_requests()
+            .await
+            .expect("failed to fetch task registration request");
+        let request_body = requests[0]
+            .body_json::<serde_json::Value>()
+            .expect("task registration request should be JSON");
+        let request_body = request_body
+            .as_object()
+            .expect("request body should be object");
+        assert!(request_body.get("timestamp").is_some());
+        assert!(request_body.get("signature").is_some());
+        assert_eq!(request_body.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn run_task_is_shared_across_clones() {
+        let auth = AgentIdentityAuth::from_initialized_record(
+            agent_identity_record_with_generated_key(),
+            "task-run-1".to_string(),
+        );
+        let cloned = auth.clone();
+
+        assert!(Arc::ptr_eq(&auth.inner, &cloned.inner));
+        assert_eq!(cloned.run_task_id(), "task-run-1");
+    }
+
+    #[tokio::test]
+    async fn failed_run_task_registration_can_be_retried_on_next_call() -> anyhow::Result<()> {
+        let server = MockServer::start().await;
+        let request_count = Arc::new(AtomicUsize::new(0));
+        let response_count = Arc::clone(&request_count);
+        Mock::given(method("POST"))
+            .and(path("/v1/agent/agent-runtime-1/task/register"))
+            .respond_with(move |_request: &wiremock::Request| {
+                if response_count.fetch_add(1, Ordering::SeqCst) == 0 {
+                    ResponseTemplate::new(500)
+                } else {
+                    ResponseTemplate::new(200).set_body_json(json!({
+                        "task_id": "task-run-1",
+                    }))
                 }
-            }
-        }
+            })
+            .expect(2)
+            .mount(&server)
+            .await;
+        let record = agent_identity_record_with_generated_key();
+        AgentIdentityAuth::load(record.clone(), &server.uri())
+            .await
+            .expect_err("first registration should fail");
+        let auth = AgentIdentityAuth::load(record, &server.uri()).await?;
+
+        assert_eq!(request_count.load(Ordering::SeqCst), 2);
+        assert_eq!(auth.run_task_id(), "task-run-1");
+        Ok(())
     }
 }

codex-rs/login/src/auth/auth_tests.rs

@@ -348,6 +348,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .unwrap()
@@ -408,6 +409,7 @@ async fn loads_api_key_from_auth_json() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .unwrap()
@@ -500,6 +502,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -519,6 +522,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
         AuthCredentialsStoreMode::File,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("updated auth should parse");
@@ -824,6 +828,7 @@ async fn build_config(
         forced_login_method,
         forced_chatgpt_workspace_id,
         chatgpt_base_url: None,
+        agent_identity_authapi_base_url: None,
     }
 }
 
@@ -886,7 +891,7 @@ async fn load_auth_reads_access_token_from_env() {
         .mount(&server)
         .await;
     Mock::given(method("POST"))
-        .and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
+        .and(path("/v1/agent/agent-runtime-id/task/register"))
         .respond_with(ResponseTemplate::new(200).set_body_json(json!({
             "task_id": "task-123",
         })))
@@ -895,9 +900,8 @@ async fn load_auth_reads_access_token_from_env() {
         .await;
     let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity);
 
-    let chatgpt_base_url = format!("{}/backend-api", server.uri());
-    let _authapi_guard =
-        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
+    let authapi_base_url = server.uri();
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
     let auth = super::load_auth(
         codex_home.path(),
         /*enable_codex_api_key_env*/ false,
@@ -905,6 +909,7 @@ async fn load_auth_reads_access_token_from_env() {
         /*forced_chatgpt_workspace_id*/ None,
         Some(&chatgpt_base_url),
         AuthKeyringBackendKind::Direct,
+        Some(&authapi_base_url),
     )
     .await
     .expect("env auth should load")
@@ -914,7 +919,7 @@ async fn load_auth_reads_access_token_from_env() {
         panic!("env auth should load as agent identity");
     };
     assert_eq!(agent_identity.record(), &expected_record);
-    assert_eq!(agent_identity.process_task_id(), "task-123");
+    assert_eq!(agent_identity.run_task_id(), "task-123");
     assert!(
         !get_auth_file(codex_home.path()).exists(),
         "env auth should not write auth.json"
@@ -951,6 +956,7 @@ async fn load_auth_reads_personal_access_token_from_env() {
             /*forced_chatgpt_workspace_id*/ None,
             /*chatgpt_base_url*/ None,
             AuthKeyringBackendKind::default(),
+            /*agent_identity_authapi_base_url*/ None,
         )
         .await
         .expect("env auth should load")
@@ -1119,6 +1125,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("env auth should load")
@@ -1226,6 +1233,7 @@ async fn enforce_login_restrictions_logs_out_for_personal_access_token_workspace
         forced_login_method: None,
         forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
         chatgpt_base_url: None,
+        agent_identity_authapi_base_url: None,
     };
 
     let err = super::enforce_login_restrictions(&config)
@@ -1317,16 +1325,15 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat
         .mount(&server)
         .await;
     Mock::given(method("POST"))
-        .and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
+        .and(path("/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 authapi_base_url = server.uri();
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
     save_auth(
         codex_home.path(),
         &AuthDotJson {
@@ -1350,6 +1357,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat
         forced_login_method: None,
         forced_chatgpt_workspace_id: Some(vec![WORKSPACE_ID_ALLOWED.to_string()]),
         chatgpt_base_url: Some(chatgpt_base_url),
+        agent_identity_authapi_base_url: Some(authapi_base_url),
     };
 
     let err = super::enforce_login_restrictions(&config)
@@ -1563,19 +1571,19 @@ async fn assert_agent_identity_plan_alias(
         .mount(&server)
         .await;
     Mock::given(method("POST"))
-        .and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
+        .and(path("/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");
+    let authapi_base_url = server.uri();
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
+    let auth =
+        CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url), Some(&authapi_base_url))
+            .await
+            .expect("agent identity auth");
 
     pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type));
     server.verify().await;
@@ -1603,6 +1611,7 @@ async fn plan_type_maps_known_plan() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -1633,6 +1642,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -1666,6 +1676,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -1699,6 +1710,7 @@ async fn plan_type_maps_unknown_to_unknown() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")
@@ -1729,6 +1741,7 @@ async fn missing_plan_type_maps_to_unknown() {
         /*forced_chatgpt_workspace_id*/ None,
         /*chatgpt_base_url*/ None,
         AuthKeyringBackendKind::Direct,
+        /*agent_identity_authapi_base_url*/ None,
     )
     .await
     .expect("load auth")

codex-rs/login/src/auth/manager.rs

@@ -18,6 +18,7 @@ use std::sync::atomic::Ordering;
 use tokio::sync::Semaphore;
 use tokio::sync::watch;
 
+use codex_agent_identity::ChatGptEnvironment;
 use codex_agent_identity::decode_agent_identity_jwt;
 use codex_agent_identity::fetch_agent_identity_jwks;
 use codex_app_server_protocol::AuthMode;
@@ -104,7 +105,6 @@ const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be
 const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str =
     "Your access token could not be refreshed. Please log out and sign in again.";
 const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again.";
-const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
 const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
 pub(super) const REVOKE_TOKEN_URL: &str = "https://auth.openai.com/oauth/revoke";
 pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE";
@@ -218,6 +218,7 @@ impl CodexAuth {
         auth_credentials_store_mode: AuthCredentialsStoreMode,
         chatgpt_base_url: Option<&str>,
         keyring_backend_kind: AuthKeyringBackendKind,
+        agent_identity_authapi_base_url: Option<&str>,
     ) -> std::io::Result<Self> {
         let auth_mode = auth_dot_json.resolved_mode();
         let client = create_client();
@@ -233,7 +234,12 @@ impl CodexAuth {
                     "agent identity auth is missing an agent identity token.",
                 ));
             };
-            return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await;
+            return Self::from_agent_identity_jwt(
+                &agent_identity,
+                chatgpt_base_url,
+                agent_identity_authapi_base_url,
+            )
+            .await;
         }
         if auth_mode == ApiAuthMode::PersonalAccessToken {
             let Some(personal_access_token) = auth_dot_json.personal_access_token.as_deref() else {
@@ -292,6 +298,7 @@ impl CodexAuth {
             /*forced_chatgpt_workspace_id*/ None,
             chatgpt_base_url,
             keyring_backend_kind,
+            /*agent_identity_authapi_base_url*/ None,
         )
         .await
     }
@@ -299,13 +306,19 @@ impl CodexAuth {
     pub async fn from_agent_identity_jwt(
         jwt: &str,
         chatgpt_base_url: Option<&str>,
+        agent_identity_authapi_base_url_override: Option<&str>,
     ) -> std::io::Result<Self> {
         let base_url = chatgpt_base_url
-            .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
+            .unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
             .trim_end_matches('/')
             .to_string();
+        let resolved_authapi_base_url = resolve_agent_identity_authapi_base_url(
+            Some(&base_url),
+            agent_identity_authapi_base_url_override,
+        )?;
         let record = verified_agent_identity_record(jwt, &base_url).await?;
-        Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
+        let auth = AgentIdentityAuth::load(record, &resolved_authapi_base_url).await?;
+        Ok(Self::AgentIdentity(auth))
     }
 
     pub async fn from_personal_access_token(access_token: &str) -> std::io::Result<Self> {
@@ -580,6 +593,22 @@ async fn verified_agent_identity_record(
     Ok(claims.into())
 }
 
+fn resolve_agent_identity_authapi_base_url(
+    chatgpt_base_url: Option<&str>,
+    agent_identity_authapi_base_url_override: Option<&str>,
+) -> std::io::Result<String> {
+    if let Some(base_url) = agent_identity_authapi_base_url_override {
+        return Ok(base_url.trim_end_matches('/').to_string());
+    }
+
+    let environment = match chatgpt_base_url {
+        Some(chatgpt_base_url) => ChatGptEnvironment::from_chatgpt_base_url(chatgpt_base_url)
+            .map_err(std::io::Error::other)?,
+        None => ChatGptEnvironment::default(),
+    };
+    Ok(environment.agent_identity_authapi_base_url().to_string())
+}
+
 /// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)`
 /// if a file was removed, `Ok(false)` if no auth file was present.
 pub fn logout(
@@ -672,7 +701,7 @@ pub async fn login_with_access_token(
         }
         CodexAccessToken::AgentIdentityJwt(jwt) => {
             let base_url = chatgpt_base_url
-                .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
+                .unwrap_or(ChatGptEnvironment::default().chatgpt_base_url())
                 .trim_end_matches('/')
                 .to_string();
             verified_agent_identity_record(jwt, &base_url).await?;
@@ -763,6 +792,7 @@ pub struct AuthConfig {
     pub keyring_backend_kind: AuthKeyringBackendKind,
     pub forced_login_method: Option<ForcedLoginMethod>,
     pub chatgpt_base_url: Option<String>,
+    pub agent_identity_authapi_base_url: Option<String>,
     pub forced_chatgpt_workspace_id: Option<Vec<String>>,
 }
 
@@ -774,6 +804,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<
         /*forced_chatgpt_workspace_id*/ None,
         config.chatgpt_base_url.as_deref(),
         config.keyring_backend_kind,
+        config.agent_identity_authapi_base_url.as_deref(),
     )
     .await?
     else {
@@ -916,6 +947,7 @@ async fn load_auth(
     forced_chatgpt_workspace_id: Option<&[String]>,
     chatgpt_base_url: Option<&str>,
     keyring_backend_kind: AuthKeyringBackendKind,
+    agent_identity_authapi_base_url: Option<&str>,
 ) -> std::io::Result<Option<CodexAuth>> {
     // API key via env var takes precedence over any other auth method.
     if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() {
@@ -936,6 +968,7 @@ async fn load_auth(
             AuthCredentialsStoreMode::Ephemeral,
             chatgpt_base_url,
             keyring_backend_kind,
+            agent_identity_authapi_base_url,
         )
         .await?;
         if let CodexAuth::PersonalAccessToken(auth) = &auth {
@@ -951,11 +984,13 @@ async fn load_auth(
                 ensure_personal_access_token_workspace_allowed(forced_chatgpt_workspace_id, &auth)?;
                 Ok(Some(CodexAuth::PersonalAccessToken(auth)))
             }
-            CodexAccessToken::AgentIdentityJwt(jwt) => {
-                CodexAuth::from_agent_identity_jwt(jwt, chatgpt_base_url)
-                    .await
-                    .map(Some)
-            }
+            CodexAccessToken::AgentIdentityJwt(jwt) => CodexAuth::from_agent_identity_jwt(
+                jwt,
+                chatgpt_base_url,
+                agent_identity_authapi_base_url,
+            )
+            .await
+            .map(Some),
         };
     }
 
@@ -981,6 +1016,7 @@ async fn load_auth(
         auth_credentials_store_mode,
         chatgpt_base_url,
         keyring_backend_kind,
+        agent_identity_authapi_base_url,
     )
     .await?;
     if let CodexAuth::PersonalAccessToken(auth) = &auth {
@@ -1573,6 +1609,7 @@ impl AuthManager {
             forced_chatgpt_workspace_id.as_deref(),
             chatgpt_base_url.as_deref(),
             keyring_backend_kind,
+            /*agent_identity_authapi_base_url*/ None,
         )
         .await
         .ok()
@@ -1799,6 +1836,7 @@ impl AuthManager {
             forced_chatgpt_workspace_id.as_deref(),
             self.chatgpt_base_url.as_deref(),
             self.keyring_backend_kind,
+            /*agent_identity_authapi_base_url*/ None,
         )
         .await
         .ok()