diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0fae5a41b..55d02cd8b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2418,6 +2418,7 @@ version = "0.0.0" dependencies = [ "base64 0.22.1", "chrono", + "codex-agent-identity", "codex-backend-client", "codex-config", "codex-core", diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 8e75c0fc8..6c453e6f9 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -1038,6 +1038,7 @@ async fn remote_control_start_allows_missing_auth_when_enabled() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -1866,6 +1867,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -1961,6 +1963,7 @@ async fn persisted_enable_does_not_follow_auth_to_an_account_without_a_preferenc codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs index 630ce8c6b..6354e435c 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs @@ -182,6 +182,7 @@ async fn list_remote_control_clients_recovers_auth_after_unauthorized() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -266,6 +267,7 @@ async fn list_remote_control_clients_retries_unauthorized_only_once() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs index ff146492f..190526d72 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs @@ -536,6 +536,7 @@ async fn remote_control_handle_recovers_auth_before_refreshing_pairing() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -804,6 +805,7 @@ async fn remote_control_handle_discards_pairing_response_after_auth_change() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 3f594a1b4..239d5d9ea 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -2211,6 +2211,7 @@ mod tests { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -2308,6 +2309,7 @@ mod tests { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -2430,6 +2432,7 @@ mod tests { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index c0f310e70..5e5a1e94b 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -1399,8 +1399,8 @@ fn stored_auth_issues( codex_app_server_protocol::AuthMode::AgentIdentity => { if auth .agent_identity - .as_deref() - .is_none_or(|token| token.trim().is_empty()) + .as_ref() + .is_none_or(|agent_identity| !agent_identity.has_auth_material()) { issues.push("agent identity auth is missing an agent identity token"); } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 8b633265e..9d278c212 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1738,12 +1738,9 @@ 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 = CodexAuth::from_agent_identity_jwt( - &agent_identity_jwt, - Some(&config.chatgpt_base_url), - /*agent_identity_authapi_base_url_override*/ None, - ) - .await?; + let auth = + CodexAuth::from_agent_identity_jwt(&agent_identity_jwt, Some(&config.chatgpt_base_url)) + .await?; return Ok(codex_model_provider::auth_provider_from_auth(&auth)); } diff --git a/codex-rs/cloud-config/Cargo.toml b/codex-rs/cloud-config/Cargo.toml index 6bf58c839..363cb21b6 100644 --- a/codex-rs/cloud-config/Cargo.toml +++ b/codex-rs/cloud-config/Cargo.toml @@ -25,6 +25,7 @@ tokio = { workspace = true, features = ["fs", "rt", "sync", "time"] } tracing = { workspace = true } [dev-dependencies] +codex-agent-identity = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "test-util", "time"] } diff --git a/codex-rs/cloud-config/src/bundle_loader.rs b/codex-rs/cloud-config/src/bundle_loader.rs index 2d76d29a6..5ec1ab260 100644 --- a/codex-rs/cloud-config/src/bundle_loader.rs +++ b/codex-rs/cloud-config/src/bundle_loader.rs @@ -63,6 +63,7 @@ pub async fn cloud_config_bundle_loader_for_storage( codex_home.clone(), enable_codex_api_key_env, credentials_store_mode, + /*forced_chatgpt_workspace_id*/ None, Some(chatgpt_base_url.clone()), keyring_backend_kind, ) diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index d63fef4c5..96fed2f58 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -17,6 +17,8 @@ use codex_config::CloudRequirementsFragment; use codex_config::CloudRequirementsTomlBundle; use codex_config::types::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; +use codex_login::auth::AgentIdentityAuth; +use codex_login::auth::AgentIdentityAuthRecord; use pretty_assertions::assert_eq; use serde_json::json; use std::collections::VecDeque; @@ -48,6 +50,7 @@ async fn auth_manager_with_api_key() -> Arc { tmp.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -77,6 +80,7 @@ async fn auth_manager_with_plan_and_identity( tmp.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -88,6 +92,28 @@ async fn auth_manager_with_plan(plan_type: &str) -> Arc { auth_manager_with_plan_and_identity(plan_type, Some("user-12345"), Some("account-12345")).await } +async fn auth_manager_with_agent_identity_business_plan() -> Arc { + let key_material = + codex_agent_identity::generate_agent_key_material().expect("generate agent key material"); + AuthManager::from_auth_for_testing(CodexAuth::AgentIdentity( + AgentIdentityAuth::from_record( + AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-123".to_string(), + agent_private_key: key_material.private_key_pkcs8_base64, + account_id: "account-12345".to_string(), + chatgpt_user_id: "user-12345".to_string(), + email: "user@example.com".to_string(), + plan_type: PlanType::Business, + chatgpt_account_is_fedramp: false, + task_id: Some("task-123".to_string()), + }, + "https://auth.openai.com/api/accounts", + ) + .await + .expect("agent identity record should be complete"), + )) +} + fn chatgpt_auth_json( plan_type: &str, chatgpt_user_id: Option<&str>, @@ -408,6 +434,28 @@ async fn get_bundle_allows_eligible_workspace_plans_and_writes_cache() { } } +#[tokio::test] +async fn get_bundle_allows_agent_identity_business_plan() { + let bundle = test_bundle(); + let fetcher = Arc::new(StaticBundleClient::new(bundle.clone())); + let codex_home = tempdir().expect("tempdir"); + let service = CloudConfigBundleService::new( + auth_manager_with_agent_identity_business_plan().await, + fetcher.clone(), + codex_home.path().to_path_buf(), + CLOUD_CONFIG_BUNDLE_TIMEOUT, + ); + + assert_eq!(service.load_startup_bundle().await, Ok(Some(bundle))); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 1); + assert!( + codex_home + .path() + .join(CLOUD_CONFIG_BUNDLE_CACHE_FILENAME) + .exists() + ); +} + #[tokio::test] async fn get_bundle_skips_team_like_usage_based_plan() { let fetcher = Arc::new(StaticBundleClient::new(test_bundle())); @@ -635,6 +683,7 @@ async fn get_bundle_recovers_after_unauthorized_reload() { auth_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -690,6 +739,7 @@ async fn get_bundle_recovers_after_unauthorized_reload_updates_cache_identity() auth_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -753,6 +803,7 @@ async fn get_bundle_surfaces_auth_recovery_message() { auth_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -818,6 +869,7 @@ async fn get_bundle_unauthorized_without_recovery_uses_generic_message() { auth_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index f683c58a4..65db955f1 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -49,6 +49,7 @@ pub async fn load_auth_manager(chatgpt_base_url: Option) -> Option 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 { diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 58f291407..56acb213e 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -228,6 +228,8 @@ pub enum Feature { PreventIdleSleep, /// Enable remote compaction v2 over the normal Responses API. RemoteCompactionV2, + /// Use Agent Identity for ChatGPT-authenticated sessions. + UseAgentIdentity, /// Enable workspace dependency support. WorkspaceDependencies, @@ -1314,6 +1316,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::UseAgentIdentity, + key: "use_agent_identity", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::WorkspaceDependencies, key: "workspace_dependencies", diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index d4c4f3e51..370af883c 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -1,83 +1,318 @@ +use std::future::Future; use std::sync::Arc; use codex_agent_identity::AgentIdentityKey; +use codex_agent_identity::ChatGptEnvironment; +use codex_agent_identity::build_abom; +use codex_agent_identity::decode_agent_identity_jwt; +use codex_agent_identity::fetch_agent_identity_jwks; +use codex_agent_identity::generate_agent_key_material; +use codex_agent_identity::is_retryable_registration_error; +use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64; +use codex_agent_identity::register_agent_identity; use codex_agent_identity::register_agent_task; use codex_protocol::account::PlanType as AccountPlanType; +use codex_protocol::protocol::SessionSource; +use thiserror::Error; use crate::default_client::build_reqwest_client; use super::storage::AgentIdentityAuthRecord; -#[derive(Clone, Debug)] -pub struct AgentIdentityAuth { - inner: Arc, +pub(super) const MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS: usize = 3; + +pub(super) fn agent_identity_authapi_base_url( + chatgpt_base_url: Option<&str>, +) -> std::io::Result { + 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()) } -#[derive(Debug)] -struct AgentIdentityAuthInner { - record: AgentIdentityAuthRecord, - run_task_id: String, +pub(super) fn require_agent_identity_authapi_base_url( + agent_identity_authapi_base_url: Option<&str>, +) -> std::io::Result<&str> { + agent_identity_authapi_base_url.ok_or_else(|| { + std::io::Error::other( + "Agent Identity only supports production and staging ChatGPT environments", + ) + }) +} + +#[derive(Debug, Error)] +pub enum AgentIdentityAuthError { + #[error( + "agent identity bootstrap unavailable after {attempts} attempts during {operation}: {message}" + )] + BootstrapUnavailable { + operation: &'static str, + attempts: usize, + message: String, + }, +} + +impl AgentIdentityAuthError { + pub fn is_bootstrap_unavailable(error: &std::io::Error) -> bool { + matches!( + error + .get_ref() + .and_then(|source| source.downcast_ref::()), + Some(Self::BootstrapUnavailable { .. }) + ) + } +} + +#[derive(Debug, Error)] +#[error("retryable agent identity registration failure: {message}")] +pub(super) struct RetryableAgentIdentityRegistrationError { + message: String, +} + +impl RetryableAgentIdentityRegistrationError { + pub(super) fn new(message: String) -> Self { + Self { message } + } +} + +#[derive(Clone, Debug)] +pub struct AgentIdentityAuth { + record: Arc, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct ManagedChatGptAgentIdentityBinding { + pub(super) account_id: String, + pub(super) chatgpt_user_id: String, + pub(super) email: String, + pub(super) plan_type: AccountPlanType, + pub(super) chatgpt_account_is_fedramp: bool, + pub(super) access_token: String, } impl AgentIdentityAuth { - pub async fn load( - record: AgentIdentityAuthRecord, + pub async fn from_record( + mut record: AgentIdentityAuthRecord, agent_identity_authapi_base_url: &str, ) -> std::io::Result { - let run_task_id = register_agent_task( - &build_reqwest_client(), - agent_identity_authapi_base_url, - key_for_record(&record), - ) - .await - .map_err(std::io::Error::other)?; + public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key) + .map_err(std::io::Error::other)?; + if record_needs_task_registration(&record) { + record.task_id = Some( + register_task_for_record_with_retries(&record, agent_identity_authapi_base_url) + .await?, + ); + } Ok(Self { - inner: Arc::new(AgentIdentityAuthInner { - record, - run_task_id, - }), + record: Arc::new(record), }) } + pub async fn from_jwt( + jwt: &str, + chatgpt_base_url: &str, + agent_identity_authapi_base_url: &str, + ) -> std::io::Result { + let record = verified_record_from_jwt(jwt, chatgpt_base_url).await?; + Self::from_record(record, agent_identity_authapi_base_url).await + } + #[cfg(test)] - fn from_initialized_record(record: AgentIdentityAuthRecord, run_task_id: String) -> Self { + fn from_initialized_record(mut record: AgentIdentityAuthRecord, run_task_id: String) -> Self { + record.task_id = Some(run_task_id); Self { - inner: Arc::new(AgentIdentityAuthInner { - record, - run_task_id, - }), + record: Arc::new(record), } } pub fn record(&self) -> &AgentIdentityAuthRecord { - &self.inner.record + self.record.as_ref() } pub fn run_task_id(&self) -> &str { - &self.inner.run_task_id + match self.record.task_id.as_deref() { + Some(task_id) => task_id, + None => unreachable!("AgentIdentityAuth should only be constructed with a task_id"), + } } pub fn account_id(&self) -> &str { - &self.inner.record.account_id + &self.record.account_id } pub fn chatgpt_user_id(&self) -> &str { - &self.inner.record.chatgpt_user_id + &self.record.chatgpt_user_id } pub fn email(&self) -> &str { - &self.inner.record.email + &self.record.email } pub fn plan_type(&self) -> AccountPlanType { - self.inner.record.plan_type + self.record.plan_type } pub fn is_fedramp_account(&self) -> bool { - self.inner.record.chatgpt_account_is_fedramp + self.record.chatgpt_account_is_fedramp } } +pub(super) async fn register_managed_chatgpt_agent_identity( + binding: ManagedChatGptAgentIdentityBinding, + agent_identity_authapi_base_url: &str, + session_source: SessionSource, +) -> std::io::Result { + let key_material = generate_agent_key_material().map_err(std::io::Error::other)?; + let client = build_reqwest_client(); + let runtime_id = retry_registration(|| async { + register_agent_identity( + &client, + agent_identity_authapi_base_url, + &binding.access_token, + binding.chatgpt_account_is_fedramp, + &key_material, + build_abom(session_source.clone()), + vec!["responsesapi".to_string()], + ) + .await + .map_err(|err| { + if is_retryable_registration_error(&err) { + std::io::Error::other(RetryableAgentIdentityRegistrationError::new( + err.to_string(), + )) + } else { + std::io::Error::other(err) + } + }) + }) + .await + .map_err(|err| classify_bootstrap_error("agent identity registration", err))?; + + let record = AgentIdentityAuthRecord { + agent_runtime_id: runtime_id, + agent_private_key: key_material.private_key_pkcs8_base64, + account_id: binding.account_id, + chatgpt_user_id: binding.chatgpt_user_id, + email: binding.email, + plan_type: binding.plan_type, + chatgpt_account_is_fedramp: binding.chatgpt_account_is_fedramp, + task_id: None, + }; + AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url) + .await + .map_err(|err| classify_bootstrap_error("agent task registration", err)) +} + +pub(super) async fn verified_record_from_jwt( + jwt: &str, + chatgpt_base_url: &str, +) -> std::io::Result { + AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; + let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url) + .await + .map_err(std::io::Error::other)?; + let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?; + Ok(claims.into()) +} + +pub(super) fn record_needs_task_registration(record: &AgentIdentityAuthRecord) -> bool { + record + .task_id + .as_deref() + .is_none_or(|task_id| task_id.trim().is_empty()) +} + +pub(super) fn record_matches_managed_chatgpt_binding( + record: &AgentIdentityAuthRecord, + binding: &ManagedChatGptAgentIdentityBinding, +) -> bool { + record.account_id == binding.account_id + && record.chatgpt_user_id == binding.chatgpt_user_id + && public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key).is_ok() +} + +pub(super) fn classify_bootstrap_error( + operation: &'static str, + err: std::io::Error, +) -> std::io::Error { + if is_retryable_io_registration_error(&err) { + std::io::Error::other(AgentIdentityAuthError::BootstrapUnavailable { + operation, + attempts: MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS, + message: err.to_string(), + }) + } else { + err + } +} + +pub(super) fn is_retryable_io_registration_error(err: &std::io::Error) -> bool { + err.get_ref().is_some_and( + ::is::< + RetryableAgentIdentityRegistrationError, + >, + ) +} + +pub(super) async fn retry_registration(mut operation: F) -> std::io::Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 1; + loop { + match operation().await { + Ok(value) => return Ok(value), + Err(err) + if attempt < MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS + && is_retryable_io_registration_error(&err) => + { + tracing::warn!( + attempt, + max_attempts = MAX_AGENT_IDENTITY_BOOTSTRAP_ATTEMPTS, + error = %err, + "agent identity registration attempt failed; retrying" + ); + attempt += 1; + } + Err(err) => return Err(err), + } + } +} + +async fn register_task_for_record_with_retries( + record: &AgentIdentityAuthRecord, + agent_identity_authapi_base_url: &str, +) -> std::io::Result { + retry_registration(|| async { + register_task_for_record(record, agent_identity_authapi_base_url).await + }) + .await +} + +async fn register_task_for_record( + record: &AgentIdentityAuthRecord, + agent_identity_authapi_base_url: &str, +) -> std::io::Result { + register_agent_task( + &build_reqwest_client(), + agent_identity_authapi_base_url, + key_for_record(record), + ) + .await + .map_err(|err| { + if is_retryable_registration_error(&err) { + std::io::Error::other(RetryableAgentIdentityRegistrationError::new( + err.to_string(), + )) + } else { + std::io::Error::other(err) + } + }) +} + fn key_for_record(record: &AgentIdentityAuthRecord) -> AgentIdentityKey<'_> { AgentIdentityKey { agent_runtime_id: &record.agent_runtime_id, @@ -111,6 +346,7 @@ mod tests { email: "agent@example.com".to_string(), plan_type: AccountPlanType::Plus, chatgpt_account_is_fedramp: false, + task_id: None, } } @@ -120,7 +356,7 @@ mod tests { } #[tokio::test] - async fn load_registers_run_task() -> anyhow::Result<()> { + async fn from_record_registers_task() -> anyhow::Result<()> { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/agent/agent-runtime-1/task/register")) @@ -131,9 +367,11 @@ mod tests { .mount(&server) .await; - let auth = - AgentIdentityAuth::load(agent_identity_record_with_generated_key(), &server.uri()) - .await?; + let auth = AgentIdentityAuth::from_record( + agent_identity_record_with_generated_key(), + &server.uri(), + ) + .await?; assert_eq!(auth.run_task_id(), "task-run-1"); let requests = server @@ -152,6 +390,38 @@ mod tests { Ok(()) } + #[tokio::test] + async fn from_jwt_registers_task() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/agent-identities/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) + .expect(1) + .mount(&server) + .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 record = agent_identity_record_with_generated_key(); + let jwt = signed_agent_identity_jwt(&record)?; + let auth = AgentIdentityAuth::from_jwt( + &jwt, + &format!("{}/backend-api", server.uri()), + &server.uri(), + ) + .await?; + + assert_eq!(auth.record().agent_runtime_id, "agent-runtime-1"); + assert_eq!(auth.run_task_id(), "task-run-1"); + Ok(()) + } + #[test] fn run_task_is_shared_across_clones() { let auth = AgentIdentityAuth::from_initialized_record( @@ -160,12 +430,12 @@ mod tests { ); let cloned = auth.clone(); - assert!(Arc::ptr_eq(&auth.inner, &cloned.inner)); + assert!(Arc::ptr_eq(&auth.record, &cloned.record)); 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<()> { + async fn from_record_retries_transient_registration() -> anyhow::Result<()> { let server = MockServer::start().await; let request_count = Arc::new(AtomicUsize::new(0)); let response_count = Arc::clone(&request_count); @@ -183,14 +453,80 @@ mod tests { .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?; + let auth = AgentIdentityAuth::from_record( + agent_identity_record_with_generated_key(), + &server.uri(), + ) + .await?; assert_eq!(request_count.load(Ordering::SeqCst), 2); assert_eq!(auth.run_task_id(), "task-run-1"); Ok(()) } + + fn signed_agent_identity_jwt( + record: &AgentIdentityAuthRecord, + ) -> jsonwebtoken::errors::Result { + let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some("test-key".to_string()); + jsonwebtoken::encode( + &header, + &json!({ + "iss": "https://chatgpt.com/codex-backend/agent-identity", + "aud": "codex-app-server", + "iat": 1_700_000_000usize, + "exp": 4_000_000_000usize, + "agent_runtime_id": record.agent_runtime_id, + "agent_private_key": record.agent_private_key, + "account_id": record.account_id, + "chatgpt_user_id": record.chatgpt_user_id, + "email": record.email, + "plan_type": record.plan_type, + "chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp, + }), + &jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?, + ) + } + + fn test_jwks_body() -> serde_json::Value { + json!({ + "keys": [{ + "kty": "RSA", + "kid": "test-key", + "use": "sig", + "alg": "RS256", + "n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q", + "e": "AQAB", + }] + }) + } + + const TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM: &[u8] = br#"-----BEGIN PRIVATE KEY----- +MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO +bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9 +FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9 +mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg +0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs +zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa +uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2 +1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu +RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8 +NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ +rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj +Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL +aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL +iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN +YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL +3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq +sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg +gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP +sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5 +BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD +KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c +r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3 +29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc +J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN +5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR +8U4M2TSWCKUY/A6sT4W8+mT9 +-----END PRIVATE KEY-----"#; } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index aa1dcd9c8..b52f03b70 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -6,6 +6,7 @@ use codex_app_server_protocol::AuthMode; use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; +use codex_protocol::protocol::SessionSource; use base64::Engine; use codex_protocol::config_types::ForcedLoginMethod; @@ -14,11 +15,14 @@ use pretty_assertions::assert_eq; use serde::Serialize; use serde_json::json; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tempfile::TempDir; use tempfile::tempdir; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -95,7 +99,7 @@ fn login_with_api_key_overwrites_existing_auth_json() { } #[tokio::test] -async fn login_with_access_token_writes_only_token() { +async fn login_with_access_token_writes_agent_identity_jwt() { let dir = tempdir().unwrap(); let auth_path = dir.path().join("auth.json"); let record = agent_identity_record(WORKSPACE_ID_ALLOWED); @@ -108,7 +112,8 @@ async fn login_with_access_token_writes_only_token() { .expect(1) .mount(&server) .await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); + let authapi_base_url = server.uri(); + let chatgpt_base_url = format!("{authapi_base_url}/backend-api"); super::login_with_access_token( dir.path(), @@ -127,14 +132,75 @@ async fn login_with_access_token_writes_only_token() { .expect("auth.json should parse"); assert_eq!(auth.auth_mode, Some(AuthMode::AgentIdentity)); assert_eq!( - auth.agent_identity.as_deref(), - Some(agent_identity.as_str()) + auth.agent_identity, + Some(AgentIdentityStorage::Jwt(agent_identity)) ); assert!(auth.tokens.is_none(), "tokens should be cleared"); assert!(auth.openai_api_key.is_none(), "API key should be cleared"); server.verify().await; } +#[tokio::test] +#[serial(codex_auth_env)] +async fn stored_agent_identity_jwt_keeps_auth_json_unchanged() -> anyhow::Result<()> { + let _access_token_guard = remove_access_token_env_var(); + let codex_home = tempdir()?; + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/agent-identities/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) + .expect(1) + .mount(&server) + .await; + mock_agent_task_registration(&server, "", &record.agent_runtime_id, "task-id").await; + let authapi_base_url = server.uri(); + let chatgpt_base_url = format!("{authapi_base_url}/backend-api"); + save_auth( + codex_home.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::AgentIdentity), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity.clone())), + personal_access_token: None, + bedrock_api_key: None, + }, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + Some(&chatgpt_base_url), + AuthKeyringBackendKind::Direct, + Some(&authapi_base_url), + ) + .await? + .expect("auth should load"); + + let CodexAuth::AgentIdentity(agent_identity_auth) = auth else { + panic!("stored JWT should load as agent identity auth"); + }; + assert_eq!(agent_identity_auth.run_task_id(), "task-id"); + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth = storage + .try_read_auth_json(&get_auth_file(codex_home.path())) + .expect("auth.json should parse"); + assert_eq!( + auth.agent_identity, + Some(AgentIdentityStorage::Jwt(agent_identity)) + ); + server.verify().await; + Ok(()) +} + #[tokio::test] #[serial(codex_auth_env)] async fn login_with_access_token_writes_only_personal_access_token() { @@ -278,6 +344,387 @@ async fn login_with_access_token_rejects_invalid_jwt() { ); } +#[tokio::test] +#[serial(codex_auth_env)] +async fn chatgpt_auth_registers_agent_identity_when_enabled() -> anyhow::Result<()> { + let codex_home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + )?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should load"); + + assert!( + auth.agent_identity_auth( + AgentIdentityAuthPolicy::JwtOnly, + /*agent_identity_authapi_base_url*/ None, + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .is_none() + ); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("authorization", "Bearer test-access-token")) + .and(body_partial_json(json!({ + "abom": { + "agent_harness_id": "codex-cli", + }, + "capabilities": ["responsesapi"], + "ttl": null, + }))) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "agent_runtime_id": "agent-runtime-123", + }))) + .expect(/*r*/ 1) + .mount(&server) + .await; + mock_agent_task_registration(&server, "", "agent-runtime-123", "task-123").await; + + let agent_auth = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should register"); + let reused = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should be reused"); + + assert_eq!( + agent_auth.record().agent_runtime_id, + reused.record().agent_runtime_id + ); + assert_eq!(agent_auth.run_task_id(), "task-123"); + assert_eq!(reused.run_task_id(), "task-123"); + assert_eq!(agent_auth.record().agent_runtime_id, "agent-runtime-123"); + assert_eq!(agent_auth.record().account_id, "account-123"); + assert_eq!(agent_auth.record().chatgpt_user_id, "user-12345"); + assert_eq!(agent_auth.record().task_id.as_deref(), Some("task-123")); + assert_eq!(reused.record().task_id.as_deref(), Some("task-123")); + let persisted = auth + .stored_managed_chatgpt_agent_identity_record("account-123") + .expect("identity should persist"); + assert_eq!(persisted.agent_runtime_id, "agent-runtime-123"); + assert_eq!(persisted.task_id.as_deref(), Some("task-123")); + + let reloaded = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should reload"); + let reloaded_agent_auth = reloaded + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should reload from storage"); + assert_eq!( + reloaded_agent_auth.record().agent_runtime_id, + "agent-runtime-123" + ); + assert_eq!(reloaded_agent_auth.run_task_id(), "task-123"); + Ok(()) +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn chatgpt_auth_retries_transient_agent_identity_registration() -> anyhow::Result<()> { + let codex_home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + )?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should load"); + + let server = MockServer::start().await; + let registration_count = Arc::new(AtomicUsize::new(0)); + let response_count = Arc::clone(®istration_count); + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(move |_request: &wiremock::Request| { + if response_count.fetch_add(1, Ordering::SeqCst) < 2 { + ResponseTemplate::new(/*status*/ 503) + } else { + ResponseTemplate::new(/*status*/ 200).set_body_json(json!({ + "agent_runtime_id": "agent-runtime-123", + })) + } + }) + .expect(/*requests*/ 3) + .mount(&server) + .await; + mock_agent_task_registration(&server, "", "agent-runtime-123", "task-123").await; + + let agent_auth = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await? + .expect("agent identity should register after retries"); + + assert_eq!(registration_count.load(Ordering::SeqCst), 3); + assert_eq!(agent_auth.record().agent_runtime_id, "agent-runtime-123"); + assert_eq!(agent_auth.record().task_id.as_deref(), Some("task-123")); + assert_eq!( + auth.stored_managed_chatgpt_agent_identity_record("account-123") + .and_then(|record| record.task_id), + Some("task-123".to_string()) + ); + Ok(()) +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn chatgpt_auth_registration_retry_exhaustion_is_fallback_eligible() -> anyhow::Result<()> { + let codex_home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + )?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should load"); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(ResponseTemplate::new(/*status*/ 503)) + .expect(/*requests*/ 3) + .mount(&server) + .await; + + let err = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await + .expect_err("retry exhaustion should return an error"); + + assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!( + auth.stored_managed_chatgpt_agent_identity_record("account-123") + .is_none() + ); + Ok(()) +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn chatgpt_auth_task_registration_retry_exhaustion_is_fallback_eligible() -> anyhow::Result<()> +{ + let codex_home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + )?; + let mut record = agent_identity_record("account-123"); + record.chatgpt_user_id = "user-12345".to_string(); + record.email = "user@example.com".to_string(); + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth_path = get_auth_file(codex_home.path()); + let mut auth_json = storage.try_read_auth_json(&auth_path)?; + auth_json.agent_identity = Some(AgentIdentityStorage::Record(record.clone())); + storage.save(&auth_json)?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should load"); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1/agent/{}/task/register", + record.agent_runtime_id + ))) + .respond_with(ResponseTemplate::new(/*status*/ 503)) + .expect(/*requests*/ 3) + .mount(&server) + .await; + + let err = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await + .expect_err("task retry exhaustion should return an error"); + + assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + record.task_id = None; + assert_eq!( + auth.stored_managed_chatgpt_agent_identity_record("account-123"), + Some(record) + ); + Ok(()) +} + +#[tokio::test] +#[serial(codex_auth_env)] +async fn chatgpt_auth_non_retryable_registration_error_is_hard_failure() -> anyhow::Result<()> { + let codex_home = tempdir()?; + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + )?; + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + /*agent_identity_authapi_base_url*/ None, + ) + .await? + .expect("auth should load"); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(ResponseTemplate::new(/*status*/ 403)) + .expect(/*requests*/ 1) + .mount(&server) + .await; + + let err = auth + .agent_identity_auth( + AgentIdentityAuthPolicy::ChatGptAuth, + Some(&server.uri()), + /*forced_chatgpt_workspace_id*/ None, + SessionSource::Cli, + ) + .await + .expect_err("hard registration failure should return an error"); + + assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!( + auth.stored_managed_chatgpt_agent_identity_record("account-123") + .is_none() + ); + Ok(()) +} + +#[tokio::test] +async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> anyhow::Result<()> { + let record = agent_identity_record(WORKSPACE_ID_ALLOWED); + let agent_identity = + signed_agent_identity_jwt(&record, json!(record.plan_type)).expect("signed agent identity"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/agent-identities/jwks")) + .respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body())) + .expect(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path(format!( + "/v1/agent/{}/task/register", + record.agent_runtime_id + ))) + .respond_with(ResponseTemplate::new(/*status*/ 503)) + .expect(/*requests*/ 3) + .mount(&server) + .await; + let authapi_base_url = server.uri(); + let chatgpt_base_url = format!("{authapi_base_url}/backend-api"); + + let err = CodexAuth::from_agent_identity_jwt_with_authapi_base_url( + &agent_identity, + Some(&chatgpt_base_url), + &authapi_base_url, + ) + .await + .expect_err("agent identity jwt task retry exhaustion should fail"); + + assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + Ok(()) +} + #[tokio::test] async fn login_with_access_token_rejects_unsigned_jwt() { let dir = tempdir().unwrap(); @@ -290,7 +737,8 @@ async fn login_with_access_token_rejects_unsigned_jwt() { .expect(1) .mount(&server) .await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); + let authapi_base_url = server.uri(); + let chatgpt_base_url = format!("{authapi_base_url}/backend-api"); super::login_with_access_token( dir.path(), @@ -457,6 +905,7 @@ async fn unauthorized_recovery_reports_mode_and_step_names() { dir.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -828,7 +1277,6 @@ async fn build_config( forced_login_method, forced_chatgpt_workspace_id, chatgpt_base_url: None, - agent_identity_authapi_base_url: None, } } @@ -879,7 +1327,7 @@ fn remove_access_token_env_var() -> EnvVarGuard { #[serial(codex_auth_env)] async fn load_auth_reads_access_token_from_env() { let codex_home = tempdir().unwrap(); - let expected_record = agent_identity_record(WORKSPACE_ID_ALLOWED); + let mut expected_record = agent_identity_record(WORKSPACE_ID_ALLOWED); let agent_identity = signed_agent_identity_jwt(&expected_record, json!(expected_record.plan_type)) .expect("signed agent identity"); @@ -898,6 +1346,7 @@ async fn load_auth_reads_access_token_from_env() { .expect(1) .mount(&server) .await; + expected_record.task_id = Some("task-123".to_string()); let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, &agent_identity); let authapi_base_url = server.uri(); @@ -1003,11 +1452,10 @@ async fn auth_manager_rejects_env_personal_access_token_workspace_mismatch() { let _access_token_guard = EnvVarGuard::set(CODEX_ACCESS_TOKEN_ENV_VAR, "at-env-workspace-mismatch"); - let manager = AuthManager::new_with_workspace_restriction( + let manager = AuthManager::new( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, - /*forced_chatgpt_workspace_id*/ Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), @@ -1054,11 +1502,10 @@ async fn auth_manager_rejects_stored_personal_access_token_workspace_mismatch() .await .expect("personal access token login should succeed"); - let manager = AuthManager::new_with_workspace_restriction( + let manager = AuthManager::new( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, - /*forced_chatgpt_workspace_id*/ Some(vec![WORKSPACE_ID_ALLOWED.to_string()]), /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), @@ -1092,6 +1539,7 @@ async fn personal_access_token_does_not_offer_unauthorized_recovery() { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -1233,7 +1681,6 @@ 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) @@ -1341,7 +1788,7 @@ async fn enforce_login_restrictions_logs_out_for_agent_identity_workspace_mismat openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(agent_identity), + agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)), personal_access_token: None, bedrock_api_key: None, }, @@ -1357,15 +1804,21 @@ 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) - .await - .expect_err("expected workspace mismatch to error"); - assert!(err.to_string().contains(&format!( - "current credentials belong to {WORKSPACE_ID_DISALLOWED}" - ))); + let err = super::enforce_login_restrictions_with_agent_identity_authapi_base_url( + &config, + Some(&authapi_base_url), + ) + .await + .expect_err("expected workspace mismatch to error"); + let message = err.to_string(); + assert!( + message.contains(&format!( + "current credentials belong to {WORKSPACE_ID_DISALLOWED}" + )), + "{message}" + ); assert!( !codex_home.path().join("auth.json").exists(), "auth.json should be removed on mismatch" @@ -1437,9 +1890,28 @@ fn agent_identity_record(account_id: &str) -> AgentIdentityAuthRecord { email: "user@example.com".to_string(), plan_type: AccountPlanType::Pro, chatgpt_account_is_fedramp: false, + task_id: None, } } +async fn mock_agent_task_registration( + server: &MockServer, + path_prefix: &str, + agent_runtime_id: &str, + task_id: &str, +) { + Mock::given(method("POST")) + .and(path(format!( + "{path_prefix}/v1/agent/{agent_runtime_id}/task/register" + ))) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({ + "task_id": task_id, + }))) + .expect(/*r*/ 1) + .mount(server) + .await; +} + fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result { fake_agent_identity_jwt_with_plan_type(record, serde_json::to_value(record.plan_type)?) } @@ -1580,10 +2052,13 @@ async fn assert_agent_identity_plan_alias( .await; 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"); + let auth = CodexAuth::from_agent_identity_jwt_with_authapi_base_url( + &jwt, + Some(&chatgpt_base_url), + &authapi_base_url, + ) + .await + .expect("agent identity auth"); pretty_assertions::assert_eq!(auth.account_plan_type(), Some(expected_plan_type)); server.verify().await; diff --git a/codex-rs/login/src/auth/bedrock_api_key_tests.rs b/codex-rs/login/src/auth/bedrock_api_key_tests.rs index 35ace9fea..691d8bbc5 100644 --- a/codex-rs/login/src/auth/bedrock_api_key_tests.rs +++ b/codex-rs/login/src/auth/bedrock_api_key_tests.rs @@ -60,6 +60,7 @@ async fn login_with_bedrock_api_key_replaces_openai_auth() -> anyhow::Result<()> codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -107,6 +108,7 @@ async fn logout_removes_bedrock_auth() -> anyhow::Result<()> { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) @@ -130,6 +132,7 @@ async fn bedrock_only_auth_storage_creates_primary_auth() -> anyhow::Result<()> codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index bb548ffe7..138bdad5b 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -20,8 +20,6 @@ use tokio::sync::watch; use tracing::instrument; 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; use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -29,18 +27,27 @@ use codex_protocol::config_types::ModelProviderAuthInfo; use super::access_token::CodexAccessToken; use super::access_token::classify_codex_access_token; +use super::agent_identity::ManagedChatGptAgentIdentityBinding; +use super::agent_identity::agent_identity_authapi_base_url; +use super::agent_identity::classify_bootstrap_error; +use super::agent_identity::record_matches_managed_chatgpt_binding; +use super::agent_identity::record_needs_task_registration; +use super::agent_identity::register_managed_chatgpt_agent_identity; +use super::agent_identity::require_agent_identity_authapi_base_url; +use super::agent_identity::verified_record_from_jwt; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::agent_identity::AgentIdentityAuth; +pub use crate::auth::agent_identity::AgentIdentityAuthError; pub use crate::auth::bedrock_api_key::BedrockApiKeyAuth; pub use crate::auth::personal_access_token::PersonalAccessTokenAuth; pub use crate::auth::storage::AgentIdentityAuthRecord; +pub use crate::auth::storage::AgentIdentityStorage; pub use crate::auth::storage::AuthDotJson; pub use crate::auth::storage::AuthKeyringBackendKind; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::create_auth_storage; use crate::auth::util::try_parse_error_message; -use crate::default_client::build_reqwest_client; use crate::default_client::create_client; use crate::token_data::TokenData; use crate::token_data::parse_chatgpt_jwt_claims; @@ -51,6 +58,7 @@ use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; +use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; @@ -65,6 +73,15 @@ pub enum CodexAuth { BedrockApiKey(BedrockApiKeyAuth), } +/// Policy for resolving Agent Identity auth from a broader Codex auth snapshot. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentIdentityAuthPolicy { + /// Use Agent Identity auth only when the current auth is already Agent Identity. + JwtOnly, + /// Allow managed ChatGPT auth to register or reuse Agent Identity auth. + ChatGptAuth, +} + impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { match (self, other) { @@ -222,7 +239,6 @@ impl CodexAuth { agent_identity_authapi_base_url: Option<&str>, ) -> std::io::Result { let auth_mode = auth_dot_json.resolved_mode(); - let client = create_client(); if auth_mode == ApiAuthMode::ApiKey { let Some(api_key) = auth_dot_json.openai_api_key.as_deref() else { return Err(std::io::Error::other("API key auth is missing a key.")); @@ -230,17 +246,34 @@ impl CodexAuth { return Ok(Self::from_api_key(api_key)); } if auth_mode == ApiAuthMode::AgentIdentity { - let Some(agent_identity) = auth_dot_json.agent_identity else { + let Some(agent_identity) = auth_dot_json.agent_identity.clone() else { return Err(std::io::Error::other( - "agent identity auth is missing an agent identity token.", + "agent identity auth is missing agent identity auth material.", )); }; - return Self::from_agent_identity_jwt( - &agent_identity, - chatgpt_base_url, - agent_identity_authapi_base_url, - ) - .await; + let base_url = chatgpt_base_url + .unwrap_or(ChatGptEnvironment::default().chatgpt_base_url()) + .trim_end_matches('/') + .to_string(); + let agent_identity_authapi_base_url = + require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?; + match agent_identity { + AgentIdentityStorage::Jwt(jwt) => { + let auth = AgentIdentityAuth::from_jwt( + &jwt, + &base_url, + agent_identity_authapi_base_url, + ) + .await?; + return Ok(Self::AgentIdentity(auth)); + } + AgentIdentityStorage::Record(record) => { + let auth = + AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url) + .await?; + return Ok(Self::AgentIdentity(auth)); + } + } } if auth_mode == ApiAuthMode::PersonalAccessToken { let Some(personal_access_token) = auth_dot_json.personal_access_token.as_deref() else { @@ -262,7 +295,7 @@ impl CodexAuth { let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); let state = ChatgptAuthState { auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, + client: create_client(), }; match auth_mode { @@ -292,6 +325,8 @@ impl CodexAuth { chatgpt_base_url: Option<&str>, keyring_backend_kind: AuthKeyringBackendKind, ) -> std::io::Result> { + let agent_identity_authapi_base_url = + agent_identity_authapi_base_url(chatgpt_base_url).ok(); load_auth( codex_home, /*enable_codex_api_key_env*/ false, @@ -299,7 +334,7 @@ impl CodexAuth { /*forced_chatgpt_workspace_id*/ None, chatgpt_base_url, keyring_backend_kind, - /*agent_identity_authapi_base_url*/ None, + agent_identity_authapi_base_url.as_deref(), ) .await } @@ -307,19 +342,28 @@ 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 { + let agent_identity_authapi_base_url = agent_identity_authapi_base_url(chatgpt_base_url)?; + Self::from_agent_identity_jwt_with_authapi_base_url( + jwt, + chatgpt_base_url, + &agent_identity_authapi_base_url, + ) + .await + } + + async fn from_agent_identity_jwt_with_authapi_base_url( + jwt: &str, + chatgpt_base_url: Option<&str>, + agent_identity_authapi_base_url: &str, ) -> std::io::Result { let base_url = chatgpt_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?; - let auth = AgentIdentityAuth::load(record, &resolved_authapi_base_url).await?; - Ok(Self::AgentIdentity(auth)) + Ok(Self::AgentIdentity( + AgentIdentityAuth::from_jwt(jwt, &base_url, agent_identity_authapi_base_url).await?, + )) } pub async fn from_personal_access_token(access_token: &str) -> std::io::Result { @@ -499,6 +543,89 @@ impl CodexAuth { self.get_current_auth_json().and_then(|t| t.tokens) } + fn stored_managed_chatgpt_agent_identity_record( + &self, + account_id: &str, + ) -> Option { + self.get_current_auth_json() + .and_then(|auth| auth.agent_identity) + .and_then(|identity| identity.as_record().cloned()) + .filter(|identity| identity.account_id == account_id) + } + + fn persist_managed_chatgpt_agent_identity_record( + &self, + record: AgentIdentityAuthRecord, + ) -> std::io::Result<()> { + if let Self::Chatgpt(chatgpt_auth) = self { + chatgpt_auth.persist_agent_identity_record(record)?; + } + Ok(()) + } + + async fn agent_identity_auth( + &self, + policy: AgentIdentityAuthPolicy, + agent_identity_authapi_base_url: Option<&str>, + forced_chatgpt_workspace_id: Option>, + session_source: SessionSource, + ) -> std::io::Result> { + match self { + Self::AgentIdentity(auth) => Ok(Some(auth.clone())), + Self::ApiKey(_) + | Self::ChatgptAuthTokens(_) + | Self::PersonalAccessToken(_) + | Self::BedrockApiKey(_) => Ok(None), + Self::Chatgpt(_) => { + if policy == AgentIdentityAuthPolicy::JwtOnly { + return Ok(None); + } + self.ensure_managed_chatgpt_agent_identity( + require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?, + forced_chatgpt_workspace_id, + session_source, + ) + .await + .map(Some) + } + } + } + + async fn ensure_managed_chatgpt_agent_identity( + &self, + agent_identity_authapi_base_url: &str, + forced_chatgpt_workspace_id: Option>, + session_source: SessionSource, + ) -> std::io::Result { + let binding = + ManagedChatGptAgentIdentityBinding::from_auth(self, forced_chatgpt_workspace_id) + .ok_or_else(|| std::io::Error::other("ChatGPT auth is unavailable"))?; + + // JWT auth is loaded as CodexAuth::AgentIdentity; this path only reuses + // records created by the managed ChatGPT Agent Identity bootstrap. + if let Some(record) = self.stored_managed_chatgpt_agent_identity_record(&binding.account_id) + && record_matches_managed_chatgpt_binding(&record, &binding) + { + let should_persist = record_needs_task_registration(&record); + let auth = AgentIdentityAuth::from_record(record, agent_identity_authapi_base_url) + .await + .map_err(|err| classify_bootstrap_error("agent task registration", err))?; + if should_persist { + self.persist_managed_chatgpt_agent_identity_record(auth.record().clone())?; + } + return Ok(auth); + } + + let auth = register_managed_chatgpt_agent_identity( + binding, + agent_identity_authapi_base_url, + session_source, + ) + .await?; + self.persist_managed_chatgpt_agent_identity_record(auth.record().clone())?; + Ok(auth) + } + /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { @@ -516,10 +643,9 @@ impl CodexAuth { bedrock_api_key: None, }; - let client = create_client(); let state = ChatgptAuthState { auth_dot_json: Arc::new(Mutex::new(Some(auth_dot_json))), - client, + client: create_client(), }; let dummy_auth_id = NEXT_DUMMY_AUTH_ID.fetch_add(1, Ordering::Relaxed); let storage = create_auth_storage( @@ -537,6 +663,43 @@ impl CodexAuth { } } +impl ManagedChatGptAgentIdentityBinding { + fn from_auth(auth: &CodexAuth, forced_workspace_id: Option>) -> Option { + if !auth.is_chatgpt_auth() { + return None; + } + + let token_data = auth.get_token_data().ok()?; + let forced_workspace_id = + forced_workspace_id + .as_deref() + .and_then(|workspace_ids| match workspace_ids { + [workspace_id] if !workspace_id.is_empty() => Some(workspace_id.clone()), + _ => None, + }); + let account_id = forced_workspace_id + .or(token_data + .account_id + .clone() + .filter(|value| !value.is_empty())) + .or(token_data.id_token.chatgpt_account_id.clone())?; + let chatgpt_user_id = token_data + .id_token + .chatgpt_user_id + .clone() + .filter(|value| !value.is_empty())?; + + Some(Self { + account_id, + chatgpt_user_id, + email: token_data.id_token.email.clone().unwrap_or_default(), + plan_type: auth.account_plan_type().unwrap_or(AccountPlanType::Unknown), + chatgpt_account_is_fedramp: auth.is_fedramp_account(), + access_token: token_data.access_token, + }) + } +} + impl ChatgptAuth { fn current_auth_json(&self) -> Option { #[expect(clippy::unwrap_used)] @@ -554,6 +717,31 @@ impl ChatgptAuth { fn client(&self) -> &CodexHttpClient { &self.state.client } + + fn persist_agent_identity_record( + &self, + record: AgentIdentityAuthRecord, + ) -> std::io::Result<()> { + persist_agent_identity_record(&self.state.auth_dot_json, &self.storage, record) + } +} + +fn persist_agent_identity_record( + auth_dot_json: &Arc>>, + storage: &Arc, + record: AgentIdentityAuthRecord, +) -> std::io::Result<()> { + let mut guard = auth_dot_json + .lock() + .map_err(|_| std::io::Error::other("failed to lock auth state"))?; + let mut auth = storage + .load()? + .or_else(|| guard.clone()) + .ok_or_else(|| std::io::Error::other("auth data is not available"))?; + auth.agent_identity = Some(AgentIdentityStorage::Record(record)); + storage.save(&auth)?; + *guard = Some(auth); + Ok(()) } pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; @@ -582,34 +770,6 @@ fn read_non_empty_env_var(key: &str) -> Option { .filter(|value| !value.is_empty()) } -async fn verified_agent_identity_record( - jwt: &str, - chatgpt_base_url: &str, -) -> std::io::Result { - AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?; - let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url) - .await - .map_err(std::io::Error::other)?; - let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?; - 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 { - 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( @@ -705,13 +865,13 @@ pub async fn login_with_access_token( .unwrap_or(ChatGptEnvironment::default().chatgpt_base_url()) .trim_end_matches('/') .to_string(); - verified_agent_identity_record(jwt, &base_url).await?; + verified_record_from_jwt(jwt, &base_url).await?; AuthDotJson { auth_mode: Some(ApiAuthMode::AgentIdentity), openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(jwt.to_string()), + agent_identity: Some(AgentIdentityStorage::Jwt(jwt.to_string())), personal_access_token: None, bedrock_api_key: None, } @@ -793,11 +953,23 @@ pub struct AuthConfig { pub keyring_backend_kind: AuthKeyringBackendKind, pub forced_login_method: Option, pub chatgpt_base_url: Option, - pub agent_identity_authapi_base_url: Option, pub forced_chatgpt_workspace_id: Option>, } pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { + let agent_identity_authapi_base_url = + agent_identity_authapi_base_url(config.chatgpt_base_url.as_deref()).ok(); + enforce_login_restrictions_with_agent_identity_authapi_base_url( + config, + agent_identity_authapi_base_url.as_deref(), + ) + .await +} + +async fn enforce_login_restrictions_with_agent_identity_authapi_base_url( + config: &AuthConfig, + agent_identity_authapi_base_url: Option<&str>, +) -> std::io::Result<()> { let Some(auth) = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, @@ -805,7 +977,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(), + agent_identity_authapi_base_url, ) .await? else { @@ -985,11 +1157,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, - agent_identity_authapi_base_url, - ) + CodexAccessToken::AgentIdentityJwt(jwt) => { + CodexAuth::from_agent_identity_jwt_with_authapi_base_url( + jwt, + chatgpt_base_url, + require_agent_identity_authapi_base_url(agent_identity_authapi_base_url)?, + ) + } .await .map(Some), }; @@ -1524,7 +1698,9 @@ pub struct AuthManager { keyring_backend_kind: AuthKeyringBackendKind, forced_chatgpt_workspace_id: RwLock>>, chatgpt_base_url: Option, + agent_identity_authapi_base_url: Option, refresh_lock: Semaphore, + agent_identity_lock: Semaphore, external_auth: RwLock>>, } @@ -1572,30 +1748,16 @@ impl Debug for AuthManager { } } +fn default_agent_identity_authapi_base_url() -> Option { + agent_identity_authapi_base_url(/*chatgpt_base_url*/ None).ok() +} + impl AuthManager { /// Create a new manager loading the initial auth using the provided /// preferred auth method. Errors loading auth are swallowed; `auth()` will /// simply return `None` in that case so callers can treat it as an /// unauthenticated state. pub async fn new( - codex_home: PathBuf, - enable_codex_api_key_env: bool, - auth_credentials_store_mode: AuthCredentialsStoreMode, - chatgpt_base_url: Option, - keyring_backend_kind: AuthKeyringBackendKind, - ) -> Self { - Self::new_with_workspace_restriction( - codex_home, - enable_codex_api_key_env, - auth_credentials_store_mode, - /*forced_chatgpt_workspace_id*/ None, - chatgpt_base_url, - keyring_backend_kind, - ) - .await - } - - async fn new_with_workspace_restriction( codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, @@ -1603,6 +1765,8 @@ impl AuthManager { chatgpt_base_url: Option, keyring_backend_kind: AuthKeyringBackendKind, ) -> Self { + let agent_identity_authapi_base_url = + agent_identity_authapi_base_url(chatgpt_base_url.as_deref()).ok(); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1610,7 +1774,7 @@ impl AuthManager { forced_chatgpt_workspace_id.as_deref(), chatgpt_base_url.as_deref(), keyring_backend_kind, - /*agent_identity_authapi_base_url*/ None, + agent_identity_authapi_base_url.as_deref(), ) .await .ok() @@ -1628,7 +1792,9 @@ impl AuthManager { keyring_backend_kind, forced_chatgpt_workspace_id: RwLock::new(forced_chatgpt_workspace_id), chatgpt_base_url, + agent_identity_authapi_base_url, refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), } } @@ -1650,7 +1816,9 @@ impl AuthManager { keyring_backend_kind: AuthKeyringBackendKind::default(), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, + agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), }) } @@ -1671,7 +1839,40 @@ impl AuthManager { keyring_backend_kind: AuthKeyringBackendKind::default(), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, + agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), + external_auth: RwLock::new(None), + }) + } + + /// Create an AuthManager with a specific CodexAuth and Agent Identity AuthAPI base URL, for testing only. + #[doc(hidden)] + pub fn from_auth_for_testing_with_agent_identity_authapi_base_url( + auth: CodexAuth, + agent_identity_authapi_base_url: String, + ) -> Arc { + let cached = CachedAuth { + auth: Some(auth), + permanent_refresh_failure: None, + }; + let (auth_change_tx, _auth_change_rx) = watch::channel(0); + Arc::new(Self { + codex_home: PathBuf::from("non-existent"), + inner: RwLock::new(cached), + auth_change_tx, + enable_codex_api_key_env: false, + auth_credentials_store_mode: AuthCredentialsStoreMode::File, + keyring_backend_kind: AuthKeyringBackendKind::default(), + forced_chatgpt_workspace_id: RwLock::new(None), + chatgpt_base_url: None, + agent_identity_authapi_base_url: Some( + agent_identity_authapi_base_url + .trim_end_matches('/') + .to_string(), + ), + refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), }) } @@ -1690,7 +1891,9 @@ impl AuthManager { keyring_backend_kind: AuthKeyringBackendKind::default(), forced_chatgpt_workspace_id: RwLock::new(None), chatgpt_base_url: None, + agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), + agent_identity_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), @@ -1736,6 +1939,38 @@ impl AuthManager { self.auth_cached() } + pub async fn agent_identity_auth( + &self, + policy: AgentIdentityAuthPolicy, + session_source: SessionSource, + ) -> std::io::Result> { + let Some(auth) = self.auth().await else { + return Ok(None); + }; + if policy == AgentIdentityAuthPolicy::ChatGptAuth && matches!(auth, CodexAuth::Chatgpt(_)) { + let _bootstrap_permit = self + .agent_identity_lock + .acquire() + .await + .map_err(std::io::Error::other)?; + return auth + .agent_identity_auth( + policy, + self.agent_identity_authapi_base_url.as_deref(), + self.forced_chatgpt_workspace_id(), + session_source, + ) + .await; + } + auth.agent_identity_auth( + policy, + self.agent_identity_authapi_base_url.as_deref(), + self.forced_chatgpt_workspace_id(), + session_source, + ) + .await + } + /// Force a reload of the auth information from auth.json. Returns /// whether the auth value changed. pub async fn reload(&self) -> bool { @@ -1838,7 +2073,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, + self.agent_identity_authapi_base_url.as_deref(), ) .await .ok() @@ -1911,6 +2146,7 @@ impl AuthManager { codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, + forced_chatgpt_workspace_id: Option>, chatgpt_base_url: Option, keyring_backend_kind: AuthKeyringBackendKind, ) -> Arc { @@ -1919,6 +2155,7 @@ impl AuthManager { codex_home, enable_codex_api_key_env, auth_credentials_store_mode, + forced_chatgpt_workspace_id, chatgpt_base_url, keyring_backend_kind, ) @@ -1931,17 +2168,15 @@ impl AuthManager { config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { - Arc::new( - Self::new_with_workspace_restriction( - config.codex_home(), - enable_codex_api_key_env, - config.cli_auth_credentials_store_mode(), - config.forced_chatgpt_workspace_id(), - Some(config.chatgpt_base_url()), - config.auth_keyring_backend_kind(), - ) - .await, + Self::shared( + config.codex_home(), + enable_codex_api_key_env, + config.cli_auth_credentials_store_mode(), + config.forced_chatgpt_workspace_id(), + Some(config.chatgpt_base_url()), + config.auth_keyring_backend_kind(), ) + .await } pub fn unauthorized_recovery(self: &Arc) -> UnauthorizedRecovery { diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index ebec06417..19ba018f3 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -51,7 +51,7 @@ pub struct AuthDotJson { pub last_refresh: Option>, #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_identity: Option, + pub agent_identity: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub personal_access_token: Option, @@ -60,6 +60,32 @@ pub struct AuthDotJson { pub bedrock_api_key: Option, } +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +#[serde(untagged)] +pub enum AgentIdentityStorage { + Jwt(String), + Record(AgentIdentityAuthRecord), +} + +impl AgentIdentityStorage { + pub fn has_auth_material(&self) -> bool { + match self { + Self::Jwt(jwt) => !jwt.trim().is_empty(), + Self::Record(record) => { + !record.agent_runtime_id.trim().is_empty() + && !record.agent_private_key.trim().is_empty() + } + } + } + + pub(crate) fn as_record(&self) -> Option<&AgentIdentityAuthRecord> { + match self { + Self::Jwt(_) => None, + Self::Record(record) => Some(record), + } + } +} + #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] pub struct AgentIdentityAuthRecord { pub agent_runtime_id: String, @@ -69,6 +95,8 @@ pub struct AgentIdentityAuthRecord { pub email: String, pub plan_type: AccountPlanType, pub chatgpt_account_is_fedramp: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub task_id: Option, } impl AgentIdentityAuthRecord { @@ -90,6 +118,7 @@ impl From for AgentIdentityAuthRecord { email: claims.email, plan_type: claims.plan_type.into(), chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, + task_id: None, } } } diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 72d59b660..db51a7746 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -81,7 +81,38 @@ async fn file_storage_round_trips_agent_identity_auth() -> anyhow::Result<()> { openai_api_key: None, tokens: None, last_refresh: None, - agent_identity: Some(agent_identity), + agent_identity: Some(AgentIdentityStorage::Jwt(agent_identity)), + personal_access_token: None, + bedrock_api_key: None, + }; + + storage.save(&auth_dot_json)?; + + let loaded = storage.load()?; + assert_eq!(Some(auth_dot_json), loaded); + Ok(()) +} + +#[tokio::test] +async fn file_storage_round_trips_registered_agent_identity_auth() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let record = AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-id".to_string(), + agent_private_key: "private-key".to_string(), + account_id: "account-id".to_string(), + chatgpt_user_id: "user-id".to_string(), + email: "user@example.com".to_string(), + plan_type: AccountPlanType::Pro, + chatgpt_account_is_fedramp: false, + task_id: Some("task-id".to_string()), + }; + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: None, + last_refresh: None, + agent_identity: Some(AgentIdentityStorage::Record(record)), personal_access_token: None, bedrock_api_key: None, }; @@ -139,8 +170,8 @@ async fn file_storage_loads_agent_identity_as_jwt() -> anyhow::Result<()> { let loaded = storage.load()?; assert_eq!( - loaded.expect("auth should load").agent_identity.as_deref(), - Some(agent_identity_jwt.as_str()) + loaded.expect("auth should load").agent_identity, + Some(AgentIdentityStorage::Jwt(agent_identity_jwt)) ); Ok(()) } diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 367f27630..2b8e48b4a 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -17,6 +17,7 @@ pub use server::ServerOptions; pub use server::ShutdownHandle; pub use server::run_login_server; +pub use auth::AgentIdentityAuthPolicy; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthKeyringBackendKind; diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index 406572229..af892d76a 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -1217,6 +1217,7 @@ impl RefreshTokenTestContext { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index 5df73e932..3c47bc88b 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -201,6 +201,7 @@ async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> { codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 55cf89307..f189b2314 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -1163,7 +1163,6 @@ 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 {