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/agent-identity/src/lib.rs

@@ -1,4 +1,6 @@
 use std::collections::BTreeMap;
+use std::error::Error as StdError;
+use std::fmt;
 use std::time::Duration;
 
 use anyhow::Context;
@@ -34,21 +36,66 @@ const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30);
 const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10);
 const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server";
 const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity";
+const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
+const PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
+const STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.api.openai.org/api/accounts";
+const AGENT_IDENTITY_KEY_SEED_BYTES: usize = 64;
+const AGENT_IDENTITY_KEY_DERIVATION_CONTEXT: &[u8] = b"codex-agent-identity-ed25519-v1";
 
-/// Stored key material for a registered agent identity.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum ChatGptEnvironment {
+    #[default]
+    Production,
+    Staging,
+}
+
+impl ChatGptEnvironment {
+    pub fn from_chatgpt_base_url(chatgpt_base_url: &str) -> Result<Self> {
+        match chatgpt_base_url.trim_end_matches('/') {
+            "https://chatgpt.com"
+            | "https://chatgpt.com/backend-api"
+            | "https://chatgpt.com/codex"
+            | "https://chatgpt.com/backend-api/codex"
+            | "https://chat.openai.com"
+            | "https://chat.openai.com/backend-api"
+            | "https://chat.openai.com/codex"
+            | "https://chat.openai.com/backend-api/codex" => Ok(Self::Production),
+            "https://chatgpt-staging.com"
+            | "https://chatgpt-staging.com/backend-api"
+            | "https://chatgpt-staging.com/codex"
+            | "https://chatgpt-staging.com/backend-api/codex" => Ok(Self::Staging),
+            _ => anyhow::bail!(
+                "Agent Identity only supports production and staging ChatGPT environments"
+            ),
+        }
+    }
+
+    pub fn chatgpt_base_url(self) -> &'static str {
+        match self {
+            Self::Production => "https://chatgpt.com/backend-api",
+            Self::Staging => "https://chatgpt-staging.com/backend-api",
+        }
+    }
+
+    pub fn agent_identity_authapi_base_url(self) -> &'static str {
+        match self {
+            Self::Production => PROD_AGENT_IDENTITY_AUTHAPI_BASE_URL,
+            Self::Staging => STAGING_AGENT_IDENTITY_AUTHAPI_BASE_URL,
+        }
+    }
+}
+
+/// Borrowed durable signing material for a registered agent identity.
+///
+/// This intentionally does not include a task id. Task ids are scoped to a
+/// single Codex run, while the agent runtime id and private key are the
+/// reusable identity material used to register and sign that run task.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct AgentIdentityKey<'a> {
     pub agent_runtime_id: &'a str,
     pub private_key_pkcs8_base64: &'a str,
 }
 
