Forward workload identity context during token exchange (#38767)

## What changed

- Read optional workload identity context from `OPENAI_WORKLOAD_IDENTITY_CONTEXT` and forward it unchanged as the `workload_identity_context` token exchange field.
- Treat the context as sensitive by redacting it from session configuration debug output and removing it from model-reachable child environments.
- Include the context in workload identity session fingerprints so sessions with different values cannot share an exchange.

## Testing

- Cover request forwarding, debug redaction, session compatibility, and child-environment scrubbing.

GitOrigin-RevId: fb50700478cf54d9d604944a4ed3e77acc928a0f

codex-rs/login/src/auth/workload_identity.rs

@@ -11,6 +11,7 @@ use codex_http_client::HttpClientFactory;
 use codex_protocol::config_types::ForcedLoginMethod;
 use codex_protocol::shell_environment::OPENAI_FEDERATION_RULE_ID_ENV_VAR;
 use codex_protocol::shell_environment::OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR;
+use codex_protocol::shell_environment::OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR;
 use codex_workload_identity::WorkloadIdentityConfig;
 use codex_workload_identity::WorkloadIdentityError;
 use codex_workload_identity::WorkloadIdentityExchange;
@@ -47,13 +48,34 @@ impl WorkloadIdentityEnvironment {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 struct WorkloadIdentitySessionConfig {
     assertion_file: PathBuf,
     environment: WorkloadIdentityEnvironment,
     federation_rule_id: String,
     http_client_factory: HttpClientFactory,
     token_url: Url,
+    workload_identity_context: Option<String>,
+}
+
+impl std::fmt::Debug for WorkloadIdentitySessionConfig {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("WorkloadIdentitySessionConfig")
+            .field("assertion_file", &self.assertion_file)
+            .field("environment", &self.environment)
+            .field("federation_rule_id", &self.federation_rule_id)
+            .field("http_client_factory", &self.http_client_factory)
+            .field("token_url", &self.token_url)
+            .field(
+                "workload_identity_context",
+                &self
+                    .workload_identity_context
+                    .as_ref()
+                    .map(|_| "[redacted]"),
+            )
+            .finish()
+    }
 }
 
 impl WorkloadIdentitySessionConfig {
@@ -63,15 +85,17 @@ impl WorkloadIdentitySessionConfig {
             environment: self.environment,
             federation_rule_id: self.federation_rule_id.trim().to_string(),
             token_url: self.token_url.to_string(),
+            workload_identity_context: self.workload_identity_context.clone(),
         }
     }
 
     fn into_exchange(self) -> Result<WorkloadIdentityExchange, WorkloadIdentityError> {
-        WorkloadIdentityExchange::new(
-            WorkloadIdentityConfig::new(self.federation_rule_id, self.assertion_file)?,
-            self.token_url,
-            self.http_client_factory,
-        )
+        let config = WorkloadIdentityConfig::new(
+            self.federation_rule_id,
+            self.assertion_file,
+            self.workload_identity_context,
+        )?;
+        WorkloadIdentityExchange::new(config, self.token_url, self.http_client_factory)
     }
 }
 
@@ -81,6 +105,7 @@ struct WorkloadIdentityFingerprint {
     environment: WorkloadIdentityEnvironment,
     federation_rule_id: String,
     token_url: String,
+    workload_identity_context: Option<String>,
 }
 
 #[derive(Debug, Error)]
@@ -126,6 +151,10 @@ fn resolve_config(
         environment.identity_token_file,
         OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR,
     )?;
+    let workload_identity_context = optional_unicode(
+        environment.workload_identity_context,
+        OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR,
+    )?;
     let auth_environment = classify_auth_environment(chatgpt_base_url)?;
 
     Ok(Some(WorkloadIdentitySessionConfig {
@@ -134,6 +163,7 @@ fn resolve_config(
         federation_rule_id,
         http_client_factory: auth_route_config.http_client_factory().clone(),
         token_url: auth_environment.token_url()?,
+        workload_identity_context,
     }))
 }
 
@@ -169,6 +199,19 @@ fn required_unicode(
     Ok(value.to_string())
 }
 
