From fb5aa093e0657c940e8ffd1dd055d22b92d67dfc Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Fri, 14 Aug 2026 17:22:54 +0000 Subject: [PATCH] Support workload identity in remote exec-server auth (#38610) ## Why Remote exec-server registry requests need to refresh managed credentials before sending a request. Static auth-header resolution cannot perform the asynchronous token exchange required by workload identity. ## What changed - Add asynchronous auth-header resolution to `AuthProvider`, with the existing static-header behavior as the default. - Resolve fresh managed credentials for each remote environment registry request while preserving the expected account and workspace identity. - Load the cloud configuration bundle during remote exec-server startup when workload identity is selected. ## Testing - Update the managed-auth and environment-registry auth tests to exercise asynchronous header resolution. GitOrigin-RevId: 5d60f1127467aaacdb5d1a8f3d92278bc4bf2e29 --- codex-rs/cli/src/main.rs | 65 +++++++++++++--- codex-rs/codex-api/src/auth.rs | 14 +++- codex-rs/codex-api/src/lib.rs | 1 + codex-rs/exec-server/src/remote.rs | 91 ++++++++++++++-------- codex-rs/model-provider/src/auth.rs | 52 +++++++++---- codex-rs/workload-identity/src/exchange.rs | 1 - 6 files changed, 163 insertions(+), 61 deletions(-) diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 4cfe030484..102f40aec6 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -18,6 +18,7 @@ use codex_cli::run_login_with_api_key; use codex_cli::run_login_with_chatgpt; use codex_cli::run_login_with_device_code; use codex_cli::run_logout; +use codex_cloud_config::cloud_config_bundle_loader_for_storage; use codex_cloud_tasks::Cli as CloudTasksCli; use codex_exec::Cli as ExecCli; use codex_exec::Command as ExecCommand; @@ -72,9 +73,12 @@ use codex_config::LoaderOverrides; use codex_core::build_models_manager; use codex_core::config::Config; use codex_core::config::ConfigBuilder; +use codex_core::config::ConfigLoadOptions; use codex_core::config::ConfigOverrides; +use codex_core::config::bootstrap_auth_config; use codex_core::config::edit::ConfigEditsBuilder; use codex_core::config::find_codex_home; +use codex_core::config::load_config_toml_with_layer_stack; use codex_core::config::resolve_profile_v2_config_path; use codex_features::FEATURES; use codex_features::Stage; @@ -82,6 +86,7 @@ use codex_features::is_known_feature_key; use codex_home::CodexHomeUserInstructionsProvider; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::is_workload_identity_selected; use codex_login::read_codex_access_token_from_env; use codex_memories_write::clear_memory_roots_contents; use codex_models_manager::bundled_models_response; @@ -1735,7 +1740,12 @@ async fn run_exec_server_command( let environment_id = cmd .environment_id .ok_or_else(|| anyhow::anyhow!("--environment-id is required when --remote is set"))?; - let config = load_exec_server_config(root_config_overrides, strict_config).await?; + let config = load_exec_server_config( + root_config_overrides, + strict_config, + /*enable_workload_identity*/ true, + ) + .await?; let (_otel, telemetry) = exec_server_telemetry::init(Some(&config)); let auth_provider = load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth) @@ -1774,7 +1784,12 @@ async fn run_exec_server_command( .await?; Ok(()) } else { - let config_result = load_exec_server_config(root_config_overrides, strict_config).await; + let config_result = load_exec_server_config( + root_config_overrides, + strict_config, + /*enable_workload_identity*/ false, + ) + .await; let config = if strict_config { Some(config_result?) } else { @@ -1833,7 +1848,7 @@ async fn load_exec_server_remote_auth_provider( return Ok(codex_model_provider::auth_provider_from_auth(&auth)); } - let auth = load_exec_server_remote_auth( + let (auth_manager, auth) = load_exec_server_remote_auth( config, "remote exec-server registration requires ChatGPT authentication or API key authentication; run `codex login` or set CODEX_API_KEY", ) @@ -1849,7 +1864,14 @@ async fn load_exec_server_remote_auth_provider( validate_api_key_remote_host(base_url)?; } - Ok(codex_model_provider::auth_provider_from_auth(&auth)) + if auth_manager.is_workload_identity_selected() { + Ok(codex_model_provider::auth_provider_from_auth_manager( + auth_manager, + &auth, + )) + } else { + Ok(codex_model_provider::auth_provider_from_auth(&auth)) + } } fn is_supported_exec_server_remote_auth(auth: &CodexAuth) -> bool { @@ -1893,21 +1915,44 @@ fn validate_api_key_remote_host(base_url: &str) -> anyhow::Result<()> { async fn load_exec_server_config( root_config_overrides: &CliConfigOverrides, strict_config: bool, + enable_workload_identity: bool, ) -> anyhow::Result { let cli_kv_overrides = root_config_overrides .parse_overrides() .map_err(anyhow::Error::msg)?; - Ok(ConfigBuilder::default() + let bootstrap_cli_overrides = cli_kv_overrides.clone(); + let mut builder = ConfigBuilder::default() .cli_overrides(cli_kv_overrides) - .strict_config(strict_config) - .build() - .await?) + .strict_config(strict_config); + if enable_workload_identity && is_workload_identity_selected() { + let codex_home = find_codex_home()?; + let bootstrap_cwd = AbsolutePathBuf::current_dir()?; + let bootstrap_config = load_config_toml_with_layer_stack( + &codex_home, + Some(&bootstrap_cwd), + bootstrap_cli_overrides, + ConfigLoadOptions { + loader_overrides: LoaderOverrides::default(), + strict_config, + cloud_config_bundle: Default::default(), + }, + ) + .await?; + let bootstrap_auth_config = bootstrap_auth_config(&codex_home, &bootstrap_config)?; + let cloud_config_bundle = cloud_config_bundle_loader_for_storage( + bootstrap_auth_config, + /*enable_codex_api_key_env*/ false, + ) + .await?; + builder = builder.cloud_config_bundle(cloud_config_bundle); + } + Ok(builder.build().await?) } async fn load_exec_server_remote_auth( config: &codex_core::config::Config, missing_auth_error: &'static str, -) -> anyhow::Result { +) -> anyhow::Result<(Arc, codex_login::CodexAuth)> { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ true).await?; @@ -1922,7 +1967,7 @@ async fn load_exec_server_remote_auth( } }; - Ok(auth) + Ok((auth_manager, auth)) } async fn enable_feature_in_config(feature: &str) -> anyhow::Result<()> { diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index b889c359e5..5cd8007af0 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -41,6 +41,15 @@ pub trait AuthProvider: Send + Sync { headers } + /// Resolves auth headers for an outbound request. + /// + /// Unlike [`Self::to_auth_headers`], implementations may perform asynchronous work to refresh + /// credentials before returning. Header-only providers with static credentials can rely on the + /// default implementation. + fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { + Box::pin(async { Ok(self.to_auth_headers()) }) + } + /// Applies auth to a complete outbound request and returns the request to send. /// /// The input `request` is moved into this method. Implementations may mutate @@ -55,7 +64,7 @@ pub trait AuthProvider: Send + Sync { fn apply_auth(&self, request: Request) -> AuthProviderFuture<'_> { Box::pin(async move { let mut request = request; - self.add_auth_headers(&mut request.headers); + request.headers.extend(self.resolve_auth_headers().await?); Ok(request) }) } @@ -64,6 +73,9 @@ pub trait AuthProvider: Send + Sync { pub type AuthProviderFuture<'a> = Pin> + Send + 'a>>; +pub type AuthHeadersFuture<'a> = + Pin> + Send + 'a>>; + /// Shared auth handle passed through API clients. pub type SharedAuthProvider = Arc; diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 98ce06817a..03055b75bb 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -22,6 +22,7 @@ pub use crate::api_bridge::map_api_error; pub use crate::auth::AgentIdentityTelemetry; pub use crate::auth::AuthError; pub use crate::auth::AuthHeaderTelemetry; +pub use crate::auth::AuthHeadersFuture; pub use crate::auth::AuthProvider; pub use crate::auth::AuthProviderFuture; pub use crate::auth::SharedAuthProvider; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index eb26e69fb1..16deaa092c 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -129,18 +129,20 @@ impl EnvironmentRegistryClient { environment_id: &str, executor_public_key: &NoiseChannelPublicKey, ) -> Result { + let url = endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/register"), + ); + let body = EnvironmentRegistryRegistrationRequest { + security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), + executor_public_key: executor_public_key.clone(), + }; let response = self .http - .post(endpoint_url( - &self.base_url, - &format!("/cloud/environment/{environment_id}/register"), - )) - .headers(self.auth_provider.to_auth_headers()) + .post(url) + .headers(self.resolve_auth_headers().await?) .headers(current_trace_context_headers()) - .json(&EnvironmentRegistryRegistrationRequest { - security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), - executor_public_key: executor_public_key.clone(), - }) + .json(&body) .send() .await?; let response: EnvironmentRegistryRegistrationResponse = @@ -185,15 +187,17 @@ impl EnvironmentRegistryClient { environment_id: &str, harness_public_key: NoiseChannelPublicKey, ) -> Result { + let url = endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/connect"), + ); + let body = EnvironmentRegistryConnectRequest { harness_public_key }; let response = self .http - .post(endpoint_url( - &self.base_url, - &format!("/cloud/environment/{environment_id}/connect"), - )) - .headers(self.auth_provider.to_auth_headers()) + .post(url) + .headers(self.resolve_auth_headers().await?) .headers(current_trace_context_headers()) - .json(&EnvironmentRegistryConnectRequest { harness_public_key }) + .json(&body) .timeout(self.connect_timeout) .send() .await?; @@ -227,6 +231,17 @@ impl EnvironmentRegistryClient { }) } + async fn resolve_auth_headers(&self) -> Result { + self.auth_provider + .resolve_auth_headers() + .await + .map_err(|error| { + ExecServerError::EnvironmentRegistryAuth(format!( + "failed to resolve environment registry authentication: {error}" + )) + }) + } + async fn parse_json_response(&self, response: HttpResponse) -> Result where R: for<'de> Deserialize<'de>, @@ -275,20 +290,22 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { authorization: &str, ) -> Result<(), ExecServerError> { let environment_id = &self.environment_id; + let url = endpoint_url( + &self.client.base_url, + &format!("/cloud/environment/{environment_id}/validate"), + ); + let body = EnvironmentRegistryHarnessKeyValidationRequest { + executor_registration_id: self.executor_registration_id.clone(), + harness_public_key: harness_public_key.clone(), + harness_key_authorization: authorization.to_string(), + }; let response = self .client .http - .post(endpoint_url( - &self.client.base_url, - &format!("/cloud/environment/{environment_id}/validate"), - )) - .headers(self.client.auth_provider.to_auth_headers()) + .post(url) + .headers(self.client.resolve_auth_headers().await?) .headers(current_trace_context_headers()) - .json(&EnvironmentRegistryHarnessKeyValidationRequest { - executor_registration_id: self.executor_registration_id.clone(), - harness_public_key: harness_public_key.clone(), - harness_key_authorization: authorization.to_string(), - }) + .json(&body) .send() .await?; let status = response.status(); @@ -791,15 +808,21 @@ mod tests { struct StaticRegistryAuthProvider; impl AuthProvider for StaticRegistryAuthProvider { - fn add_auth_headers(&self, headers: &mut HeaderMap) { - let _ = headers.insert( - http::header::AUTHORIZATION, - HeaderValue::from_static("Bearer registry-token"), - ); - let _ = headers.insert( - "ChatGPT-Account-ID", - HeaderValue::from_static("workspace-123"), - ); + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn resolve_auth_headers(&self) -> codex_api::AuthHeadersFuture<'_> { + Box::pin(async { + let mut headers = HeaderMap::new(); + let _ = headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer registry-token"), + ); + let _ = headers.insert( + "ChatGPT-Account-ID", + HeaderValue::from_static("workspace-123"), + ); + Ok(headers) + }) } } diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 85d01ae216..3f342dbf52 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -5,6 +5,8 @@ use std::sync::atomic::Ordering; use codex_agent_identity::AgentIdentityKey; use codex_agent_identity::authorization_header_for_agent_task; use codex_api::AgentIdentityTelemetry; +use codex_api::AuthError; +use codex_api::AuthHeadersFuture; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; use codex_login::AuthHeaders; @@ -128,26 +130,42 @@ struct AuthManagerAuthProvider { expected_auth: CodexAuth, } +impl AuthManagerAuthProvider { + fn is_expected_auth(&self, auth: &CodexAuth) -> bool { + auth.uses_codex_backend() + && auth.get_account_id() == self.expected_auth.get_account_id() + && auth.get_chatgpt_user_id() == self.expected_auth.get_chatgpt_user_id() + && auth.is_workspace_account() == self.expected_auth.is_workspace_account() + } + + fn current_auth(&self) -> Option { + self.auth_manager + .auth_cached() + .filter(|auth| self.is_expected_auth(auth)) + } +} + impl AuthProvider for AuthManagerAuthProvider { fn add_auth_headers(&self, headers: &mut HeaderMap) { - let Some(auth) = self - .auth_manager - .auth_cached() - .filter(CodexAuth::uses_codex_backend) - else { + let Some(auth) = self.current_auth() else { return; }; - // The caller's account-scoped state was built for the expected - // identity. Follow token refreshes for that identity, but never cross - // an account or workspace boundary without rebuilding that state. - if auth.get_account_id() != self.expected_auth.get_account_id() - || auth.get_chatgpt_user_id() != self.expected_auth.get_chatgpt_user_id() - || auth.is_workspace_account() != self.expected_auth.is_workspace_account() - { - return; - } auth_provider_from_auth(&auth).add_auth_headers(headers); } + + fn resolve_auth_headers(&self) -> AuthHeadersFuture<'_> { + Box::pin(async move { + let auth = self + .auth_manager + .auth() + .await + .filter(|auth| self.is_expected_auth(auth)) + .ok_or_else(|| { + AuthError::Transient("managed authentication is unavailable".to_string()) + })?; + Ok(auth_provider_from_auth(&auth).to_auth_headers()) + }) + } } // Some providers are meant to send no auth headers. Examples include local OSS @@ -523,8 +541,12 @@ mod tests { .expect("save reloaded auth"); auth_manager.reload().await; + let resolved_headers = provider + .resolve_auth_headers() + .await + .expect("managed auth headers should resolve"); assert_eq!( - provider.to_auth_headers().get(AUTHORIZATION), + resolved_headers.get(AUTHORIZATION), Some(&HeaderValue::from_static("Bearer header.e30.reloaded")) ); diff --git a/codex-rs/workload-identity/src/exchange.rs b/codex-rs/workload-identity/src/exchange.rs index 0427eed221..bd9606385a 100644 --- a/codex-rs/workload-identity/src/exchange.rs +++ b/codex-rs/workload-identity/src/exchange.rs @@ -96,7 +96,6 @@ impl WorkloadIdentityExchange { return Ok(token); } } - let valid_from = Instant::now(); let result = match self.exchange_uncached().await { Ok(token) => state.store(token, valid_from, Instant::now()),