-/// Task binding to use when constructing a task-scoped AgentAssertion.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub struct AgentTaskAuthorizationTarget<'a> {
-    pub agent_runtime_id: &'a str,
-    pub task_id: &'a str,
-}
-
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub struct AgentBillOfMaterials {
     pub agent_version: String,
@@ -103,23 +150,92 @@ struct RegisterTaskResponse {
     encrypted_task_id_camel: Option<String>,
 }
 
+#[derive(Debug, Serialize)]
+struct RegisterAgentRequest {
+    abom: AgentBillOfMaterials,
+    agent_public_key: String,
+    capabilities: Vec<String>,
+    ttl: Option<u64>,
+}
+
+#[derive(Debug, Deserialize)]
+struct RegisterAgentResponse {
+    agent_runtime_id: String,
+}
+
+/// HTTP status failure returned by Agent Identity registration endpoints.
+#[derive(Debug)]
+pub struct AgentIdentityRegistrationHttpError {
+    operation: &'static str,
+    status: reqwest::StatusCode,
+    body: String,
+}
+
+impl AgentIdentityRegistrationHttpError {
+    fn new(operation: &'static str, status: reqwest::StatusCode, body: String) -> Self {
+        Self {
+            operation,
+            status,
+            body,
+        }
+    }
+
+    /// HTTP status returned by the registration endpoint.
+    pub fn status(&self) -> reqwest::StatusCode {
+        self.status
+    }
+}
+
+impl fmt::Display for AgentIdentityRegistrationHttpError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if self.body.is_empty() {
+            write!(f, "{} failed with status {}", self.operation, self.status)
+        } else {
+            write!(
+                f,
+                "{} failed with status {}: {}",
+                self.operation, self.status, self.body
+            )
+        }
+    }
+}
+
+impl StdError for AgentIdentityRegistrationHttpError {}
+
+/// Returns whether an Agent Identity registration error is safe to retry.
+pub fn is_retryable_registration_error(error: &anyhow::Error) -> bool {
+    error.chain().any(is_retryable_registration_cause)
+}
+
+fn is_retryable_registration_cause(cause: &(dyn StdError + 'static)) -> bool {
+    if let Some(error) = cause.downcast_ref::<AgentIdentityRegistrationHttpError>() {
+        return is_retryable_registration_status(error.status());
+    }
+
+    if let Some(error) = cause.downcast_ref::<reqwest::Error>() {
+        if let Some(status) = error.status() {
+            return is_retryable_registration_status(status);
+        }
+        return error.is_timeout() || error.is_connect() || error.is_request();
+    }
+
+    false
+}
+
+fn is_retryable_registration_status(status: reqwest::StatusCode) -> bool {
+    status == reqwest::StatusCode::TOO_MANY_REQUESTS || status.is_server_error()
+}
+
 pub fn authorization_header_for_agent_task(
     key: AgentIdentityKey<'_>,
-    target: AgentTaskAuthorizationTarget<'_>,
+    task_id: &str,
 ) -> Result<String> {
-    anyhow::ensure!(
-        key.agent_runtime_id == target.agent_runtime_id,
-        "agent task runtime {} does not match stored agent identity {}",
-        target.agent_runtime_id,
-        key.agent_runtime_id
-    );
-
     let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
     let envelope = AgentAssertionEnvelope {
-        agent_runtime_id: target.agent_runtime_id.to_string(),
+        agent_runtime_id: key.agent_runtime_id.to_string(),
-        task_id: target.task_id.to_string(),
+        task_id: task_id.to_string(),
         timestamp: timestamp.clone(),
-        signature: sign_agent_assertion_payload(key, target.task_id, &timestamp)?,
+        signature: sign_agent_assertion_payload(key, task_id, &timestamp)?,
     };
     let serialized_assertion = serialize_agent_assertion(&envelope)?;
     Ok(format!("AgentAssertion {serialized_assertion}"))
@@ -127,10 +243,10 @@ pub fn authorization_header_for_agent_task(
 
 pub async fn fetch_agent_identity_jwks(
     client: &reqwest::Client,
-    chatgpt_base_url: &str,
+    agent_identity_jwt_base_url: &str,
 ) -> Result<JwkSet> {
     let response = client
-        .get(agent_identity_jwks_url(chatgpt_base_url))
+        .get(agent_identity_jwks_url(agent_identity_jwt_base_url))
         .timeout(AGENT_IDENTITY_JWKS_TIMEOUT)
         .send()
         .await
@@ -195,7 +311,7 @@ pub fn sign_task_registration_payload(
 
 pub async fn register_agent_task(
     client: &reqwest::Client,
-    chatgpt_base_url: &str,
+    agent_identity_authapi_base_url: &str,
     key: AgentIdentityKey<'_>,
 ) -> Result<String> {
     let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
@@ -203,7 +319,7 @@ pub async fn register_agent_task(
         signature: sign_task_registration_payload(key, &timestamp)?,
         timestamp,
     };
-    let url = agent_task_registration_url(chatgpt_base_url, key.agent_runtime_id);
+    let url = agent_task_registration_url(agent_identity_authapi_base_url, key.agent_runtime_id);
 
     let response = client
         .post(url)
@@ -220,7 +336,12 @@ pub async fn register_agent_task(
         } else {
             body
         };
-        anyhow::bail!("failed to register agent task with status {status}: {body}");
+        return Err(AgentIdentityRegistrationHttpError::new(
+            "agent task registration",
+            status,
+            body,
+        )
+        .into());
     }
 
     let response = response
@@ -231,6 +352,45 @@ pub async fn register_agent_task(
     task_id_from_register_task_response(key, response)
 }
 
+pub async fn register_agent_identity(
+    client: &reqwest::Client,
+    agent_identity_authapi_base_url: &str,
+    access_token: &str,
+    is_fedramp_account: bool,
+    key_material: &GeneratedAgentKeyMaterial,
+    abom: AgentBillOfMaterials,
+    capabilities: Vec<String>,
+) -> Result<String> {
+    let url = agent_registration_url(agent_identity_authapi_base_url);
+    let request = RegisterAgentRequest {
+        abom,
+        agent_public_key: key_material.public_key_ssh.clone(),
+        capabilities,
+        ttl: None,
+    };
+
+    let mut request_builder = client
+        .post(&url)
+        .bearer_auth(access_token)
+        .json(&request)
+        .timeout(AGENT_REGISTRATION_TIMEOUT);
+    if is_fedramp_account {
+        request_builder = request_builder.header("X-OpenAI-Fedramp", "true");
+    }
+
+    let response = request_builder
+        .send()
+        .await
+        .with_context(|| format!("failed to send agent identity registration request to {url}"))?
+        .error_for_status()
+        .with_context(|| format!("agent identity registration failed for {url}"))?
+        .json::<RegisterAgentResponse>()
+        .await
+        .with_context(|| format!("failed to parse agent identity response from {url}"))?;
+
+    Ok(response.agent_runtime_id)
+}
+
 fn task_id_from_register_task_response(
     key: AgentIdentityKey<'_>,
     response: RegisterTaskResponse,
@@ -260,10 +420,17 @@ pub fn decrypt_task_id_response(
 }
 
 pub fn generate_agent_key_material() -> Result<GeneratedAgentKeyMaterial> {
-    let mut secret_key_bytes = [0u8; 32];
+    let mut seed_material = [0u8; AGENT_IDENTITY_KEY_SEED_BYTES];
     OsRng
-        .try_fill_bytes(&mut secret_key_bytes)
+        .try_fill_bytes(&mut seed_material)
-        .context("failed to generate agent identity private key bytes")?;
+        .context("failed to generate agent identity private key seed material")?;
+    // Ed25519 stores a 32-byte seed, so derive it from all sampled seed material.
+    let mut digest = Sha512::new();
+    digest.update(AGENT_IDENTITY_KEY_DERIVATION_CONTEXT);
+    digest.update(seed_material);
+    let digest = digest.finalize();
+    let mut secret_key_bytes = [0u8; 32];
+    secret_key_bytes.copy_from_slice(&digest[..32]);
     let signing_key = SigningKey::from_bytes(&secret_key_bytes);
     let private_key_pkcs8 = signing_key
         .to_pkcs8_der()
@@ -296,23 +463,22 @@ pub fn curve25519_secret_key_from_private_key_pkcs8_base64(
     Ok(curve25519_secret_key_from_signing_key(&signing_key))
 }
 
-pub fn agent_registration_url(chatgpt_base_url: &str) -> String {
+pub fn agent_registration_url(agent_identity_authapi_base_url: &str) -> String {
-    let trimmed = chatgpt_base_url.trim_end_matches('/');
+    agent_identity_authapi_url(agent_identity_authapi_base_url, "/v1/agent/register")
-    format!("{trimmed}/v1/agent/register")
 }
 
-pub fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String {
+pub fn agent_task_registration_url(
-    let trimmed = chatgpt_base_url.trim_end_matches('/');
+    agent_identity_authapi_base_url: &str,
-    format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register")
+    agent_runtime_id: &str,
+) -> String {
+    agent_identity_authapi_url(
+        agent_identity_authapi_base_url,
+        &format!("/v1/agent/{agent_runtime_id}/task/register"),
+    )
 }
 
-pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String {
+pub fn agent_identity_jwks_url(agent_identity_jwt_base_url: &str) -> String {
-    let trimmed = chatgpt_base_url.trim_end_matches('/');
+    let trimmed = agent_identity_jwt_base_url.trim_end_matches('/');
-    format!("{trimmed}/authenticate_app_v2")
-}
-
-pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
-    let trimmed = chatgpt_base_url.trim_end_matches('/');
     if trimmed.contains("/backend-api") {
         format!("{trimmed}/wham/agent-identities/jwks")
     } else {
@@ -320,15 +486,9 @@ pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
     }
 }
 
-pub fn agent_identity_request_id() -> Result<String> {
+fn agent_identity_authapi_url(agent_identity_authapi_base_url: &str, api_path: &str) -> String {
-    let mut request_id_bytes = [0u8; 16];
+    let base_url = agent_identity_authapi_base_url.trim_end_matches('/');
-    OsRng
+    format!("{base_url}{api_path}")
-        .try_fill_bytes(&mut request_id_bytes)
-        .context("failed to generate agent identity request id")?;
-    Ok(format!(
-        "codex-agent-identity-{}",
-        URL_SAFE_NO_PAD.encode(request_id_bytes)
-    ))
 }
 
 pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials {
@@ -412,6 +572,24 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn register_task_request_uses_single_run_task_shape() {
+        let request = RegisterTaskRequest {
+            timestamp: "2026-04-23T00:00:00Z".to_string(),
+            signature: "signature".to_string(),
+        };
+
+        let serialized = serde_json::to_value(request).expect("serialize request");
+
+        assert_eq!(
+            serialized,
+            serde_json::json!({
+                "timestamp": "2026-04-23T00:00:00Z",
+                "signature": "signature",
+            })
+        );
+    }
+
     #[test]
     fn authorization_header_for_agent_task_serializes_signed_agent_assertion() {
         let signing_key = SigningKey::from_bytes(&[7u8; 32]);
@@ -422,13 +600,9 @@ mod tests {
             agent_runtime_id: "agent-123",
             private_key_pkcs8_base64: &BASE64_STANDARD.encode(private_key.as_bytes()),
         };
-        let target = AgentTaskAuthorizationTarget {
-            agent_runtime_id: "agent-123",
-            task_id: "task-123",
-        };
 
-        let header =
+        let header = authorization_header_for_agent_task(key, "task-123")
-            authorization_header_for_agent_task(key, target).expect("build agent assertion header");
+            .expect("build agent assertion header");
         let token = header
             .strip_prefix("AgentAssertion ")
             .expect("agent assertion scheme");
@@ -464,31 +638,6 @@ mod tests {
             .expect("signature should verify");
     }
 
-    #[test]
-    fn authorization_header_for_agent_task_rejects_mismatched_runtime() {
-        let signing_key = SigningKey::from_bytes(&[7u8; 32]);
-        let private_key = signing_key
-            .to_pkcs8_der()
-            .expect("encode test key material");
-        let private_key_pkcs8_base64 = BASE64_STANDARD.encode(private_key.as_bytes());
-        let key = AgentIdentityKey {
-            agent_runtime_id: "agent-123",
-            private_key_pkcs8_base64: &private_key_pkcs8_base64,
-        };
-        let target = AgentTaskAuthorizationTarget {
-            agent_runtime_id: "agent-456",
-            task_id: "task-123",
-        };
-
-        let error = authorization_header_for_agent_task(key, target)
-            .expect_err("runtime mismatch should fail");
-
-        assert_eq!(
-            error.to_string(),
-            "agent task runtime agent-456 does not match stored agent identity agent-123"
-        );
-    }
-
     #[test]
     fn decode_agent_identity_jwt_reads_claims() {
         let jwt = jwt_with_payload(serde_json::json!({
@@ -704,7 +853,98 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
     }
 
     #[test]
-    fn agent_identity_jwks_url_uses_backend_api_base_url() {
+    fn chatgpt_environment_maps_known_urls_to_authapi() -> anyhow::Result<()> {
+        assert_eq!(
+            ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt.com/backend-api/codex")?,
+            ChatGptEnvironment::Production
+        );
+        assert_eq!(
+            ChatGptEnvironment::Production.agent_identity_authapi_base_url(),
+            "https://auth.openai.com/api/accounts"
+        );
+        assert_eq!(
+            ChatGptEnvironment::from_chatgpt_base_url("https://chatgpt-staging.com/backend-api")?,
+            ChatGptEnvironment::Staging
+        );
+        assert_eq!(
+            ChatGptEnvironment::Staging.agent_identity_authapi_base_url(),
+            "https://auth.api.openai.org/api/accounts"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn chatgpt_environment_rejects_custom_urls() {
+        assert!(ChatGptEnvironment::from_chatgpt_base_url("http://localhost:8080").is_err(),);
+    }
+
+    #[test]
+    fn agent_registration_url_appends_to_authapi_base_url() {
+        assert_eq!(
+            agent_registration_url("https://auth.openai.com/api/accounts"),
+            "https://auth.openai.com/api/accounts/v1/agent/register"
+        );
+        assert_eq!(
+            agent_registration_url("http://localhost:8080"),
+            "http://localhost:8080/v1/agent/register"
+        );
+        assert_eq!(
+            agent_registration_url("http://localhost:8080/backend-api"),
+            "http://localhost:8080/backend-api/v1/agent/register"
+        );
+    }
+
+    #[test]
+    fn agent_task_registration_url_appends_to_authapi_base_url() {
+        assert_eq!(
+            agent_task_registration_url("https://auth.openai.com/api/accounts", "agent-runtime-id"),
+            "https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
+        );
+        assert_eq!(
+            agent_task_registration_url(
+                "https://auth.openai.com/api/accounts/",
+                "agent-runtime-id"
+            ),
+            "https://auth.openai.com/api/accounts/v1/agent/agent-runtime-id/task/register"
+        );
+        assert_eq!(
+            agent_task_registration_url("http://localhost:8080", "agent-runtime-id"),
+            "http://localhost:8080/v1/agent/agent-runtime-id/task/register"
+        );
+    }
+
+    #[test]
+    fn retryable_registration_error_accepts_429_and_5xx() {
+        let too_many_requests = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
+            "agent registration",
+            reqwest::StatusCode::TOO_MANY_REQUESTS,
+            "rate limited".to_string(),
+        ));
+        let unavailable = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
+            "agent registration",
+            reqwest::StatusCode::SERVICE_UNAVAILABLE,
+            "try later".to_string(),
+        ));
+
+        assert!(is_retryable_registration_error(&too_many_requests));
+        assert!(is_retryable_registration_error(&unavailable));
+    }
+
+    #[test]
+    fn retryable_registration_error_rejects_hard_failures() {
+        let forbidden = anyhow::Error::new(AgentIdentityRegistrationHttpError::new(
+            "agent registration",
+            reqwest::StatusCode::FORBIDDEN,
+            "not allowed".to_string(),
+        ));
+        let malformed = anyhow::anyhow!("failed to sign registration request");
+
+        assert!(!is_retryable_registration_error(&forbidden));
+        assert!(!is_retryable_registration_error(&malformed));
+    }
+
+    #[test]
+    fn agent_identity_jwks_url_uses_agent_identity_jwt_route() {
         assert_eq!(
             agent_identity_jwks_url("https://chatgpt.com/backend-api"),
             "https://chatgpt.com/backend-api/wham/agent-identities/jwks"
@@ -716,7 +956,7 @@ J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
     }
 
     #[test]
-    fn agent_identity_jwks_url_uses_codex_api_base_url() {
+    fn agent_identity_jwks_url_uses_jwt_issuer_base_url() {
         assert_eq!(
             agent_identity_jwks_url("http://localhost:8080/api/codex"),
             "http://localhost:8080/api/codex/agent-identities/jwks"

codex-rs/cli/src/main.rs

@@ -1731,9 +1731,12 @@ async fn load_exec_server_remote_auth_provider(
         let agent_identity_jwt = read_codex_access_token_from_env().ok_or_else(|| {
             anyhow::anyhow!("CODEX_ACCESS_TOKEN is required when --use-agent-identity-auth is set")
         })?;
-        let auth =
+        let auth = CodexAuth::from_agent_identity_jwt(
-            CodexAuth::from_agent_identity_jwt(&agent_identity_jwt, Some(&config.chatgpt_base_url))
+            &agent_identity_jwt,
-                .await?;
+            Some(&config.chatgpt_base_url),
+            /*agent_identity_authapi_base_url_override*/ None,
+        )
+        .await?;
         return Ok(codex_model_provider::auth_provider_from_auth(&auth));
     }
 

codex-rs/exec/src/lib.rs

@@ -475,6 +475,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result
         forced_login_method: config.forced_login_method,
         forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
         chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
+        agent_identity_authapi_base_url: None,
     })
     .await
     {

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> {
+    pub async fn load(
-        let agent_identity_authapi_base_url = agent_identity_authapi_base_url();
+        record: AgentIdentityAuthRecord,
-        let process_task_id = register_agent_task(
+        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,
+            agent_identity_authapi_base_url,
-            key(&record),
+            key_for_record(&record),
         )
         .await
         .map_err(std::io::Error::other)?;
         Ok(Self {
-            record,
+            inner: Arc::new(AgentIdentityAuthInner {
-            process_task_id,
+                record,
+                run_task_id,
+            }),
         })
     }
 
-    pub fn record(&self) -> &AgentIdentityAuthRecord {
+    #[cfg(test)]
-        &self.record
+    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 {
+    pub fn record(&self) -> &AgentIdentityAuthRecord {
-        &self.process_task_id
+        &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 {
+fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
-    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<'_> {
     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]
+    fn agent_identity_record(private_key: String) -> AgentIdentityAuthRecord {
-    #[serial(codex_auth_env)]
+        AgentIdentityAuthRecord {
-    fn agent_identity_authapi_base_url_prefers_env_value() {
+            agent_runtime_id: "agent-runtime-1".to_string(),
-        let _guard = EnvVarGuard::set(
+            agent_private_key: private_key,
-            CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR,
+            account_id: "account-1".to_string(),
-            "https://authapi.example.test/api/accounts/",
+            chatgpt_user_id: "user-1".to_string(),
-        );
+            email: "agent@example.com".to_string(),
-        assert_eq!(
+            plan_type: AccountPlanType::Plus,
-            agent_identity_authapi_base_url(),
+            chatgpt_account_is_fedramp: false,
-            "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 }
         }
     }
 
-    impl Drop for EnvVarGuard {
+    fn agent_identity_record_with_generated_key() -> AgentIdentityAuthRecord {
-        fn drop(&mut self) {
+        let key_material = generate_agent_key_material().expect("generate key material");
-            unsafe {
+        agent_identity_record(key_material.private_key_pkcs8_base64)
-                match &self.original {
+    }
-                    Some(value) => env::set_var(self.key, value),
+
-                    None => env::remove_var(self.key),
+    #[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_base_url = server.uri();
-    let _authapi_guard =
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
-        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
     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_base_url = server.uri();
-    let _authapi_guard =
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
-        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
     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_base_url = server.uri();
-    let _authapi_guard =
+    let chatgpt_base_url = format!("{authapi_base_url}/backend-api");
-        EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
+    let auth =
-    let auth = CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url))
+        CodexAuth::from_agent_identity_jwt(&jwt, Some(&chatgpt_base_url), Some(&authapi_base_url))
-        .await
+            .await
-        .expect("agent identity auth");
+            .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) => {
+            CodexAccessToken::AgentIdentityJwt(jwt) => CodexAuth::from_agent_identity_jwt(
-                CodexAuth::from_agent_identity_jwt(jwt, chatgpt_base_url)
+                jwt,
-                    .await
+                chatgpt_base_url,
-                    .map(Some)
+                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()

codex-rs/model-provider/src/auth.rs

@@ -1,7 +1,6 @@
 use std::sync::Arc;
 
 use codex_agent_identity::AgentIdentityKey;
-use codex_agent_identity::AgentTaskAuthorizationTarget;
 use codex_agent_identity::authorization_header_for_agent_task;
 use codex_api::AuthProvider;
 use codex_api::SharedAuthProvider;
@@ -30,10 +29,7 @@ impl AuthProvider for AgentIdentityAuthProvider {
                 agent_runtime_id: &record.agent_runtime_id,
                 private_key_pkcs8_base64: &record.agent_private_key,
             },
-            AgentTaskAuthorizationTarget {
+            self.auth.run_task_id(),
-                agent_runtime_id: &record.agent_runtime_id,
-                task_id: self.auth.process_task_id(),
-            },
         )
         .map_err(std::io::Error::other);
 

codex-rs/tui/src/lib.rs

@@ -1161,6 +1161,7 @@ pub async fn run_main(
             forced_login_method: config.forced_login_method,
             forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
             chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
+            agent_identity_authapi_base_url: None,
         })
         .await
         {