From 272a5aeb653d1ca7f01aead112ebc4d6b7641c57 Mon Sep 17 00:00:00 2001
From: Winston Howes <winston@openai.com>
Date: Tue, 7 Jul 2026 11:37:49 -0700
Subject: [PATCH] fix(network-proxy): preserve broker context across snapshots

---
 codex-rs/core/src/tools/runtimes/mod_tests.rs | 44 +++++++++++++++++++
 .../network-proxy/src/credential_broker.rs    | 12 ++++-
 .../src/credential_broker/providers.rs        |  6 +--
 3 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs
index f090171a65..a3a3d226c6 100644
--- a/codex-rs/core/src/tools/runtimes/mod_tests.rs
+++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs
@@ -706,6 +706,50 @@ async fn snapshot_wrapper_restores_prepared_broker_credentials() -> anyhow::Resu
     Ok(())
 }
 
+#[tokio::test]
+async fn snapshot_wrapper_keeps_context_and_strips_unbrokered_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(
+        &snapshot_path,
+        "# Snapshot file\n\
+         export GH_HOST='github.example.com'\n\
+         export GITHUB_TOKEN='ghp-snapshot-real'\n",
+    )
+    .expect("write snapshot");
+    let (session_shell, shell_snapshot) =
+        shell_with_snapshot(ShellType::Bash, "/bin/bash", snapshot_path.abs());
+    let command = vec![
+        "/bin/bash".to_string(),
+        "-lc".to_string(),
+        "printf '%s\\n%s' \"$GH_HOST\" \"${GITHUB_TOKEN-unset}\"".to_string(),
+    ];
+    let env = proxy
+        .prepare_for_optional_environment(HashMap::new(), /*environment_id*/ None)?
+        .env;
+    let rewritten = maybe_wrap_shell_lc_with_snapshot(
+        &command,
+        &session_shell,
+        Some(&shell_snapshot),
+        &HashMap::new(),
+        &env,
+        &RuntimePathPrepends::default(),
+    );
+    let output = Command::new(&rewritten[0])
+        .args(&rewritten[1..])
+        .envs(&env)
+        .output()
+        .expect("run rewritten command");
+
+    assert!(output.status.success(), "command failed: {output:?}");
+    assert_eq!(
+        String::from_utf8_lossy(&output.stdout),
+        "github.example.com\nunset"
+    );
+    Ok(())
+}
+
 #[test]
 fn broker_inactive_snapshot_exports_omit_credentials() {
     let (captures, restores) = build_proxy_env_exports(&HashMap::new());
diff --git a/codex-rs/network-proxy/src/credential_broker.rs b/codex-rs/network-proxy/src/credential_broker.rs
index 1b884cc67b..270f614d9a 100644
--- a/codex-rs/network-proxy/src/credential_broker.rs
+++ b/codex-rs/network-proxy/src/credential_broker.rs
@@ -293,12 +293,20 @@ pub fn brokered_credential_dummy_env_keys(env: &HashMap<String, String>) -> Vec<
     keys
 }
 
-/// Returns supported credential keys only for an environment with an active broker.
+/// Returns credential keys plus provider context keys already present in an environment with an
+/// active broker.
 pub fn brokered_credential_env_keys(
     env: &HashMap<String, String>,
 ) -> impl Iterator<Item = &'static str> {
     let active = env_value(env, CREDENTIAL_BROKER_ACTIVE_ENV_KEY).is_some_and(|value| value == "1");
-    providers::credential_broker_env_keys().filter(move |_| active)
+    let mut keys = Vec::new();
+    if active {
+        keys.extend(providers::credential_env_keys());
+        keys.extend(
+            providers::credential_context_env_keys().filter(|key| env_value(env, key).is_some()),
+        );
+    }
+    keys.into_iter()
 }
 
 #[cfg(test)]
diff --git a/codex-rs/network-proxy/src/credential_broker/providers.rs b/codex-rs/network-proxy/src/credential_broker/providers.rs
index 68d55173e9..5d8d823a46 100644
--- a/codex-rs/network-proxy/src/credential_broker/providers.rs
+++ b/codex-rs/network-proxy/src/credential_broker/providers.rs
@@ -75,10 +75,8 @@ impl CredentialHostBinding {
     }
 }
 
-pub(super) fn credential_broker_env_keys() -> impl Iterator<Item = &'static str> {
-    credential_providers()
-        .flat_map(|provider| provider.context_env_vars.iter().copied())
-        .chain(credential_env_keys())
+pub(super) fn credential_context_env_keys() -> impl Iterator<Item = &'static str> {
+    credential_providers().flat_map(|provider| provider.context_env_vars.iter().copied())
 }
 
 pub(super) fn credential_env_keys() -> impl Iterator<Item = &'static str> {