+fn optional_unicode(
+    value: Option<OsString>,
+    variable: &'static str,
+) -> Result<Option<String>, WorkloadIdentitySessionError> {
+    value
+        .map(|value| {
+            value.into_string().map_err(|_| {
+                invalid_config(format!("workload identity variable {variable} is invalid"))
+            })
+        })
+        .transpose()
+}
+
 fn classify_auth_environment(
     base_url: &str,
 ) -> Result<WorkloadIdentityEnvironment, WorkloadIdentitySessionError> {
@@ -203,6 +246,7 @@ fn invalid_config(message: impl Into<String>) -> WorkloadIdentitySessionError {
 struct ProcessEnvironment {
     federation_rule_id: Option<OsString>,
     identity_token_file: Option<OsString>,
+    workload_identity_context: Option<OsString>,
 }
 
 impl ProcessEnvironment {
@@ -210,6 +254,7 @@ impl ProcessEnvironment {
         Self {
             federation_rule_id: std::env::var_os(OPENAI_FEDERATION_RULE_ID_ENV_VAR),
             identity_token_file: std::env::var_os(OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR),
+            workload_identity_context: std::env::var_os(OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR),
         }
     }
 

codex-rs/login/src/auth/workload_identity_tests.rs

@@ -21,6 +21,7 @@ fn complete_environment() -> ProcessEnvironment {
     ProcessEnvironment {
         federation_rule_id: Some("rule-one".into()),
         identity_token_file: Some(std::env::temp_dir().join("identity-token").into_os_string()),
+        workload_identity_context: None,
     }
 }
 
@@ -48,6 +49,18 @@ fn markers_select_wif_and_partial_configuration_fails_closed() {
         .expect("no markers")
         .is_none()
     );
+    assert!(
+        resolve_for_test(
+            ProcessEnvironment {
+                workload_identity_context: Some(r#"{"instance_id":"box-one"}"#.into()),
+                ..ProcessEnvironment::default()
+            },
+            /*chatgpt_login_allowed*/ true,
+            "https://chatgpt.com/backend-api",
+        )
+        .expect("context alone is not a WIF marker")
+        .is_none()
+    );
     for (environment, missing) in [
         (
             ProcessEnvironment {
@@ -131,6 +144,24 @@ fn auth_policy_and_app_environment_are_enforced() {
     assert!(error.to_string().contains("app routing"));
 }
 
+#[test]
+fn workload_context_is_preserved_without_logging_its_value() {
+    let context = r#"{"instance_id":"box-one"}"#;
+    let config = resolve_for_test(
+        ProcessEnvironment {
+            workload_identity_context: Some(context.into()),
+            ..complete_environment()
+        },
+        /*chatgpt_login_allowed*/ true,
+        "https://chatgpt.com/backend-api",
+    )
+    .expect("valid configuration")
+    .expect("WIF selected");
+
+    assert_eq!(config.workload_identity_context.as_deref(), Some(context));
+    assert!(!format!("{config:?}").contains(context));
+}
+
 fn session_config(directory: &Path, server: &MockServer) -> WorkloadIdentitySessionConfig {
     let assertion_file = directory.join("identity-token");
     std::fs::write(&assertion_file, "assertion-one").expect("write assertion");
@@ -140,6 +171,7 @@ fn session_config(directory: &Path, server: &MockServer) -> WorkloadIdentitySess
         federation_rule_id: "rule-one".to_string(),
         http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
         token_url: Url::parse(&format!("{}/oauth/token", server.uri())).expect("token URL"),
+        workload_identity_context: None,
     }
 }
 
@@ -189,7 +221,9 @@ async fn compatible_adapters_share_exchange() {
         .mount(&server)
         .await;
     let registry = WorkloadIdentitySessionRegistry::default();
-    let first_config = session_config(temp_dir.path(), &server);
+    let context = r#"{"instance_id":"box-one"}"#;
+    let mut first_config = session_config(temp_dir.path(), &server);
+    first_config.workload_identity_context = Some(context.into());
     let second_config = first_config.clone();
     let first = WorkloadIdentityExternalAuth::from_config_with_registry(first_config, &registry)
         .expect("first adapter");
@@ -211,6 +245,15 @@ async fn compatible_adapters_share_exchange() {
             .get_token()
             .expect("second token")
     );
+
+    let requests = server.received_requests().await.expect("requests");
+    assert_eq!(requests.len(), 1);
+    assert_eq!(
+        url::form_urlencoded::parse(&requests[0].body)
+            .find(|(name, _)| name == "workload_identity_context")
+            .map(|(_, value)| value.into_owned()),
+        Some(context.to_string())
+    );
 }
 
 #[tokio::test]
@@ -229,6 +272,8 @@ async fn incompatible_process_session_settings_are_rejected() {
     std::fs::write(&different_file.assertion_file, "assertion-two").expect("write assertion");
     let mut different_environment = base.clone();
     different_environment.environment = WorkloadIdentityEnvironment::Production;
+    let mut different_context = base.clone();
+    different_context.workload_identity_context = Some(r#"{"instance_id":"box-two"}"#.into());
     let mut different_route = base;
     different_route.http_client_factory =
         HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy);
@@ -241,7 +286,12 @@ async fn incompatible_process_session_settings_are_rejected() {
         &different_route_adapter.session
     ));
 
-    for config in [different_rule, different_file, different_environment] {
+    for config in [
+        different_rule,
+        different_file,
+        different_environment,
+        different_context,
+    ] {
         assert!(matches!(
             WorkloadIdentityExternalAuth::from_config_with_registry(config, &registry),
             Err(WorkloadIdentitySessionError::ConflictingConfiguration)

codex-rs/protocol/src/shell_environment.rs

@@ -7,11 +7,13 @@ pub const CODEX_SESSION_ID_ENV_VAR: &str = "CODEX_SESSION_ID";
 pub const CODEX_THREAD_ID_ENV_VAR: &str = "CODEX_THREAD_ID";
 pub const OPENAI_FEDERATION_RULE_ID_ENV_VAR: &str = "OPENAI_FEDERATION_RULE_ID";
 pub const OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR: &str = "OPENAI_IDENTITY_TOKEN_FILE";
+pub const OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR: &str = "OPENAI_WORKLOAD_IDENTITY_CONTEXT";
 
 /// Environment variables that model-reachable child processes must not inherit.
 pub const NON_INHERITABLE_ENV_VARS: &[&str] = &[
     OPENAI_FEDERATION_RULE_ID_ENV_VAR,
     OPENAI_IDENTITY_TOKEN_FILE_ENV_VAR,
+    OPENAI_WORKLOAD_IDENTITY_CONTEXT_ENV_VAR,
 ];
 
 pub fn is_non_inheritable_env_var(name: &str) -> bool {

codex-rs/protocol/src/shell_environment_tests.rs

@@ -17,6 +17,10 @@ fn non_inheritable_environment_is_removed_after_policy_overrides() {
             "openai_federation_rule_id".to_string(),
             "inherited-rule".to_string(),
         ),
+        (
+            "OPENAI_WORKLOAD_IDENTITY_CONTEXT".to_string(),
+            r#"{"instance_id":"box-one"}"#.to_string(),
+        ),
     ];
     let policy = ShellEnvironmentPolicy {
         inherit: ShellEnvironmentPolicyInherit::All,
@@ -44,6 +48,10 @@ fn command_scrubber_removes_names_from_real_child_environment() {
             .args([TEST_NAME, "--exact", "--nocapture"])
             .env(CHILD_MODE_ENV_VAR, "1")
             .env("OpenAI_Federation_Rule_Id", "inherited-rule")
+            .env(
+                "OpenAI_Workload_Identity_Context",
+                r#"{"instance_id":"box-one"}"#,
+            )
             .output()
             .expect("run inherited-environment test process");
         assert!(

codex-rs/workload-identity/src/exchange.rs

@@ -171,11 +171,17 @@ impl WorkloadIdentityExchange {
 
     async fn exchange_uncached(&self) -> Result<WorkloadIdentityToken, WorkloadIdentityError> {
         let assertion = read_assertion(&self.config.assertion_file).await?;
-        let body = url::form_urlencoded::Serializer::new(String::new())
-            .append_pair("grant_type", JWT_BEARER_GRANT_TYPE)
-            .append_pair("assertion", &assertion)
-            .append_pair("federation_rule_id", &self.config.federation_rule_id)
-            .finish();
+        let body = {
+            let mut serializer = url::form_urlencoded::Serializer::new(String::new());
+            serializer
+                .append_pair("grant_type", JWT_BEARER_GRANT_TYPE)
+                .append_pair("assertion", &assertion)
+                .append_pair("federation_rule_id", &self.config.federation_rule_id);
+            if let Some(context) = &self.config.workload_identity_context {
+                serializer.append_pair("workload_identity_context", context);
+            }
+            serializer.finish()
+        };
         let response = self
             .client
             .post(self.token_url.as_str())

codex-rs/workload-identity/src/lib.rs

@@ -13,12 +13,14 @@ use thiserror::Error;
 pub struct WorkloadIdentityConfig {
     pub(crate) assertion_file: PathBuf,
     pub(crate) federation_rule_id: String,
+    pub(crate) workload_identity_context: Option<String>,
 }
 
 impl WorkloadIdentityConfig {
     pub fn new(
         federation_rule_id: String,
         assertion_file: PathBuf,
+        workload_identity_context: Option<String>,
     ) -> Result<Self, WorkloadIdentityError> {
         let federation_rule_id = federation_rule_id.trim();
         if federation_rule_id.is_empty() {
@@ -30,6 +32,7 @@ impl WorkloadIdentityConfig {
         Ok(Self {
             assertion_file,
             federation_rule_id: federation_rule_id.to_string(),
+            workload_identity_context,
         })
     }
 }

codex-rs/workload-identity/src/workload_identity_tests.rs

@@ -32,7 +32,12 @@ fn assertion_file(assertion: &str) -> (TempDir, PathBuf) {
 
 fn make_exchange(path: PathBuf, server: &MockServer) -> WorkloadIdentityExchange {
     WorkloadIdentityExchange::new(
-        WorkloadIdentityConfig::new("idpm_rule_one".to_string(), path).expect("valid config"),
+        WorkloadIdentityConfig::new(
+            "idpm_rule_one".to_string(),
+            path,
+            /*workload_identity_context*/ None,
+        )
+        .expect("valid config"),
         Url::parse(&format!("{}/oauth/token", server.uri())).expect("valid token URL"),
         HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
     )
@@ -97,6 +102,50 @@ async fn exchange_sends_three_field_contract_and_caches_valid_response() {
     );
 }
 
+#[tokio::test]
+async fn exchange_forwards_optional_workload_context_without_parsing() {
+    let server = MockServer::start().await;
+    Mock::given(method("POST"))
+        .and(path("/oauth/token"))
+        .and(header("content-type", "application/x-www-form-urlencoded"))
+        .respond_with(success("sensitive-access-token", /*expires_in*/ 600))
+        .mount(&server)
+        .await;
+    let (_temp_dir, assertion_path) = assertion_file("assertion-one");
+    let context = "server-validates-this-raw-value";
+    let config = WorkloadIdentityConfig::new(
+        "idpm_rule_one".to_string(),
+        assertion_path,
+        /*workload_identity_context*/ Some(context.to_string()),
+    )
+    .expect("valid config");
+    let exchange = WorkloadIdentityExchange::new(
+        config,
+        Url::parse(&format!("{}/oauth/token", server.uri())).expect("valid token URL"),
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+    )
+    .expect("valid exchange");
+
+    exchange.resolve().await.expect("exchange");
+
+    let requests = server.received_requests().await.expect("requests");
+    assert_eq!(requests.len(), 1);
+    assert_eq!(
+        url::form_urlencoded::parse(&requests[0].body)
+            .into_owned()
+            .collect::<Vec<_>>(),
+        vec![
+            ("grant_type".to_string(), JWT_BEARER_GRANT_TYPE.to_string()),
+            ("assertion".to_string(), "assertion-one".to_string()),
+            (
+                "federation_rule_id".to_string(),
+                "idpm_rule_one".to_string()
+            ),
+            ("workload_identity_context".to_string(), context.to_string()),
+        ]
+    );
+}
+
 #[tokio::test]
 async fn concurrent_resolve_and_rejected_token_refresh_are_single_flight() {
     let server = MockServer::start().await;
@@ -264,13 +313,21 @@ async fn transient_proactive_refresh_failure_uses_still_valid_token() {
 #[test]
 fn configuration_requires_an_absolute_file_and_secure_token_url() {
     assert!(matches!(
-        WorkloadIdentityConfig::new("idpm_rule_one".to_string(), PathBuf::from("relative.jwt")),
+        WorkloadIdentityConfig::new(
+            "idpm_rule_one".to_string(),
+            PathBuf::from("relative.jwt"),
+            /*workload_identity_context*/ None,
+        ),
         Err(WorkloadIdentityError::AssertionFileMustBeAbsolute)
     ));
 
     let (_temp_dir, assertion_path) = assertion_file("assertion-one");
-    let config = WorkloadIdentityConfig::new("idpm_rule_one".to_string(), assertion_path)
-        .expect("valid config");
+    let config = WorkloadIdentityConfig::new(
+        "idpm_rule_one".to_string(),
+        assertion_path,
+        /*workload_identity_context*/ None,
+    )
+    .expect("valid config");
     assert!(matches!(
         WorkloadIdentityExchange::new(
             config,