From ed623147c1c79443031ab5eadb87a780b54e7dd3 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Wed, 24 Jun 2026 14:03:28 -0400 Subject: [PATCH] fix(core): preserve brokered credentials across shell snapshots --- codex-rs/core/src/tools/runtimes/mod_tests.rs | 62 ++++++++++++++----- codex-rs/core/src/tools/runtimes/shell.rs | 17 ++++- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index 4ce62fbd2d..f090171a65 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -7,7 +7,6 @@ use crate::tools::sandboxing::SandboxAttempt; use crate::tools::sandboxing::managed_network_for_sandbox_permissions; #[cfg(target_os = "macos")] use codex_network_proxy::CODEX_PROXY_GIT_SSH_COMMAND_MARKER; -use codex_network_proxy::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; use codex_network_proxy::CUSTOM_CA_ENV_KEYS; use codex_network_proxy::ConfigReloader; use codex_network_proxy::ConfigReloaderFuture; @@ -65,7 +64,7 @@ fn shell_with_snapshot( async fn test_network_proxy() -> anyhow::Result { let mut config = NetworkProxyConfig::default(); - config.set_credential_broker_enabled(true); + config.set_credential_broker_enabled(/*enabled*/ true); let state = codex_network_proxy::build_config_state(config, NetworkProxyConstraints::default())?; NetworkProxy::builder() @@ -625,8 +624,9 @@ fn maybe_wrap_shell_lc_with_snapshot_unsets_absent_permission_profile() { assert_eq!(output.stdout, b""); } -#[test] -fn maybe_wrap_shell_lc_with_snapshot_restores_proxy_env_from_process_env() { +#[tokio::test] +async fn snapshot_wrapper_restores_prepared_broker_credentials() -> anyhow::Result<()> { + let proxy = test_network_proxy().await?; let dir = tempdir().expect("create temp dir"); let snapshot_path = dir.path().join("snapshot.sh"); std::fs::write( @@ -647,35 +647,63 @@ fn maybe_wrap_shell_lc_with_snapshot_restores_proxy_env_from_process_env() { "printf '%s\\n%s\\n%s\\n%s\\n%s' \"$PIP_PROXY\" \"$HTTP_PROXY\" \"$http_proxy\" \"$GIT_SSH_COMMAND\" \"$OPENAI_API_KEY\"" .to_string(), ]; + let real_openai_api_key = + "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_"; + let mut env = proxy + .prepare_for_optional_environment( + HashMap::from([( + "OPENAI_API_KEY".to_string(), + real_openai_api_key.to_string(), + )]), + /*environment_id*/ None, + )? + .env; + let brokered_openai_api_key = env + .get("OPENAI_API_KEY") + .expect("brokered OpenAI API key") + .clone(); + assert_ne!(brokered_openai_api_key, real_openai_api_key); + env.extend([ + ("PIP_PROXY".to_string(), "http://127.0.0.1:4321".to_string()), + ( + "HTTP_PROXY".to_string(), + "http://127.0.0.1:4321".to_string(), + ), + ( + "http_proxy".to_string(), + "http://127.0.0.1:4321".to_string(), + ), + ( + "GIT_SSH_COMMAND".to_string(), + "ssh -o ProxyCommand=fresh".to_string(), + ), + ]); let rewritten = maybe_wrap_shell_lc_with_snapshot( &command, &session_shell, Some(&shell_snapshot), &HashMap::new(), - &HashMap::from([(CREDENTIAL_BROKER_ACTIVE_ENV_KEY.into(), "1".into())]), + &env, &RuntimePathPrepends::default(), ); let output = Command::new(&rewritten[0]) .args(&rewritten[1..]) - .env(PROXY_ACTIVE_ENV_KEY, "1") - .env(CREDENTIAL_BROKER_ACTIVE_ENV_KEY, "1") - .env("PIP_PROXY", "http://127.0.0.1:4321") - .env("HTTP_PROXY", "http://127.0.0.1:4321") - .env("http_proxy", "http://127.0.0.1:4321") - .env("GIT_SSH_COMMAND", "ssh -o ProxyCommand=fresh") - .env("OPENAI_API_KEY", "sk-brokered") + .envs(&env) .output() .expect("run rewritten command"); assert!(output.status.success(), "command failed: {output:?}"); assert_eq!( String::from_utf8_lossy(&output.stdout), - "http://127.0.0.1:4321\n\ - http://127.0.0.1:4321\n\ - http://127.0.0.1:4321\n\ - ssh -o ProxyCommand=stale\n\ - sk-brokered" + format!( + "http://127.0.0.1:4321\n\ + http://127.0.0.1:4321\n\ + http://127.0.0.1:4321\n\ + ssh -o ProxyCommand=stale\n\ + {brokered_openai_api_key}" + ) ); + Ok(()) } #[test] diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index 3ff50f3f51..29fd0d3136 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -40,6 +40,7 @@ use crate::tools::sandboxing::managed_network_for_sandbox_permissions; use crate::tools::sandboxing::sandbox_permissions_preserving_denied_reads; use crate::tools::sandboxing::with_cached_approval; use codex_network_proxy::NetworkProxy; +use codex_protocol::error::CodexErr; use codex_protocol::exec_output::ExecToolCallOutput; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::ReviewDecision; @@ -48,6 +49,7 @@ use codex_shell_command::powershell::prefix_powershell_script_with_utf8; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::collections::HashMap; +use std::io; use tokio_util::sync::CancellationToken; #[derive(Clone, Debug)] @@ -258,7 +260,20 @@ impl ToolRuntime for ShellRuntime { ); let managed_network = managed_network_for_sandbox_permissions(req.network.as_ref(), sandbox_permissions); - let env = exec_env_for_sandbox_permissions(&req.env, sandbox_permissions); + let mut env = exec_env_for_sandbox_permissions(&req.env, sandbox_permissions); + if let Some(network) = managed_network { + network + .apply_to_env_for_optional_environment( + &mut env, + Some(&req.turn_environment.environment_id), + ) + .map_err(|err| { + ToolError::Codex(CodexErr::Io(io::Error::other(format!( + "failed to prepare network proxy for environment `{}`: {err}", + req.turn_environment.environment_id + )))) + })?; + } let explicit_env_overrides = req.explicit_env_overrides.clone(); #[cfg(unix)] let (env, runtime_path_prepends) = {