refactor: load agent identity runtime eagerly

codex-rs/core/src/connectors.rs

@@ -148,7 +148,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools(
     let auth = auth_manager.auth().await;
     if !config
         .features
-        .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
+        .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
     {
         return Some(Vec::new());
     }
@@ -220,7 +220,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager(
     let auth = auth_manager.auth().await;
     if !config
         .features
-        .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth))
+        .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend))
     {
         return Ok(AccessibleConnectorsStatus {
             connectors: Vec::new(),

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

@@ -1,60 +1,40 @@
-use std::sync::Arc;
-
 use codex_agent_identity::AgentIdentityKey;
 use codex_agent_identity::register_agent_task;
 use codex_protocol::account::PlanType as AccountPlanType;
-use tokio::sync::OnceCell;
 
 use crate::default_client::build_reqwest_client;
 
 use super::storage::AgentIdentityAuthRecord;
 
 const AGENT_IDENTITY_AUTHAPI_BASE_URL: &str = "https://auth.openai.com/api/accounts";
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct AgentIdentityAuth {
     record: AgentIdentityAuthRecord,
-    process_task_id: Arc<OnceCell<String>>,
-}
-
-impl Clone for AgentIdentityAuth {
-    fn clone(&self) -> Self {
-        Self {
-            record: self.record.clone(),
-            process_task_id: Arc::clone(&self.process_task_id),
-        }
-    }
+    process_task_id: String,
 }
 
 impl AgentIdentityAuth {
-    pub fn new(record: AgentIdentityAuthRecord) -> Self {
-        Self {
+    pub async fn load(record: AgentIdentityAuthRecord) -> std::io::Result<Self> {
+        let process_task_id = register_agent_task(
+            &build_reqwest_client(),
+            AGENT_IDENTITY_AUTHAPI_BASE_URL,
+            key(&record),
+        )
+        .await
+        .map_err(std::io::Error::other)?;
+        Ok(Self {
             record,
-            process_task_id: Arc::new(OnceCell::new()),
-        }
+            process_task_id,
+        })
     }
 
     pub fn record(&self) -> &AgentIdentityAuthRecord {
         &self.record
     }
 
-    pub fn process_task_id(&self) -> Option<&str> {
-        self.process_task_id.get().map(String::as_str)
-    }
-
-    pub async fn ensure_runtime(&self) -> std::io::Result<()> {
-        self.process_task_id
-            .get_or_try_init(|| async {
-                register_agent_task(
-                    &build_reqwest_client(),
-                    AGENT_IDENTITY_AUTHAPI_BASE_URL,
-                    self.key(),
-                )
-                .await
-                .map_err(std::io::Error::other)
-            })
-            .await
-            .map(|_| ())
+    pub fn process_task_id(&self) -> &str {
+        &self.process_task_id
     }
 
     pub fn account_id(&self) -> &str {
@@ -76,11 +56,11 @@ impl AgentIdentityAuth {
     pub fn is_fedramp_account(&self) -> bool {
         self.record.chatgpt_account_is_fedramp
     }
+}
 
-    fn key(&self) -> AgentIdentityKey<'_> {
-        AgentIdentityKey {
-            agent_runtime_id: &self.record.agent_runtime_id,
-            private_key_pkcs8_base64: &self.record.agent_private_key,
-        }
+fn key(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> {
+    AgentIdentityKey {
+        agent_runtime_id: &record.agent_runtime_id,
+        private_key_pkcs8_base64: &record.agent_private_key,
     }
 }

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

@@ -625,33 +625,6 @@ impl Drop for EnvVarGuard {
     }
 }
 
-#[tokio::test]
-#[serial(codex_auth_env)]
-async fn load_auth_reads_agent_identity_from_env() {
-    let codex_home = tempdir().unwrap();
-    let expected_record = agent_identity_record("account-123");
-    let agent_identity = fake_agent_identity_jwt(&expected_record).expect("fake agent identity");
-    let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
-
-    let auth = super::load_auth(
-        codex_home.path(),
-        /*enable_codex_api_key_env*/ false,
-        AuthCredentialsStoreMode::File,
-    )
-    .await
-    .expect("env auth should load")
-    .expect("env auth should be present");
-
-    let CodexAuth::AgentIdentity(agent_identity) = auth else {
-        panic!("env auth should load as agent identity");
-    };
-    assert_eq!(agent_identity.record(), &expected_record);
-    assert!(
-        !get_auth_file(codex_home.path()).exists(),
-        "env auth should not write auth.json"
-    );
-}
-
 #[tokio::test]
 #[serial(codex_auth_env)]
 async fn load_auth_keeps_codex_api_key_env_precedence() {
@@ -805,9 +778,11 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
 }
 
 fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord {
+    let key_material =
+        codex_agent_identity::generate_agent_key_material().expect("generate agent key material");
     AgentIdentityAuthRecord {
         agent_runtime_id: "agent-runtime-id".to_string(),
-        agent_private_key: "private-key".to_string(),
+        agent_private_key: key_material.private_key_pkcs8_base64,
         account_id: account_id.to_string(),
         chatgpt_user_id: "user-id".to_string(),
         email: "user@example.com".to_string(),

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

@@ -212,7 +212,7 @@ impl CodexAuth {
                     "agent identity auth is missing an agent identity token.",
                 ));
             };
-            return Self::from_agent_identity_jwt(&agent_identity);
+            return Self::from_agent_identity_jwt(&agent_identity).await;
         }
 
         let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
@@ -246,9 +246,9 @@ impl CodexAuth {
         .await
     }
 
-    pub fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
+    pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
         let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
-        Ok(Self::AgentIdentity(AgentIdentityAuth::new(record)))
+        Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
     }
 
     pub fn auth_mode(&self) -> AuthMode {
@@ -322,16 +322,6 @@ impl CodexAuth {
         }
     }
 
-    pub async fn initialize_runtime(
-        &self,
-        _chatgpt_base_url: Option<String>,
-    ) -> std::io::Result<()> {
-        match self {
-            Self::AgentIdentity(auth) => auth.ensure_runtime().await,
-            Self::ApiKey(_) | Self::Chatgpt(_) | Self::ChatgptAuthTokens(_) => Ok(()),
-        }
-    }
-
     /// Returns `None` if Codex backend auth does not expose an account id.
     pub fn get_account_id(&self) -> Option<String> {
         match self {
@@ -749,7 +739,9 @@ async fn load_auth(
     }
 
     if let Some(agent_identity) = read_codex_agent_identity_from_env() {
-        return CodexAuth::from_agent_identity_jwt(&agent_identity).map(Some);
+        return CodexAuth::from_agent_identity_jwt(&agent_identity)
+            .await
+            .map(Some);
     }
 
     // Fall back to the configured persistent store (file/keyring/auto) for managed auth.
@@ -1400,12 +1392,7 @@ impl AuthManager {
             tracing::error!("Failed to refresh token: {}", err);
             return Some(auth);
         }
-        let auth = self.auth_cached()?;
-        if let Err(err) = auth.initialize_runtime(self.chatgpt_base_url.clone()).await {
-            tracing::error!("Failed to initialize auth runtime: {err}");
-            return None;
-        }
-        Some(auth)
+        self.auth_cached()
     }
 
     /// Force a reload of the auth information from auth.json. Returns

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

@@ -21,23 +21,17 @@ struct AgentIdentityAuthProvider {
 impl AuthProvider for AgentIdentityAuthProvider {
     fn add_auth_headers(&self, headers: &mut HeaderMap) {
         let record = self.auth.record();
-        let header_value = self
-            .auth
-            .process_task_id()
-            .ok_or_else(|| std::io::Error::other("agent identity process task is not initialized"))
-            .and_then(|task_id| {
-                authorization_header_for_agent_task(
-                    AgentIdentityKey {
-                        agent_runtime_id: &record.agent_runtime_id,
-                        private_key_pkcs8_base64: &record.agent_private_key,
-                    },
-                    AgentTaskAuthorizationTarget {
-                        agent_runtime_id: &record.agent_runtime_id,
-                        task_id,
-                    },
-                )
-                .map_err(std::io::Error::other)
-            });
+        let header_value = authorization_header_for_agent_task(
+            AgentIdentityKey {
+                agent_runtime_id: &record.agent_runtime_id,
+                private_key_pkcs8_base64: &record.agent_private_key,
+            },
+            AgentTaskAuthorizationTarget {
+                agent_runtime_id: &record.agent_runtime_id,
+                task_id: self.auth.process_task_id(),
+            },
+        )
+        .map_err(std::io::Error::other);
 
         if let Ok(header_value) = header_value
             && let Ok(header) = HeaderValue::from_str(&header_value)

codex-rs/models-manager/src/manager_tests.rs

@@ -9,9 +9,6 @@ use codex_login::ExternalAuth;
 use codex_login::ExternalAuthRefreshContext;
 use codex_login::ExternalAuthTokens;
 use codex_login::TokenData;
-use codex_login::auth::AgentIdentityAuth;
-use codex_login::auth::AgentIdentityAuthRecord;
-use codex_protocol::account::PlanType;
 use codex_protocol::openai_models::ModelsResponse;
 use pretty_assertions::assert_eq;
 use serde_json::json;
@@ -237,18 +234,6 @@ c2ln",
         .expect("auth should be present")
 }
 
-fn agent_identity_auth_for_tests() -> CodexAuth {
-    CodexAuth::AgentIdentity(AgentIdentityAuth::new(AgentIdentityAuthRecord {
-        agent_runtime_id: "agent-runtime-id".to_string(),
-        agent_private_key: "agent-private-key".to_string(),
-        account_id: "account-id".to_string(),
-        chatgpt_user_id: "chatgpt-user-id".to_string(),
-        email: "agent@example.com".to_string(),
-        plan_type: PlanType::Pro,
-        chatgpt_account_is_fedramp: false,
-    }))
-}
-
 #[tokio::test]
 async fn get_model_info_tracks_fallback_usage() {
     let codex_home = tempdir().expect("temp dir");
@@ -713,43 +698,6 @@ async fn refresh_available_models_fetches_with_chatgpt_auth_tokens() {
     );
 }
 
-#[tokio::test]
-async fn refresh_available_models_fetches_with_agent_identity() {
-    let dynamic_slug = "dynamic-model-only-for-test-agent-identity";
-    let codex_home = tempdir().expect("temp dir");
-    let endpoint = TestModelsEndpoint::new(vec![vec![remote_model(
-        dynamic_slug,
-        "Agent Identity",
-        /*priority*/ 1,
-    )]]);
-    let manager = openai_manager_for_tests_with_auth(
-        codex_home.path().to_path_buf(),
-        endpoint.clone(),
-        Some(AuthManager::from_auth_for_testing(
-            agent_identity_auth_for_tests(),
-        )),
-    );
-
-    manager
-        .refresh_available_models(RefreshStrategy::Online)
-        .await
-        .expect("refresh should fetch with agent identity");
-
-    assert!(
-        manager
-            .get_remote_models()
-            .await
-            .iter()
-            .any(|candidate| candidate.slug == dynamic_slug),
-        "remote refresh should include models fetched with agent identity"
-    );
-    assert_eq!(
-        endpoint.fetch_count(),
-        1,
-        "endpoint should fetch models with agent identity"
-    );
-}
-
 #[test]
 fn build_available_models_picks_default_after_hiding_hidden_models() {
     let manager = static_manager_for_tests(ModelsResponse { models: Vec::new() });
@@ -768,35 +716,6 @@ fn build_available_models_picks_default_after_hiding_hidden_models() {
     assert_eq!(available, vec![expected_hidden, expected_visible]);
 }
 
-#[tokio::test]
-async fn static_manager_treats_agent_identity_as_backend_auth_for_filtering() {
-    let chatgpt_only_model = {
-        let mut model = remote_model("chatgpt-only", "ChatGPT Only", /*priority*/ 0);
-        model.supported_in_api = false;
-        model
-    };
-    let api_model = remote_model("api-model", "API Model", /*priority*/ 1);
-    let manager = StaticModelsManager::new(
-        Some(AuthManager::from_auth_for_testing(
-            agent_identity_auth_for_tests(),
-        )),
-        ModelsResponse {
-            models: vec![chatgpt_only_model, api_model],
-        },
-        CollaborationModesConfig::default(),
-    );
-
-    let agent_identity_models = manager.list_models(RefreshStrategy::Online).await;
-
-    assert_eq!(
-        agent_identity_models
-            .iter()
-            .map(|model| model.model.as_str())
-            .collect::<Vec<_>>(),
-        vec!["chatgpt-only", "api-model"]
-    );
-}
-
 #[tokio::test]
 async fn static_manager_reads_latest_auth_mode() {
     let auth_manager =