diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 42cd415217ee..be4f56097f7d 100644 --- a/codex-rs/core/src/connectors.rs +++ b/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(), diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 23bbdb504ad1..116bc0993999 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -1,9 +1,6 @@ -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; @@ -11,50 +8,33 @@ 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>, -} - -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 { + 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, } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index b6c86761ef02..b3fe02125c4f 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/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(), diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index ad6d644493e4..ed4c0f01d27a 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/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 { + pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result { 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, - ) -> 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 { 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 diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index d9e31e78277e..3c7f4dbd0fc2 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/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) diff --git a/codex-rs/models-manager/src/manager_tests.rs b/codex-rs/models-manager/src/manager_tests.rs index 57c42be0154d..0f09980349db 100644 --- a/codex-rs/models-manager/src/manager_tests.rs +++ b/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!["chatgpt-only", "api-model"] - ); -} - #[tokio::test] async fn static_manager_reads_latest_auth_mode() { let auth_manager =