diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs index f5f16ab906..e71a57b836 100644 --- a/codex-rs/login/src/auth/workload_identity.rs +++ b/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, +} + +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::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, } #[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, + variable: &'static str, +) -> Result, 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 { @@ -203,6 +246,7 @@ fn invalid_config(message: impl Into) -> WorkloadIdentitySessionError { struct ProcessEnvironment { federation_rule_id: Option, identity_token_file: Option, + workload_identity_context: Option, } 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), } } diff --git a/codex-rs/login/src/auth/workload_identity_tests.rs b/codex-rs/login/src/auth/workload_identity_tests.rs index 15d5115553..e11eca1116 100644 --- a/codex-rs/login/src/auth/workload_identity_tests.rs +++ b/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, ®istry) .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, ®istry), Err(WorkloadIdentitySessionError::ConflictingConfiguration) diff --git a/codex-rs/protocol/src/shell_environment.rs b/codex-rs/protocol/src/shell_environment.rs index c19a21989e..ae3ccb2251 100644 --- a/codex-rs/protocol/src/shell_environment.rs +++ b/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 { diff --git a/codex-rs/protocol/src/shell_environment_tests.rs b/codex-rs/protocol/src/shell_environment_tests.rs index 3b7c5831cc..d6c4009909 100644 --- a/codex-rs/protocol/src/shell_environment_tests.rs +++ b/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!( diff --git a/codex-rs/workload-identity/src/exchange.rs b/codex-rs/workload-identity/src/exchange.rs index bd9606385a..018a1e12c5 100644 --- a/codex-rs/workload-identity/src/exchange.rs +++ b/codex-rs/workload-identity/src/exchange.rs @@ -171,11 +171,17 @@ impl WorkloadIdentityExchange { async fn exchange_uncached(&self) -> Result { 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()) diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs index a75c9decdb..252d8e8d61 100644 --- a/codex-rs/workload-identity/src/lib.rs +++ b/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, } impl WorkloadIdentityConfig { pub fn new( federation_rule_id: String, assertion_file: PathBuf, + workload_identity_context: Option, ) -> Result { 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, }) } } diff --git a/codex-rs/workload-identity/src/workload_identity_tests.rs b/codex-rs/workload-identity/src/workload_identity_tests.rs index da587e2706..8ebd5ce481 100644 --- a/codex-rs/workload-identity/src/workload_identity_tests.rs +++ b/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![ + ("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,