diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5ff3f462f1..fd4ed6d8d9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1758,6 +1758,7 @@ dependencies = [ "codex-protocol", "crypto_box", "ed25519-dalek", + "jsonwebtoken", "pretty_assertions", "rand 0.9.3", "reqwest", diff --git a/codex-rs/agent-identity/Cargo.toml b/codex-rs/agent-identity/Cargo.toml index 7976c3354b..4610d6ec9b 100644 --- a/codex-rs/agent-identity/Cargo.toml +++ b/codex-rs/agent-identity/Cargo.toml @@ -19,6 +19,7 @@ chrono = { workspace = true } codex-protocol = { workspace = true } crypto_box = { workspace = true } ed25519-dalek = { workspace = true } +jsonwebtoken = { workspace = true } rand = { workspace = true } reqwest = { workspace = true, features = ["json"] } serde = { workspace = true, features = ["derive"] } diff --git a/codex-rs/agent-identity/src/lib.rs b/codex-rs/agent-identity/src/lib.rs index a6d7e25dfd..66bb809df1 100644 --- a/codex-rs/agent-identity/src/lib.rs +++ b/codex-rs/agent-identity/src/lib.rs @@ -8,6 +8,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use chrono::SecondsFormat; use chrono::Utc; +use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::protocol::SessionSource; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; @@ -15,6 +16,12 @@ use ed25519_dalek::SigningKey; use ed25519_dalek::VerifyingKey; use ed25519_dalek::pkcs8::DecodePrivateKey; use ed25519_dalek::pkcs8::EncodePrivateKey; +use jsonwebtoken::Algorithm; +use jsonwebtoken::DecodingKey; +use jsonwebtoken::Validation; +use jsonwebtoken::decode; +use jsonwebtoken::decode_header; +use jsonwebtoken::jwk::JwkSet; use rand::TryRngCore; use rand::rngs::OsRng; use serde::Deserialize; @@ -23,6 +30,9 @@ use sha2::Digest as _; use sha2::Sha512; const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30); +const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10); +const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server"; +const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity"; /// Stored key material for a registered agent identity. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -50,6 +60,21 @@ pub struct GeneratedAgentKeyMaterial { pub public_key_ssh: String, } +#[derive(Clone, Debug, Deserialize, PartialEq, Eq)] +pub struct AgentIdentityJwtClaims { + pub iss: String, + pub aud: String, + pub iat: usize, + pub exp: usize, + pub agent_runtime_id: String, + pub agent_private_key: String, + pub account_id: String, + pub chatgpt_user_id: String, + pub email: String, + pub plan_type: AccountPlanType, + pub chatgpt_account_is_fedramp: bool, +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] struct AgentAssertionEnvelope { agent_runtime_id: String, @@ -98,6 +123,40 @@ pub fn authorization_header_for_agent_task( Ok(format!("AgentAssertion {serialized_assertion}")) } +pub async fn fetch_agent_identity_jwks( + client: &reqwest::Client, + chatgpt_base_url: &str, +) -> Result { + client + .get(agent_identity_jwks_url(chatgpt_base_url)) + .timeout(AGENT_IDENTITY_JWKS_TIMEOUT) + .send() + .await + .context("failed to fetch agent identity JWKS")? + .error_for_status() + .context("failed to fetch agent identity JWKS")? + .json() + .await + .context("failed to decode agent identity JWKS") +} + +pub fn decode_agent_identity_jwt(jwt: &str, jwks: &JwkSet) -> Result { + let header = decode_header(jwt).context("failed to decode agent identity JWT header")?; + let kid = header + .kid + .context("agent identity JWT header does not include a kid")?; + let jwk = jwks + .find(&kid) + .with_context(|| format!("agent identity JWT kid {kid} is not trusted"))?; + let decoding_key = DecodingKey::from_jwk(jwk).context("failed to build JWT decoding key")?; + let mut validation = Validation::new(Algorithm::RS256); + validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]); + validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]); + decode::(jwt, &decoding_key, &validation) + .map(|data| data.claims) + .context("failed to verify agent identity JWT") +} + pub fn sign_task_registration_payload( key: AgentIdentityKey<'_>, timestamp: &str, @@ -217,6 +276,11 @@ pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { format!("{trimmed}/authenticate_app_v2") } +pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/api/codex/agent-identities/jwks") +} + pub fn agent_identity_request_id() -> Result { let mut request_id_bytes = [0u8; 16]; OsRng diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index cddc5d5856..ceb82ef7ed 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -1353,7 +1353,7 @@ impl CodexMessageProcessor { self.config.cli_auth_credentials_store_mode, ) { Ok(()) => { - self.auth_manager.reload(); + self.auth_manager.reload().await; Ok(()) } Err(err) => Err(JSONRPCErrorError { @@ -1504,7 +1504,7 @@ impl CodexMessageProcessor { .await; if success { - auth_manager.reload(); + auth_manager.reload().await; config_manager.replace_cloud_requirements_loader( auth_manager.clone(), chatgpt_base_url, @@ -1612,7 +1612,7 @@ impl CodexMessageProcessor { .await; if success { - auth_manager.reload(); + auth_manager.reload().await; config_manager.replace_cloud_requirements_loader( auth_manager.clone(), chatgpt_base_url, @@ -1748,7 +1748,7 @@ impl CodexMessageProcessor { self.outgoing.send_error(request_id, error).await; return; } - self.auth_manager.reload(); + self.auth_manager.reload().await; self.config_manager.replace_cloud_requirements_loader( self.auth_manager.clone(), self.config.chatgpt_base_url.clone(), diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 729f6d04af..01ced14f97 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -391,7 +391,8 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { let processor_outgoing = Arc::clone(&outgoing_message_sender); let auth_manager = - AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env); + AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) + .await; let config_manager = ConfigManager::new( args.config.codex_home.to_path_buf(), args.cli_overrides, diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index d9b403165c..0a78365f3c 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -426,7 +426,7 @@ pub async fn run_main_with_transport( config_manager .replace_thread_config_loader(Arc::clone(&discovered_thread_config_loader)); let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; config_manager.replace_cloud_requirements_loader(auth_manager, config.chatgpt_base_url); } Err(err) => { @@ -588,7 +588,7 @@ pub async fn run_main_with_transport( } let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; let remote_control_enabled = config.features.enabled(Feature::RemoteControl); if transport_accept_handles.is_empty() && !remote_control_enabled { @@ -669,7 +669,7 @@ pub async fn run_main_with_transport( let outgoing_message_sender = Arc::new(OutgoingMessageSender::new(outgoing_tx)); let outbound_control_tx = outbound_control_tx; let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { outgoing: outgoing_message_sender, arg0_paths, diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs index 5b6690c0ba..84ef24f99c 100644 --- a/codex-rs/app-server/src/message_processor/tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor/tracing_tests.rs @@ -127,7 +127,7 @@ impl TracingHarness { let server = create_mock_responses_server_repeating_assistant("Done").await; let codex_home = TempDir::new()?; let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?); - let (processor, outgoing_rx) = build_test_processor(config); + let (processor, outgoing_rx) = build_test_processor(config).await; let tracing = init_test_tracing(); tracing.exporter.reset(); tracing::callsite::rebuild_interest_cache(); @@ -257,7 +257,7 @@ async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result, ) -> ( Arc, @@ -266,7 +266,7 @@ fn build_test_processor( let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); let auth_manager = - AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false).await; let config_manager = ConfigManager::new( config.codex_home.to_path_buf(), Vec::new(), diff --git a/codex-rs/app-server/src/transport/remote_control/tests.rs b/codex-rs/app-server/src/transport/remote_control/tests.rs index 6b0051f8db..82e02ae5de 100644 --- a/codex-rs/app-server/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server/src/transport/remote_control/tests.rs @@ -497,7 +497,8 @@ async fn remote_control_start_allows_missing_auth_when_enabled() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; let (transport_event_tx, _transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); @@ -1085,7 +1086,8 @@ async fn remote_control_waits_for_account_id_before_enrolling() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; let expected_server_name = gethostname().to_string_lossy().trim().to_string(); let expected_enrollment = RemoteControlEnrollment { account_id: "account_id".to_string(), diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 464832e34a..4cd078455a 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -706,7 +706,7 @@ pub(crate) async fn load_remote_control_auth( "remote control requires ChatGPT authentication", )); } - auth_manager.reload(); + auth_manager.reload().await; reloaded = true; continue; }; @@ -714,7 +714,7 @@ pub(crate) async fn load_remote_control_auth( break auth; } if auth.get_account_id().is_none() && !reloaded { - auth_manager.reload(); + auth_manager.reload().await; reloaded = true; continue; } @@ -1090,7 +1090,8 @@ mod tests { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut enrollment = Some(RemoteControlEnrollment { account_id: "account_id".to_string(), @@ -1172,7 +1173,8 @@ mod tests { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut enrollment = None; save_auth( diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index 42aac41138..05d8186686 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -21,7 +21,7 @@ pub(crate) async fn chatgpt_get_request_with_timeout( ) -> anyhow::Result { let chatgpt_base_url = &config.chatgpt_base_url; let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; let auth = auth_manager .auth() .await diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index 62e8040940..9dba71ce3a 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -26,7 +26,7 @@ const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); async fn apps_enabled(config: &Config) -> bool { let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; let auth = auth_manager.auth().await; config .features @@ -35,7 +35,7 @@ async fn apps_enabled(config: &Config) -> bool { async fn connector_auth(config: &Config) -> anyhow::Result { let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; let auth = auth_manager .auth() .await diff --git a/codex-rs/cli/src/login.rs b/codex-rs/cli/src/login.rs index 42241aa933..e1a950b885 100644 --- a/codex-rs/cli/src/login.rs +++ b/codex-rs/cli/src/login.rs @@ -316,7 +316,13 @@ pub async fn run_login_with_device_code_fallback_to_browser( pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! { let config = load_config_or_exit(cli_config_overrides).await; - match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode) { + match CodexAuth::from_auth_storage_with_base_url( + &config.codex_home, + config.cli_auth_credentials_store_mode, + Some(&config.chatgpt_base_url), + ) + .await + { Ok(Some(auth)) => match auth.auth_mode() { AuthMode::ApiKey => match auth.get_token() { Ok(api_key) => { diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 35005f5be3..8376a80d48 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1367,7 +1367,7 @@ async fn run_debug_models_command( .map_err(anyhow::Error::msg)?; let config = Config::load_with_cli_overrides(cli_overrides).await?; let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true); + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; let models_manager = build_models_manager(&config, auth_manager, CollaborationModesConfig::default()); models_manager diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 8c51888a16..06e50675cc 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -722,7 +722,7 @@ pub fn cloud_requirements_loader( }) } -pub fn cloud_requirements_loader_for_storage( +pub async fn cloud_requirements_loader_for_storage( codex_home: PathBuf, enable_codex_api_key_env: bool, credentials_store_mode: AuthCredentialsStoreMode, @@ -733,7 +733,8 @@ pub fn cloud_requirements_loader_for_storage( enable_codex_api_key_env, credentials_store_mode, Some(chatgpt_base_url.clone()), - ); + ) + .await; cloud_requirements_loader(auth_manager, chatgpt_base_url, codex_home) } @@ -848,7 +849,7 @@ mod tests { Ok(()) } - fn auth_manager_with_api_key() -> Arc { + async fn auth_manager_with_api_key() -> Arc { let tmp = tempdir().expect("tempdir"); let auth_json = json!({ "OPENAI_API_KEY": "sk-test-key", @@ -856,15 +857,18 @@ mod tests { "last_refresh": null, }); write_auth_json(tmp.path(), auth_json).expect("write auth"); - Arc::new(AuthManager::new( - tmp.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )) + Arc::new( + AuthManager::new( + tmp.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ) } - fn auth_manager_with_plan_and_identity( + async fn auth_manager_with_plan_and_identity( plan_type: &str, chatgpt_user_id: Option<&str>, account_id: Option<&str>, @@ -881,12 +885,15 @@ mod tests { ), ) .expect("write auth"); - Arc::new(AuthManager::new( - tmp.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )) + Arc::new( + AuthManager::new( + tmp.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ) } fn chatgpt_auth_json( @@ -970,7 +977,7 @@ mod tests { manager: Arc, } - fn managed_auth_context( + async fn managed_auth_context( plan_type: &str, chatgpt_user_id: Option<&str>, account_id: Option<&str>, @@ -990,18 +997,22 @@ mod tests { ) .expect("write auth"); ManagedAuthContext { - manager: Arc::new(AuthManager::new( - home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )), + manager: Arc::new( + AuthManager::new( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ), _home: home, } } - fn auth_manager_with_plan(plan_type: &str) -> Arc { + 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 } fn parse_for_fetch(contents: Option<&str>) -> Option { @@ -1113,7 +1124,7 @@ mod tests { #[tokio::test] async fn fetch_cloud_requirements_skips_non_chatgpt_auth() { - let auth_manager = auth_manager_with_api_key(); + let auth_manager = auth_manager_with_api_key().await; let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( auth_manager, @@ -1129,7 +1140,7 @@ mod tests { async fn fetch_cloud_requirements_skips_non_business_or_enterprise_plan() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("pro"), + auth_manager_with_plan("pro").await, Arc::new(StaticFetcher { contents: None }), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1142,7 +1153,7 @@ mod tests { async fn fetch_cloud_requirements_skips_team_like_usage_based_plan() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("self_serve_business_usage_based"), + auth_manager_with_plan("self_serve_business_usage_based").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1156,7 +1167,7 @@ mod tests { async fn fetch_cloud_requirements_allows_business_plan() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1188,7 +1199,7 @@ mod tests { async fn fetch_cloud_requirements_allows_business_like_usage_based_plan() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("enterprise_cbp_usage_based"), + auth_manager_with_plan("enterprise_cbp_usage_based").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1220,7 +1231,7 @@ mod tests { async fn fetch_cloud_requirements_allows_hc_plan_as_enterprise() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("hc"), + auth_manager_with_plan("hc").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1324,7 +1335,7 @@ enabled = false #[tokio::test(start_paused = true)] async fn fetch_cloud_requirements_times_out() { - let auth_manager = auth_manager_with_plan("enterprise"); + let auth_manager = auth_manager_with_plan("enterprise").await; let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( auth_manager, @@ -1351,7 +1362,7 @@ enabled = false ])); let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1400,12 +1411,15 @@ enabled = false ), ) .expect("write initial auth"); - let auth_manager = Arc::new(AuthManager::new( - auth_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )); + let auth_manager = Arc::new( + AuthManager::new( + auth_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ); write_auth_json( auth_home.path(), @@ -1474,12 +1488,15 @@ enabled = false ), ) .expect("write initial auth"); - let auth_manager = Arc::new(AuthManager::new( - auth_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )); + let auth_manager = Arc::new( + AuthManager::new( + auth_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ); write_auth_json( auth_home.path(), @@ -1554,7 +1571,8 @@ enabled = false Some("account-12345"), "stale-access-token", "test-refresh-token", - ); + ) + .await; write_auth_json( auth._home.path(), chatgpt_auth_json( @@ -1606,12 +1624,15 @@ enabled = false ), ) .expect("write auth"); - let auth_manager = Arc::new(AuthManager::new( - auth_home.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - /*chatgpt_base_url*/ None, - )); + let auth_manager = Arc::new( + AuthManager::new( + auth_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await, + ); let fetcher = Arc::new(UnauthorizedFetcher { message: @@ -1648,7 +1669,7 @@ enabled = false ])); let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1673,7 +1694,7 @@ enabled = false ))])); let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, fetcher, codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1695,7 +1716,7 @@ enabled = false async fn fetch_cloud_requirements_uses_cache_when_valid() { let codex_home = tempdir().expect("tempdir"); let prime_service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1706,7 +1727,7 @@ enabled = false let fetcher = Arc::new(SequenceFetcher::new(vec![Err(request_error())])); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1742,7 +1763,8 @@ enabled = false "business", /*chatgpt_user_id*/ None, Some("account-12345"), - ), + ) + .await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1785,7 +1807,7 @@ enabled = false async fn fetch_cloud_requirements_does_not_use_cache_when_auth_identity_is_incomplete() { let codex_home = tempdir().expect("tempdir"); let prime_service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1802,7 +1824,8 @@ enabled = false "business", /*chatgpt_user_id*/ None, Some("account-12345"), - ), + ) + .await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1838,7 +1861,8 @@ enabled = false "business", Some("user-12345"), Some("account-12345"), - ), + ) + .await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1855,7 +1879,8 @@ enabled = false "business", Some("user-99999"), Some("account-12345"), - ), + ) + .await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1887,7 +1912,7 @@ enabled = false async fn fetch_cloud_requirements_ignores_tampered_cache() { let codex_home = tempdir().expect("tempdir"); let prime_service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -1912,7 +1937,7 @@ enabled = false "allowed_approval_policies = [\"never\"]".to_string(), ))])); let service = CloudRequirementsService::new( - auth_manager_with_plan("enterprise"), + auth_manager_with_plan("enterprise").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -1970,7 +1995,7 @@ enabled = false "allowed_approval_policies = [\"never\"]".to_string(), ))])); let service = CloudRequirementsService::new( - auth_manager_with_plan("enterprise"), + auth_manager_with_plan("enterprise").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -2002,7 +2027,7 @@ enabled = false async fn fetch_cloud_requirements_writes_signed_cache() { let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, Arc::new(StaticFetcher { contents: Some("allowed_approval_policies = [\"never\"]".to_string()), }), @@ -2065,7 +2090,7 @@ enabled = false let fetcher = Arc::new(SequenceFetcher::new(vec![Ok(None), Err(request_error())])); let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("enterprise"), + auth_manager_with_plan("enterprise").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -2083,7 +2108,7 @@ enabled = false ])); let codex_home = tempdir().expect("tempdir"); let service = CloudRequirementsService::new( - auth_manager_with_plan("enterprise"), + auth_manager_with_plan("enterprise").await, fetcher.clone(), codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, @@ -2116,7 +2141,7 @@ enabled = false )), ])); let service = CloudRequirementsService::new( - auth_manager_with_plan("business"), + auth_manager_with_plan("business").await, fetcher, codex_home.path().to_path_buf(), CLOUD_REQUIREMENTS_TIMEOUT, diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index e433b892e5..9a5056aa66 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -44,12 +44,15 @@ pub fn normalize_base_url(input: &str) -> String { pub async fn load_auth_manager(chatgpt_base_url: Option) -> Option { // TODO: pass in cli overrides once cloud tasks properly support them. let config = Config::load_with_cli_overrides(Vec::new()).await.ok()?; - Some(AuthManager::new( - config.codex_home.to_path_buf(), - /*enable_codex_api_key_env*/ false, - config.cli_auth_credentials_store_mode, - chatgpt_base_url.or(Some(config.chatgpt_base_url)), - )) + Some( + AuthManager::new( + config.codex_home.to_path_buf(), + /*enable_codex_api_key_env*/ false, + config.cli_auth_credentials_store_mode, + chatgpt_base_url.or(Some(config.chatgpt_base_url)), + ) + .await, + ) } /// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`, diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 968b93214c..b8a5f07840 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -144,7 +144,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools( config: &Config, ) -> Option> { let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; let auth = auth_manager.auth().await; if !config .features @@ -216,7 +216,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( environment_manager: &EnvironmentManager, ) -> anyhow::Result { let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; let auth = auth_manager.auth().await; if !config .features @@ -434,7 +434,7 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( Some(auth) } else { let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; loaded_auth = auth_manager.auth().await; loaded_auth.as_ref() }; diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index b6c66410a1..f782c0d232 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -29,7 +29,7 @@ pub async fn build_prompt_input( config.ephemeral = true; let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; let local_runtime_paths = ExecServerRuntimePaths::from_optional_paths( config.codex_self_exe.clone(), diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 67751161d3..cb9602a3f7 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1096,7 +1096,8 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { config.model_provider = model_provider; let auth_manager = - match CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File) { + match CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File).await + { Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth), Ok(None) => panic!("No CodexAuth found in codex_home"), Err(e) => panic!("Failed to load CodexAuth: {e}"), diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index c96e06279b..154291389b 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -348,7 +348,8 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result /*enable_codex_api_key_env*/ false, config_toml.cli_auth_credentials_store.unwrap_or_default(), chatgpt_base_url, - ); + ) + .await; let run_cli_overrides = cli_kv_overrides.clone(); let run_loader_overrides = loader_overrides.clone(); let run_cloud_requirements = cloud_requirements.clone(); @@ -438,7 +439,9 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result auth_credentials_store_mode: config.cli_auth_credentials_store_mode, forced_login_method: config.forced_login_method, forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(), - }) { + }) + .await + { eprintln!("{err}"); std::process::exit(1); } diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 5f2dc9cfc8..6b18806831 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -1,10 +1,7 @@ -use std::sync::Arc; - use codex_agent_identity::AgentIdentityKey; use codex_agent_identity::normalize_chatgpt_base_url; 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; @@ -12,51 +9,35 @@ use super::storage::AgentIdentityAuthRecord; const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api"; -#[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, + chatgpt_base_url: Option<&str>, + ) -> std::io::Result { + let base_url = normalize_chatgpt_base_url( + chatgpt_base_url.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL), + ); + let process_task_id = register_agent_task(&build_reqwest_client(), &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, chatgpt_base_url: Option) -> std::io::Result<()> { - self.process_task_id - .get_or_try_init(|| async { - let base_url = normalize_chatgpt_base_url( - chatgpt_base_url - .as_deref() - .unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL), - ); - register_agent_task(&build_reqwest_client(), &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 { @@ -78,11 +59,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 6f17822e77..86a81b90d2 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -78,10 +78,11 @@ fn login_with_api_key_overwrites_existing_auth_json() { assert!(auth.tokens.is_none(), "tokens should be cleared"); } -#[test] -fn missing_auth_json_returns_none() { +#[tokio::test] +async fn missing_auth_json_returns_none() { let dir = tempdir().unwrap(); let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File) + .await .expect("call should succeed"); assert_eq!(auth, None); } @@ -105,6 +106,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .unwrap() .unwrap(); assert_eq!(None, auth.api_key()); @@ -158,6 +160,7 @@ async fn loads_api_key_from_auth_json() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .unwrap() .unwrap(); assert_eq!(auth.auth_mode(), AuthMode::ApiKey); @@ -184,15 +187,16 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { Ok(()) } -#[test] -fn unauthorized_recovery_reports_mode_and_step_names() { +#[tokio::test] +async fn unauthorized_recovery_reports_mode_and_step_names() { let dir = tempdir().unwrap(); let manager = AuthManager::shared( dir.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; let managed = UnauthorizedRecovery { manager: Arc::clone(&manager), step: UnauthorizedRecoveryStep::Reload, @@ -212,8 +216,8 @@ fn unauthorized_recovery_reports_mode_and_step_names() { assert_eq!(external.step_name(), "external_refresh"); } -#[test] -fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { +#[tokio::test] +async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { let codex_home = tempdir().unwrap(); write_auth_file( AuthFileParams { @@ -230,6 +234,7 @@ fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); let mut updated_auth_dot_json = auth @@ -245,7 +250,9 @@ fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { codex_home.path(), updated_auth_dot_json, AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, ) + .await .expect("updated auth should parse"); let manager = AuthManager::from_auth_for_testing(auth.clone()); @@ -594,8 +601,9 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() { ) .await; - let err = - super::enforce_login_restrictions(&config).expect_err("expected method mismatch to error"); + let err = super::enforce_login_restrictions(&config) + .await + .expect_err("expected method mismatch to error"); assert!(err.to_string().contains("ChatGPT login is required")); assert!( !codex_home.path().join("auth.json").exists(), @@ -625,6 +633,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() { .await; let err = super::enforce_login_restrictions(&config) + .await .expect_err("expected workspace mismatch to error"); assert!(err.to_string().contains("workspace org_mine")); assert!( @@ -654,7 +663,9 @@ async fn enforce_login_restrictions_allows_matching_workspace() { ) .await; - super::enforce_login_restrictions(&config).expect("matching workspace should succeed"); + super::enforce_login_restrictions(&config) + .await + .expect("matching workspace should succeed"); assert!( codex_home.path().join("auth.json").exists(), "auth.json should remain when restrictions pass" @@ -675,7 +686,9 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f ) .await; - super::enforce_login_restrictions(&config).expect("matching workspace should succeed"); + super::enforce_login_restrictions(&config) + .await + .expect("matching workspace should succeed"); assert!( codex_home.path().join("auth.json").exists(), "auth.json should remain when restrictions pass" @@ -696,6 +709,7 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() { .await; let err = super::enforce_login_restrictions(&config) + .await .expect_err("environment API key should not satisfy forced ChatGPT login"); assert!( err.to_string() @@ -703,8 +717,8 @@ async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() { ); } -#[test] -fn plan_type_maps_known_plan() { +#[tokio::test] +async fn plan_type_maps_known_plan() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( AuthFileParams { @@ -721,14 +735,15 @@ fn plan_type_maps_known_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Pro)); } -#[test] -fn plan_type_maps_self_serve_business_usage_based_plan() { +#[tokio::test] +async fn plan_type_maps_self_serve_business_usage_based_plan() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( AuthFileParams { @@ -745,6 +760,7 @@ fn plan_type_maps_self_serve_business_usage_based_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); @@ -754,8 +770,8 @@ fn plan_type_maps_self_serve_business_usage_based_plan() { ); } -#[test] -fn plan_type_maps_enterprise_cbp_usage_based_plan() { +#[tokio::test] +async fn plan_type_maps_enterprise_cbp_usage_based_plan() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( AuthFileParams { @@ -772,6 +788,7 @@ fn plan_type_maps_enterprise_cbp_usage_based_plan() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); @@ -781,8 +798,8 @@ fn plan_type_maps_enterprise_cbp_usage_based_plan() { ); } -#[test] -fn plan_type_maps_unknown_to_unknown() { +#[tokio::test] +async fn plan_type_maps_unknown_to_unknown() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( AuthFileParams { @@ -799,14 +816,15 @@ fn plan_type_maps_unknown_to_unknown() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); pretty_assertions::assert_eq!(auth.account_plan_type(), Some(AccountPlanType::Unknown)); } -#[test] -fn missing_plan_type_maps_to_unknown() { +#[tokio::test] +async fn missing_plan_type_maps_to_unknown() { let codex_home = tempdir().unwrap(); let _jwt = write_auth_file( AuthFileParams { @@ -823,6 +841,7 @@ fn missing_plan_type_maps_to_unknown() { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, ) + .await .expect("load auth") .expect("auth available"); diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 419c6a4bac..5b0daedca9 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -16,6 +16,10 @@ use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use tokio::sync::Semaphore; +use codex_agent_identity::AgentIdentityJwtClaims; +use codex_agent_identity::decode_agent_identity_jwt; +use codex_agent_identity::fetch_agent_identity_jwks; +use codex_agent_identity::normalize_chatgpt_base_url; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; @@ -29,6 +33,7 @@ pub use crate::auth::storage::AuthDotJson; 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; @@ -193,10 +198,11 @@ impl From for std::io::Error { } impl CodexAuth { - fn from_auth_dot_json( + async fn from_auth_dot_json( codex_home: &Path, auth_dot_json: AuthDotJson, auth_credentials_store_mode: AuthCredentialsStoreMode, + chatgpt_base_url: Option<&str>, ) -> std::io::Result { let auth_mode = auth_dot_json.resolved_mode(); let client = create_client(); @@ -212,7 +218,9 @@ impl CodexAuth { "agent identity auth is missing an agent identity record.", )); }; - return Ok(Self::AgentIdentity(AgentIdentityAuth::new(record))); + return Ok(Self::AgentIdentity( + AgentIdentityAuth::load(record, chatgpt_base_url).await?, + )); } let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode); @@ -234,7 +242,7 @@ impl CodexAuth { } } - pub fn from_auth_storage( + pub async fn from_auth_storage( codex_home: &Path, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result> { @@ -243,6 +251,21 @@ impl CodexAuth { /*enable_codex_api_key_env*/ false, auth_credentials_store_mode, ) + .await + } + + pub async fn from_auth_storage_with_base_url( + codex_home: &Path, + auth_credentials_store_mode: AuthCredentialsStoreMode, + chatgpt_base_url: Option<&str>, + ) -> std::io::Result> { + load_auth_with_base_url( + codex_home, + /*enable_codex_api_key_env*/ false, + auth_credentials_store_mode, + chatgpt_base_url, + ) + .await } pub fn auth_mode(&self) -> AuthMode { @@ -316,16 +339,6 @@ impl CodexAuth { } } - pub async fn initialize_runtime( - &self, - chatgpt_base_url: Option, - ) -> std::io::Result<()> { - match self { - Self::AgentIdentity(auth) => auth.ensure_runtime(chatgpt_base_url).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 { @@ -474,6 +487,7 @@ impl ChatgptAuth { pub const OPENAI_API_KEY_ENV_VAR: &str = "OPENAI_API_KEY"; pub const CODEX_API_KEY_ENV_VAR: &str = "CODEX_API_KEY"; +pub const CODEX_AGENT_IDENTITY_ENV_VAR: &str = "CODEX_AGENT_IDENTITY"; pub fn read_openai_api_key_from_env() -> Option { env::var(OPENAI_API_KEY_ENV_VAR) @@ -489,6 +503,13 @@ pub fn read_codex_api_key_from_env() -> Option { .filter(|value| !value.is_empty()) } +pub fn read_codex_agent_identity_from_env() -> Option { + env::var(CODEX_AGENT_IDENTITY_ENV_VAR) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + /// 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( @@ -509,6 +530,7 @@ pub async fn logout_with_revoke( auth_credentials_store_mode, /*chatgpt_base_url*/ None, ) + .await .logout_with_revoke() .await } @@ -579,12 +601,13 @@ pub struct AuthConfig { pub forced_chatgpt_workspace_id: Option, } -pub fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { +pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> { let Some(auth) = load_auth( &config.codex_home, /*enable_codex_api_key_env*/ true, config.auth_credentials_store_mode, - )? + ) + .await? else { return Ok(()); }; @@ -684,15 +707,54 @@ fn logout_all_stores( Ok(removed_ephemeral || removed_managed) } -fn load_auth( +async fn load_agent_identity_auth_from_jwt( + agent_identity: &str, + chatgpt_base_url: Option<&str>, +) -> std::io::Result { + let normalized_base_url = + normalize_chatgpt_base_url(chatgpt_base_url.unwrap_or("https://chatgpt.com/backend-api")); + let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), &normalized_base_url) + .await + .map_err(std::io::Error::other)?; + let claims = decode_agent_identity_jwt(agent_identity, &jwks).map_err(std::io::Error::other)?; + let record = agent_identity_record_from_claims(claims); + Ok(CodexAuth::AgentIdentity( + AgentIdentityAuth::load(record, Some(&normalized_base_url)).await?, + )) +} + +fn agent_identity_record_from_claims(claims: AgentIdentityJwtClaims) -> AgentIdentityAuthRecord { + AgentIdentityAuthRecord { + agent_runtime_id: claims.agent_runtime_id, + agent_private_key: claims.agent_private_key, + account_id: claims.account_id, + chatgpt_user_id: claims.chatgpt_user_id, + email: claims.email, + plan_type: claims.plan_type, + chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp, + } +} + +async fn load_auth( codex_home: &Path, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> std::io::Result> { - let build_auth = |auth_dot_json: AuthDotJson, storage_mode| { - CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, storage_mode) - }; + load_auth_with_base_url( + codex_home, + enable_codex_api_key_env, + auth_credentials_store_mode, + /*chatgpt_base_url*/ None, + ) + .await +} +async fn load_auth_with_base_url( + codex_home: &Path, + enable_codex_api_key_env: bool, + auth_credentials_store_mode: AuthCredentialsStoreMode, + chatgpt_base_url: Option<&str>, +) -> std::io::Result> { // API key via env var takes precedence over any other auth method. if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() { return Ok(Some(CodexAuth::from_api_key(api_key.as_str()))); @@ -705,10 +767,22 @@ fn load_auth( AuthCredentialsStoreMode::Ephemeral, ); if let Some(auth_dot_json) = ephemeral_storage.load()? { - let auth = build_auth(auth_dot_json, AuthCredentialsStoreMode::Ephemeral)?; + let auth = CodexAuth::from_auth_dot_json( + codex_home, + auth_dot_json, + AuthCredentialsStoreMode::Ephemeral, + chatgpt_base_url, + ) + .await?; return Ok(Some(auth)); } + if let Some(agent_identity) = read_codex_agent_identity_from_env() { + return load_agent_identity_auth_from_jwt(&agent_identity, chatgpt_base_url) + .await + .map(Some); + } + // If the caller explicitly requested ephemeral auth, there is no persisted fallback. if auth_credentials_store_mode == AuthCredentialsStoreMode::Ephemeral { return Ok(None); @@ -721,7 +795,13 @@ fn load_auth( None => return Ok(None), }; - let auth = build_auth(auth_dot_json, auth_credentials_store_mode)?; + let auth = CodexAuth::from_auth_dot_json( + codex_home, + auth_dot_json, + auth_credentials_store_mode, + chatgpt_base_url, + ) + .await?; Ok(Some(auth)) } @@ -1135,6 +1215,7 @@ impl UnauthorizedRecovery { match self .manager .reload_if_account_id_matches(self.expected_account_id.as_deref()) + .await { ReloadOutcome::ReloadedChanged => { self.step = UnauthorizedRecoveryStep::RefreshToken; @@ -1245,17 +1326,19 @@ impl AuthManager { /// 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 fn new( + pub async fn new( codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option, ) -> Self { - let managed_auth = load_auth( + let managed_auth = load_auth_with_base_url( &codex_home, enable_codex_api_key_env, auth_credentials_store_mode, + chatgpt_base_url.as_deref(), ) + .await .ok() .flatten(); Self { @@ -1358,23 +1441,21 @@ 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 /// whether the auth value changed. - pub fn reload(&self) -> bool { + pub async fn reload(&self) -> bool { tracing::info!("Reloading auth"); - let new_auth = self.load_auth_from_storage(); + let new_auth = self.load_auth_from_storage().await; self.set_cached_auth(new_auth) } - fn reload_if_account_id_matches(&self, expected_account_id: Option<&str>) -> ReloadOutcome { + async fn reload_if_account_id_matches( + &self, + expected_account_id: Option<&str>, + ) -> ReloadOutcome { let expected_account_id = match expected_account_id { Some(account_id) => account_id, None => { @@ -1383,7 +1464,7 @@ impl AuthManager { } }; - let new_auth = self.load_auth_from_storage(); + let new_auth = self.load_auth_from_storage().await; let new_account_id = new_auth.as_ref().and_then(CodexAuth::get_account_id); if new_account_id.as_deref() != Some(expected_account_id) { @@ -1454,12 +1535,14 @@ impl AuthManager { } } - fn load_auth_from_storage(&self) -> Option { - load_auth( + async fn load_auth_from_storage(&self) -> Option { + load_auth_with_base_url( &self.codex_home, self.enable_codex_api_key_env, self.auth_credentials_store_mode, + self.chatgpt_base_url.as_deref(), ) + .await .ok() .flatten() } @@ -1523,22 +1606,25 @@ impl AuthManager { } /// Convenience constructor returning an `Arc` wrapper. - pub fn shared( + pub async fn shared( codex_home: PathBuf, enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, chatgpt_base_url: Option, ) -> Arc { - Arc::new(Self::new( - codex_home, - enable_codex_api_key_env, - auth_credentials_store_mode, - chatgpt_base_url, - )) + Arc::new( + Self::new( + codex_home, + enable_codex_api_key_env, + auth_credentials_store_mode, + chatgpt_base_url, + ) + .await, + ) } /// Convenience constructor returning an `Arc` wrapper from resolved config. - pub fn shared_from_config( + pub async fn shared_from_config( config: &impl AuthManagerConfig, enable_codex_api_key_env: bool, ) -> Arc { @@ -1547,7 +1633,8 @@ impl AuthManager { enable_codex_api_key_env, config.cli_auth_credentials_store_mode(), Some(config.chatgpt_base_url()), - ); + ) + .await; auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id()); auth_manager } @@ -1613,7 +1700,10 @@ impl AuthManager { .as_ref() .and_then(CodexAuth::get_account_id); - match self.reload_if_account_id_matches(expected_account_id.as_deref()) { + match self + .reload_if_account_id_matches(expected_account_id.as_deref()) + .await + { ReloadOutcome::ReloadedChanged => { tracing::info!("Skipping token refresh because auth changed after guarded reload."); Ok(()) @@ -1680,10 +1770,10 @@ impl AuthManager { /// if a file was removed, Ok(false) if no auth file existed. On success, /// reloads the in‑memory auth cache so callers immediately observe the /// unauthenticated state. - pub fn logout(&self) -> std::io::Result { + pub async fn logout(&self) -> std::io::Result { let removed = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; // Always reload to clear any cached auth (even if file absent). - self.reload(); + self.reload().await; Ok(removed) } @@ -1696,7 +1786,7 @@ impl AuthManager { } let result = logout_all_stores(&self.codex_home, self.auth_credentials_store_mode)?; // Always reload to clear any cached auth (even if file absent). - self.reload(); + self.reload().await; Ok(result) } @@ -1792,7 +1882,7 @@ impl AuthManager { AuthCredentialsStoreMode::Ephemeral, ) .map_err(RefreshTokenError::Transient)?; - self.reload(); + self.reload().await; Ok(()) } @@ -1812,7 +1902,7 @@ impl AuthManager { refresh_response.refresh_token, ) .map_err(RefreshTokenError::from)?; - self.reload(); + self.reload().await; Ok(()) } diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index 3ae0eb5fa1..afdf466d09 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -46,7 +46,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -56,7 +56,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; ctx.auth_manager .refresh_token_from_authority() @@ -110,7 +110,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -120,7 +120,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; ctx.auth_manager .refresh_token() @@ -164,7 +164,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); @@ -175,7 +175,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); let disk_auth = AuthDotJson { @@ -230,7 +230,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -240,7 +240,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let mut disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); disk_tokens.account_id = Some("other-account".to_string()); @@ -299,7 +299,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let stale_refresh = Utc::now() - Duration::days(9); let fresh_access_token = access_token_with_expiration(Utc::now() + Duration::hours(1)); let initial_tokens = build_tokens(&fresh_access_token, INITIAL_REFRESH_TOKEN); @@ -310,7 +310,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { last_refresh: Some(stale_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let cached_auth = ctx .auth_manager @@ -347,7 +347,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let fresh_refresh = Utc::now() - Duration::days(1); let expired_access_token = access_token_with_expiration(Utc::now() - Duration::hours(1)); let initial_tokens = build_tokens(&expired_access_token, INITIAL_REFRESH_TOKEN); @@ -358,7 +358,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { last_refresh: Some(fresh_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let cached_auth = ctx .auth_manager @@ -398,7 +398,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { let server = MockServer::start().await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let stale_refresh = Utc::now() - Duration::days(9); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -408,7 +408,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { last_refresh: Some(stale_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let fresh_refresh = Utc::now() - Duration::days(1); let disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); @@ -461,7 +461,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let stale_refresh = Utc::now() - Duration::days(9); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -471,7 +471,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul last_refresh: Some(stale_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let fresh_refresh = Utc::now() - Duration::days(1); let disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); @@ -522,7 +522,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -532,7 +532,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let err = ctx .auth_manager @@ -575,7 +575,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -585,7 +585,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let first_err = ctx .auth_manager @@ -642,7 +642,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -652,7 +652,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let first_err = ctx .auth_manager @@ -723,7 +723,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -733,7 +733,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let err = ctx .auth_manager @@ -776,7 +776,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -786,7 +786,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); let disk_auth = AuthDotJson { @@ -870,7 +870,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { .mount(&server) .await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let initial_last_refresh = Utc::now() - Duration::days(1); let initial_tokens = build_tokens(INITIAL_ACCESS_TOKEN, INITIAL_REFRESH_TOKEN); let initial_auth = AuthDotJson { @@ -880,7 +880,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { last_refresh: Some(initial_last_refresh), agent_identity: None, }; - ctx.write_auth(&initial_auth)?; + ctx.write_auth(&initial_auth).await?; let mut disk_tokens = build_tokens("disk-access-token", "disk-refresh-token"); disk_tokens.account_id = Some("other-account".to_string()); @@ -941,7 +941,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let ctx = RefreshTokenTestContext::new(&server)?; + let ctx = RefreshTokenTestContext::new(&server).await?; let auth = AuthDotJson { auth_mode: Some(AuthMode::ApiKey), openai_api_key: Some("sk-test".to_string()), @@ -949,7 +949,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { last_refresh: None, agent_identity: None, }; - ctx.write_auth(&auth)?; + ctx.write_auth(&auth).await?; let mut recovery = ctx.auth_manager.unauthorized_recovery(); assert!(!recovery.has_next()); @@ -974,7 +974,7 @@ struct RefreshTokenTestContext { } impl RefreshTokenTestContext { - fn new(server: &MockServer) -> Result { + async fn new(server: &MockServer) -> Result { let codex_home = TempDir::new()?; let endpoint = format!("{}/oauth/token", server.uri()); @@ -985,7 +985,8 @@ impl RefreshTokenTestContext { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; Ok(Self { codex_home, @@ -1000,13 +1001,13 @@ impl RefreshTokenTestContext { .context("auth.json should exist") } - fn write_auth(&self, auth_dot_json: &AuthDotJson) -> Result<()> { + async fn write_auth(&self, auth_dot_json: &AuthDotJson) -> Result<()> { save_auth( self.codex_home.path(), auth_dot_json, AuthCredentialsStoreMode::File, )?; - self.auth_manager.reload(); + self.auth_manager.reload().await; Ok(()) } } diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index e703b15eb1..2364ee7f54 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -143,7 +143,8 @@ async fn auth_manager_logout_with_revoke_uses_cached_auth() -> Result<()> { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, - ); + ) + .await; save_auth( codex_home.path(), &chatgpt_auth_with_refresh_token("newer-disk-refresh-token"), diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 1d904e4577..95fc4341ee 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -141,7 +141,8 @@ pub async fn run_main( arg0_paths, Arc::new(config), environment_manager, - ); + ) + .await; async move { while let Some(msg) = incoming_rx.recv().await { match msg { diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index e0a1e77f68..32d05d1f93 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -49,7 +49,7 @@ pub(crate) struct MessageProcessor { impl MessageProcessor { /// Create a new `MessageProcessor`, retaining a handle to the outgoing /// `Sender` so handlers can enqueue messages to be written to stdout. - pub(crate) fn new( + pub(crate) async fn new( outgoing: OutgoingMessageSender, arg0_paths: Arg0DispatchPaths, config: Arc, @@ -59,7 +59,8 @@ impl MessageProcessor { let auth_manager = AuthManager::shared_from_config( config.as_ref(), /*enable_codex_api_key_env*/ false, - ); + ) + .await; let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager, diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index d9e31e7827..8ff0901c1c 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -21,23 +21,16 @@ 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(), + }, + ); if let Ok(header_value) = header_value && let Ok(header) = HeaderValue::from_str(&header_value) diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 7f65e3b049..1e70590e69 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -794,7 +794,8 @@ pub async fn run_main( /*enable_codex_api_key_env*/ false, config_toml.cli_auth_credentials_store.unwrap_or_default(), chatgpt_base_url, - ); + ) + .await; let model_provider_override = if cli.oss { let resolved = resolve_oss_provider( @@ -891,7 +892,9 @@ pub async fn run_main( auth_credentials_store_mode: config.cli_auth_credentials_store_mode, forced_login_method: config.forced_login_method, forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(), - }) { + }) + .await + { eprintln!("{err}"); std::process::exit(1); } @@ -1162,7 +1165,8 @@ async fn run_ratatui_app( /*enable_codex_api_key_env*/ false, initial_config.cli_auth_credentials_store_mode, initial_config.chatgpt_base_url.clone(), - ); + ) + .await; } // If the user made an explicit trust decision, or we showed the login flow, reload config diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index 1e55b5c5d2..cbe1a1fe8b 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -986,7 +986,8 @@ mod tests { /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, "https://chatgpt.com/backend-api/".to_string(), - ), + ) + .await, feedback: codex_feedback::CodexFeedback::new(), log_db: None, environment_manager: Arc::new(