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

@@ -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)?;
     }
 
+    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<codex_core::config::Config> {
     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)
+        .strict_config(strict_config);
-        .build()
+    if enable_workload_identity && is_workload_identity_selected() {
-        .await?)
+        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<codex_login::CodexAuth> {
+) -> anyhow::Result<(Arc<AuthManager>, 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<()> {

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<Box<dyn Future<Output = Result<Request, AuthError>> + Send + 'a>>;
 
+pub type AuthHeadersFuture<'a> =
+    Pin<Box<dyn Future<Output = Result<HeaderMap, AuthError>> + Send + 'a>>;
+
 /// Shared auth handle passed through API clients.
 pub type SharedAuthProvider = Arc<dyn AuthProvider>;
 

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;

codex-rs/exec-server/src/remote.rs

@@ -129,18 +129,20 @@ impl EnvironmentRegistryClient {
         environment_id: &str,
         executor_public_key: &NoiseChannelPublicKey,
     ) -> Result<EnvironmentRegistryRegistrationResponse, ExecServerError> {
-        let response = self
+        let url = endpoint_url(
-            .http
-            .post(endpoint_url(
             &self.base_url,
             &format!("/cloud/environment/{environment_id}/register"),
-            ))
+        );
-            .headers(self.auth_provider.to_auth_headers())
+        let body = EnvironmentRegistryRegistrationRequest {
-            .headers(current_trace_context_headers())
-            .json(&EnvironmentRegistryRegistrationRequest {
             security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(),
             executor_public_key: executor_public_key.clone(),
-            })
+        };
+        let response = self
+            .http
+            .post(url)
+            .headers(self.resolve_auth_headers().await?)
+            .headers(current_trace_context_headers())
+            .json(&body)
             .send()
             .await?;
         let response: EnvironmentRegistryRegistrationResponse =
@@ -185,15 +187,17 @@ impl EnvironmentRegistryClient {
         environment_id: &str,
         harness_public_key: NoiseChannelPublicKey,
     ) -> Result<NoiseRendezvousConnectBundle, ExecServerError> {
-        let response = self
+        let url = endpoint_url(
-            .http
-            .post(endpoint_url(
             &self.base_url,
             &format!("/cloud/environment/{environment_id}/connect"),
-            ))
+        );
-            .headers(self.auth_provider.to_auth_headers())
+        let body = EnvironmentRegistryConnectRequest { harness_public_key };
+        let response = self
+            .http
+            .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<HeaderMap, ExecServerError> {
+        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<R>(&self, response: HttpResponse) -> Result<R, ExecServerError>
     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 response = self
+        let url = endpoint_url(
-            .client
-            .http
-            .post(endpoint_url(
             &self.client.base_url,
             &format!("/cloud/environment/{environment_id}/validate"),
-            ))
+        );
-            .headers(self.client.auth_provider.to_auth_headers())
+        let body = EnvironmentRegistryHarnessKeyValidationRequest {
-            .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(),
-            })
+        };
+        let response = self
+            .client
+            .http
+            .post(url)
+            .headers(self.client.resolve_auth_headers().await?)
+            .headers(current_trace_context_headers())
+            .json(&body)
             .send()
             .await?;
         let status = response.status();
@@ -791,7 +808,11 @@ mod tests {
     struct StaticRegistryAuthProvider;
 
     impl AuthProvider for StaticRegistryAuthProvider {
-        fn add_auth_headers(&self, headers: &mut HeaderMap) {
+        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"),
@@ -800,6 +821,8 @@ mod tests {
                     "ChatGPT-Account-ID",
                     HeaderValue::from_static("workspace-123"),
                 );
+                Ok(headers)
+            })
         }
     }
 

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<CodexAuth> {
+        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
+        let Some(auth) = self.current_auth() else {
-            .auth_manager
-            .auth_cached()
-            .filter(CodexAuth::uses_codex_backend)
-        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"))
         );
 

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()),