github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 05:22:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cache protected shell snapshots and harden capture cleanup (#43954)

Browse Source 
## Why

Credential-brokered commands rebuild shell snapshots on every invocation. Snapshot capture failures can also include credential-bearing startup output, and descendants holding output pipes need cleanup even after the shell exits.

## What changed

- Cache successful protected snapshots per environment, keyed by working directory, shell, login mode, and sandbox configuration. Rebuild missing or stale snapshots while keeping concurrent captures independently cancellable.
- Invalidate snapshots when shell settings or credential broker configuration change, and retry capture once if broker configuration changes during startup.
- Restore credentials and associated provider context while respecting environment policy filters and explicit overrides.
- Keep full-buffer capture subject to timeout and cancellation through output draining, clean up descendants on capture expiration or drain failure, and omit startup output from snapshot errors.

## Testing

Add regression coverage for snapshot reuse, invalidation, recovery after storage failure, concurrent cancellation, credential-safe errors, and descendant cleanup. Verify that successful background startup and output beyond the shell output cap remain supported.

GitOrigin-RevId: 7c7c8455bf6e3b9ca173dbb4277783469f0b45e6
This commit is contained in:
Winston Howes
2026-09-09 00:42:57 +00:00
committed by copyberry
parent 929389f596
commit 808b3411fd
12 changed files with 1022 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/environment_selection.rs
+35 -1
View File
@@ -34,6 +34,7 @@ use tokio_util::task::AbortOnDropHandle;
use tracing::Instrument;
use tracing::instrument::WithSubscriber;
use crate::session::turn_context::ShellSnapshotCache;
use crate::session::turn_context::ShellSnapshotTask;
use crate::session::turn_context::TurnEnvironment;
use crate::shell::Shell;
@@ -145,6 +146,7 @@ struct ResolvedEnvironment {
temporary_directories: Option<Vec<PathUri>>,
shell_snapshot: ShellSnapshotTask,
shell_snapshot_builder: ShellSnapshot,
shell_snapshot_cache: ShellSnapshotCache,
shell_snapshot_v2_supported: bool,
installed_config: Option<EnvironmentConfig>,
}
@@ -275,6 +277,7 @@ impl ThreadEnvironments {
temporary_directories: environment.temporary_directories,
shell_snapshot: inherited_snapshot,
shell_snapshot_builder: shell_snapshot.clone(),
shell_snapshot_cache: Arc::default(),
shell_snapshot_v2_supported: environment.shell_snapshot_v2_supported
&& (selected_environment.is_remote()
|| !shell_snapshot.should_rebuild_inherited()),
@@ -311,7 +314,7 @@ impl ThreadEnvironments {
cwd: PathUri,
shell: Option<Shell>,
) -> ShellSnapshotTask {
// Protected snapshots are captured by the command, using its sandbox.
// Protected snapshots require the command's sandbox and are captured lazily.
if shell_snapshot.should_rebuild_inherited() {
return futures::future::ready(None).boxed().shared();
}
@@ -368,9 +371,22 @@ impl ThreadEnvironments {
);
if !failed && !restarting_as_pending {
let shell_settings_changed = matches!(
(&environment.selection.config, &selected_environment.config),
(EnvironmentConfigState::Ready(previous), EnvironmentConfigState::Ready(current))
if previous.allow_login_shell != current.allow_login_shell
|| previous.shell_environment_policy != current.shell_environment_policy
);
let mut environment = environment.clone();
environment.selection = selected_environment;
environment.config_origin = config_origin;
if shell_settings_changed
&& !environment.environment.is_remote()
&& self.shell_snapshot.should_rebuild_inherited()
&& let Some(Ok(resolved)) = environment.resolution.clone().now_or_never()
{
self.restart_shell_snapshot(&mut environment, resolved);
}
next.push(environment);
continue;
}
@@ -518,7 +534,22 @@ impl ThreadEnvironments {
.iter()
.map(|environment| {
let mut environment = environment.clone();
let shell_settings_changed = environment.config_origin
== EnvironmentConfigOrigin::Thread
&& matches!(
&environment.selection.config,
EnvironmentConfigState::Ready(previous)
if previous.allow_login_shell != config.allow_login_shell
|| previous.shell_environment_policy != config.shell_environment_policy
);
environment.refresh_thread_config(config);
if shell_settings_changed
&& !environment.environment.is_remote()
&& self.shell_snapshot.should_rebuild_inherited()
&& let Some(Ok(resolved)) = environment.resolution.clone().now_or_never()
{
self.restart_shell_snapshot(&mut environment, resolved);
}
environment
})
.collect();
@@ -558,6 +589,7 @@ impl ThreadEnvironments {
);
selected.resolution = futures::future::ready(Ok(ResolvedEnvironment {
shell_snapshot,
shell_snapshot_cache: Arc::default(),
shell_snapshot_v2_supported: cfg!(unix)
&& !self.shell_snapshot.should_rebuild_inherited(),
..resolved
@@ -768,6 +800,7 @@ impl ThreadEnvironments {
temporary_directories: temporary_dirs,
shell_snapshot: task,
shell_snapshot_builder,
shell_snapshot_cache: Arc::default(),
shell_snapshot_v2_supported,
installed_config,
})
@@ -851,6 +884,7 @@ impl TurnEnvironmentState {
turn_environment.shell_snapshot = environment.shell_snapshot;
turn_environment.shell_snapshot_builder =
Some(Box::new(environment.shell_snapshot_builder));
turn_environment.shell_snapshot_cache = environment.shell_snapshot_cache;
turn_environment.shell_snapshot_v2_supported =
environment.shell_snapshot_v2_supported;
turn_environment.user_home_dir = environment.user_home_dir;
codex-rs/core/src/exec.rs
+78 -5
View File
@@ -117,6 +117,8 @@ pub enum ExecCapturePolicy {
/// Trusted internal helpers can buffer the full child output in memory
/// without the shell-oriented output cap or exec-expiration behavior.
FullBuffer,
/// Internal helpers that need complete output but must honor timeout and cancellation.
FullBufferWithExpiration,
}
fn select_process_exec_tool_sandbox_type(
@@ -266,7 +268,7 @@ impl ExecCapturePolicy {
fn retained_bytes_cap(self) -> Option<usize> {
match self {
Self::ShellTool => Some(EXEC_OUTPUT_MAX_BYTES),
Self::FullBuffer => None,
Self::FullBuffer | Self::FullBufferWithExpiration => None,
}
}
@@ -276,7 +278,7 @@ impl ExecCapturePolicy {
fn uses_expiration(self) -> bool {
match self {
Self::ShellTool => true,
Self::ShellTool | Self::FullBufferWithExpiration => true,
Self::FullBuffer => false,
}
}
@@ -982,12 +984,15 @@ async fn consume_output(
}
};
tokio::pin!(expiration_wait);
let (exit_status, timed_out) = tokio::select! {
let process_group_id = child.id();
let mut expiration_resolved = false;
let (mut exit_status, mut timed_out) = tokio::select! {
status_result = child.wait() => {
let exit_status = status_result?;
(exit_status, false)
}
outcome = &mut expiration_wait => {
expiration_resolved = true;
match outcome {
Some(ExecExpirationOutcome::TimedOut) => {
kill_child_process_group(&mut child)?;
@@ -1065,8 +1070,76 @@ async fn consume_output(
let mut stdout_handle = stdout_handle;
let mut stderr_handle = stderr_handle;
let stdout = await_output(&mut stdout_handle, capture_policy.io_drain_timeout()).await?;
let stderr = await_output(&mut stderr_handle, capture_policy.io_drain_timeout()).await?;
let (stdout, stderr) = if capture_policy == ExecCapturePolicy::FullBufferWithExpiration {
let mut stdout_done = false;
let mut stderr_done = false;
let drain = async {
let stdout = (&mut stdout_handle).await;
stdout_done = true;
let stdout = stdout.map_err(io::Error::other)??;
let stderr = (&mut stderr_handle).await;
stderr_done = true;
let stderr = stderr.map_err(io::Error::other)??;
Ok::<_, io::Error>((stdout, stderr))
};
// The leader can exit while descendants still hold its pipes. Keep the original
// expiration active until draining finishes, without polling a completed future again.
let drained = tokio::select! {
biased;
outcome = &mut expiration_wait, if !expiration_resolved => {
match outcome {
Some(ExecExpirationOutcome::TimedOut) => {
exit_status = synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE);
timed_out = true;
}
Some(ExecExpirationOutcome::Cancelled) => {
exit_status = synthetic_exit_status_for_code(/*code*/ 1);
}
None => unreachable!("full-buffer capture always uses expiration"),
}
None
}
result = tokio::time::timeout(capture_policy.io_drain_timeout(), drain) => {
Some(result.unwrap_or_else(|_| Err(io::Error::new(
io::ErrorKind::TimedOut,
"capture output pipes did not close before the drain deadline",
))))
}
};
match drained {
Some(Ok(output)) => output,
failure => {
let cleanup = process_group_id
.map_or(Ok(()), codex_utils_pty::process_group::kill_process_group);
stdout_handle.abort();
stderr_handle.abort();
if !stdout_done {
let _ = stdout_handle.await;
}
if !stderr_done {
let _ = stderr_handle.await;
}
cleanup?;
if let Some(Err(err)) = failure {
return Err(err.into());
}
(
StreamOutput {
text: Vec::new(),
truncated_after_lines: None,
},
StreamOutput {
text: Vec::new(),
truncated_after_lines: None,
},
)
}
}
} else {
let stdout = await_output(&mut stdout_handle, capture_policy.io_drain_timeout()).await?;
let stderr = await_output(&mut stderr_handle, capture_policy.io_drain_timeout()).await?;
(stdout, stderr)
};
let aggregated_output = aggregate_output(&stdout, &stderr, retained_bytes_cap);
Ok(RawExecToolCallOutput {
codex-rs/core/src/exec_tests.rs
+144 -3
View File
@@ -8,6 +8,7 @@ use core_test_support::PathExt;
use pretty_assertions::assert_eq;
use std::collections::HashMap;
use std::time::Duration;
use test_case::test_case;
use tokio::io::AsyncWriteExt;
use tokio::time::timeout;
@@ -331,8 +332,148 @@ async fn exec_full_buffer_capture_keeps_io_drain_timeout_when_descendant_holds_p
Ok(())
}
#[cfg(unix)]
#[derive(Clone, Copy)]
enum CaptureDrainFailure {
Cancellation,
Timeout,
DrainTimeout,
}
#[cfg(unix)]
#[test_case("stdout", CaptureDrainFailure::Cancellation; "cancel_stdout")]
#[test_case("stderr", CaptureDrainFailure::Cancellation; "cancel_stderr")]
#[test_case("stdout", CaptureDrainFailure::Timeout; "timeout_stdout")]
#[test_case("stderr", CaptureDrainFailure::Timeout; "timeout_stderr")]
#[test_case("stdout", CaptureDrainFailure::DrainTimeout; "incomplete_stdout")]
#[test_case("stderr", CaptureDrainFailure::DrainTimeout; "incomplete_stderr")]
#[tokio::test]
async fn process_exec_tool_call_preserves_full_buffer_capture_policy() -> Result<()> {
async fn full_buffer_expiration_cleans_up_after_leader_exit(
pipe: &str,
failure: CaptureDrainFailure,
) -> anyhow::Result<()> {
let dir = tempfile::tempdir()?;
let redirect = match pipe {
"stdout" => "2>/dev/null",
"stderr" => ">/dev/null",
_ => unreachable!("test specifies stdout or stderr"),
};
let mut command = tokio::process::Command::new("/bin/sh");
command
.args(["-c", &format!(
"printf complete; (while [ ! -f release ]; do sleep 0.01; done; printf survived > late_write) {redirect} &"
)])
.current_dir(dir.path())
.process_group(/*pgroup*/ 0)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
let child = command.spawn()?;
let process_group_id = child.id().expect("running child");
let cancellation = CancellationToken::new();
let expiration = match failure {
CaptureDrainFailure::Cancellation => ExecExpiration::Cancellation(cancellation.clone()),
CaptureDrainFailure::Timeout => ExecExpiration::Timeout(Duration::from_millis(250)),
CaptureDrainFailure::DrainTimeout => ExecExpiration::Timeout(Duration::from_secs(30)),
};
let capture = consume_output(
child,
expiration,
ExecCapturePolicy::FullBufferWithExpiration,
/*stdout_stream*/ None,
);
let after_leader_exit = async {
timeout(Duration::from_secs(5), async {
while unsafe {
libc::kill(process_group_id as libc::pid_t, /*sig*/ 0)
} == 0
{
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("capture should reap the direct child before cancellation");
if matches!(failure, CaptureDrainFailure::Cancellation) {
cancellation.cancel();
}
};
let (result, ()) = timeout(Duration::from_secs(6), async {
tokio::join!(capture, after_leader_exit)
})
.await?;
std::fs::write(dir.path().join("release"), "")?;
tokio::time::sleep(Duration::from_millis(100)).await;
let survived = dir.path().join("late_write").exists();
codex_utils_pty::process_group::kill_process_group(process_group_id)?;
assert!(
!survived,
"pipe-holding descendant survived capture failure"
);
match failure {
CaptureDrainFailure::Cancellation | CaptureDrainFailure::Timeout => {
let output = result?;
assert_eq!(
(
output.timed_out,
output.exit_status.success(),
output.stdout.text,
output.stderr.text
),
(
matches!(failure, CaptureDrainFailure::Timeout),
false,
Vec::new(),
Vec::new()
),
);
}
CaptureDrainFailure::DrainTimeout => {
assert!(result.is_err(), "incomplete capture must fail");
}
}
Ok(())
}
#[cfg(unix)]
#[tokio::test]
async fn full_buffer_expiration_preserves_successful_background_startup() -> anyhow::Result<()> {
let dir = tempfile::tempdir()?;
let mut command = tokio::process::Command::new("/bin/sh");
command
.args([
"-c",
"printf complete; (sleep 0.05; printf survived > late_write) >/dev/null 2>&1 &",
])
.current_dir(dir.path())
.process_group(/*pgroup*/ 0)
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::piped())
.kill_on_drop(true);
let output = consume_output(
command.spawn()?,
ExecExpiration::Timeout(Duration::from_secs(5)),
ExecCapturePolicy::FullBufferWithExpiration,
/*stdout_stream*/ None,
)
.await?;
assert!(output.exit_status.success());
assert_eq!(output.stdout.text, b"complete");
timeout(Duration::from_secs(5), async {
while !dir.path().join("late_write").exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await?;
Ok(())
}
#[test_case(ExecCapturePolicy::FullBuffer, 1; "without_expiration")]
#[test_case(ExecCapturePolicy::FullBufferWithExpiration, 30_000; "with_expiration")]
#[tokio::test]
async fn process_exec_tool_call_preserves_full_buffer_capture_policy(
capture_policy: ExecCapturePolicy,
timeout_ms: u64,
) -> Result<()> {
let byte_count = EXEC_OUTPUT_MAX_BYTES.saturating_add(128 * 1024);
#[cfg(windows)]
let command = vec![
@@ -355,8 +496,8 @@ async fn process_exec_tool_call_preserves_full_buffer_capture_policy() -> Result
ExecParams {
command,
cwd: cwd.clone(),
expiration: 1.into(),
capture_policy: ExecCapturePolicy::FullBuffer,
expiration: timeout_ms.into(),
capture_policy,
env: std::env::vars().collect(),
network: None,
network_environment_id: None,
codex-rs/core/src/session/turn_context.rs
+64 -13
View File
@@ -39,11 +39,22 @@ use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::future::Shared;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use tracing::instrument;
pub(crate) type ShellSnapshotTask = Shared<BoxFuture<'static, Option<Arc<ShellSnapshotFile>>>>;
pub(crate) type ShellSnapshotCache =
Arc<Mutex<HashMap<ShellSnapshotCacheKey, Arc<ShellSnapshotFile>>>>;
#[derive(Eq, Hash, PartialEq)]
pub(crate) struct ShellSnapshotCacheKey {
cwd: AbsolutePathBuf,
shell_path: PathBuf,
allow_login_shell: bool,
sandbox: String,
}
#[derive(Clone)]
pub(crate) struct TurnEnvironment {
@@ -59,6 +70,7 @@ pub(crate) struct TurnEnvironment {
pub(crate) executor_platform_os: Option<String>,
pub(crate) shell_snapshot: ShellSnapshotTask,
pub(crate) shell_snapshot_builder: Option<Box<ShellSnapshot>>,
pub(crate) shell_snapshot_cache: ShellSnapshotCache,
pub(crate) shell_snapshot_v2_supported: bool,
}
@@ -80,6 +92,7 @@ impl TurnEnvironment {
executor_platform_os: None,
shell_snapshot: futures::future::ready(None).boxed().shared(),
shell_snapshot_builder: None,
shell_snapshot_cache: Arc::default(),
shell_snapshot_v2_supported: false,
}
}
@@ -127,20 +140,58 @@ impl TurnEnvironment {
}
if credential_broker_enabled {
let sandbox = sandbox?;
self.shell_snapshot_builder
.as_ref()?
.as_ref()
.clone()
.build(
Arc::clone(&self.environment),
PathUri::from_abs_path(cwd),
Some(shell.clone()),
command[1] == "-lc",
self.shell_environment_policy().clone(),
Some(sandbox),
)
let use_login_shell = command[1] == "-lc";
let shell_snapshot_builder = self.shell_snapshot_builder.as_ref()?;
let key = ShellSnapshotCacheKey {
cwd: cwd.clone(),
shell_path: shell.shell_path.clone(),
allow_login_shell: use_login_shell,
sandbox: sandbox
.cache_key()
.inspect_err(|err| {
tracing::warn!("Failed to identify shell snapshot sandbox: {err:?}");
})
.ok()?,
};
if let Some(snapshot) = self
.shell_snapshot_cache
.lock()
.await
.filter(|snapshot| snapshot.is_brokered_for(self.shell_environment_policy()))
.get(&key)
.filter(|snapshot| {
snapshot.path().as_path().is_file()
&& snapshot.is_brokered_for(self.shell_environment_policy())
})
.cloned()
{
return Some(snapshot);
}
// Captures keep this command's cancellation token. Retry once if brokerage
// changes during startup, without retrying failed or cancelled captures.
for _ in 0..2 {
let snapshot = shell_snapshot_builder
.as_ref()
.clone()
.build(
Arc::clone(&self.environment),
PathUri::from_abs_path(cwd),
Some(shell.clone()),
use_login_shell,
self.shell_environment_policy().clone(),
Some(sandbox.clone()),
)
.await?;
if !snapshot.is_brokered_for(self.shell_environment_policy()) {
continue;
}
let mut snapshots = self.shell_snapshot_cache.lock().await;
if snapshots.len() >= 32 && !snapshots.contains_key(&key) {
snapshots.clear();
}
snapshots.insert(key, Arc::clone(&snapshot));
return Some(snapshot);
}
None
} else {
self.shell_snapshot.peek()?.clone()
}
codex-rs/core/src/shell_snapshot.rs
+105 -13
View File
@@ -17,6 +17,7 @@ use anyhow::bail;
use codex_exec_server::Environment;
use codex_network_proxy::CREDENTIAL_BROKER_ACTIVE_ENV_KEY;
use codex_network_proxy::NetworkProxy;
use codex_network_proxy::brokered_credential_binding_env_keys;
use codex_network_proxy::brokered_credential_dummy_env_keys;
use codex_network_proxy::brokered_credential_env_keys;
use codex_network_proxy::credential_broker_provider_context_env_keys;
@@ -77,8 +78,12 @@ pub(crate) struct ShellSnapshotFile {
struct SnapshotCredentials {
network_proxy: NetworkProxy,
broker_config_revision: u64,
shell_environment_policy: ShellEnvironmentPolicy,
credential_env: HashMap<String, String>,
context_env: HashMap<String, String>,
context_credential_keys: HashMap<String, Vec<String>>,
binding_context_credential_keys: HashMap<String, Vec<String>>,
}
struct SnapshotCredentialBroker {
@@ -316,6 +321,10 @@ impl ShellSnapshotFile {
) -> bool {
self.credentials.as_ref().is_some_and(|credentials| {
&credentials.shell_environment_policy == shell_environment_policy
&& credentials.broker_config_revision
== credentials
.network_proxy
.credential_broker_config_revision()
})
}
@@ -332,20 +341,72 @@ impl ShellSnapshotFile {
return;
};
// This snapshot belongs to one command and was captured under this exact policy.
if &credentials.shell_environment_policy != shell_environment_policy {
return;
}
let mut snapshot_env = credentials.credential_env.clone();
snapshot_env.insert(
let allowed_snapshot_env = create_env_from_vars(
credentials
.credential_env
.iter()
.chain(&credentials.context_env)
.map(|(key, value)| (key.clone(), value.clone())),
shell_environment_policy,
/*thread_id*/ None,
);
let explicit_env_overrides = &shell_environment_policy.r#set;
let mut snapshot_real_credentials = credentials.credential_env.clone();
snapshot_real_credentials.insert(
CREDENTIAL_BROKER_ACTIVE_ENV_KEY.to_string(),
"1".to_string(),
);
credentials
.network_proxy
.restore_brokered_credentials(&mut snapshot_env, &mut []);
snapshot_env.remove(CREDENTIAL_BROKER_ACTIVE_ENV_KEY);
env.extend(snapshot_env);
.restore_brokered_credentials(&mut snapshot_real_credentials, &mut []);
let mut restored_credential_keys = Vec::new();
for (key, value) in &credentials.credential_env {
if explicit_env_overrides.contains_key(key) || !allowed_snapshot_env.contains_key(key) {
continue;
}
let current = env.entry(key.clone()).or_default();
if snapshot_real_credentials.get(key) != Some(current) {
current.clone_from(value);
restored_credential_keys.push(key);
}
}
for (key, value) in &credentials.context_env {
let Some(credential_keys) = credentials.context_credential_keys.get(key) else {
continue;
};
if explicit_env_overrides.contains_key(key)
|| !allowed_snapshot_env.contains_key(key)
|| !credential_keys
.iter()
.any(|credential_key| allowed_snapshot_env.contains_key(credential_key))
{
continue;
}
let restores_binding_context = credentials
.binding_context_credential_keys
.get(key)
.is_some_and(|binding_keys| {
binding_keys
.iter()
.any(|credential_key| restored_credential_keys.contains(&credential_key))
});
if restores_binding_context || !env.contains_key(key) {
env.insert(key.clone(), value.clone());
}
}
let previous = env.insert(
CREDENTIAL_BROKER_ACTIVE_ENV_KEY.to_string(),
"1".to_string(),
);
credentials
.network_proxy
.restore_brokered_credentials(env, &mut []);
if let Some(previous) = previous {
env.insert(CREDENTIAL_BROKER_ACTIVE_ENV_KEY.to_string(), previous);
} else {
env.remove(CREDENTIAL_BROKER_ACTIVE_ENV_KEY);
}
}
}
@@ -394,6 +455,9 @@ async fn capture_snapshot(
credential_broker: Option<&SnapshotCredentialBroker>,
sandbox: Option<&ShellSnapshotSandbox>,
) -> Result<(String, Option<SnapshotCredentials>)> {
let broker_config_revision = credential_broker
.map(|broker| broker.network_proxy.credential_broker_config_revision())
.unwrap_or_default();
let shell_type = shell.shell_type;
let shell_startup = if credential_broker.is_some_and(|broker| !broker.allow_login_shell) {
SnapshotStartup::NonInteractive
@@ -509,17 +573,45 @@ async fn capture_snapshot(
)
.ok_or_else(|| anyhow!("shell snapshot contains a credential outside supported exports"))?;
let credential_env = brokered_credential_env_keys(&env)
.map(str::to_string)
.chain(allowed_brokered_keys)
.filter_map(|key| env.get(&key).map(|value| (key, value.clone())))
let credential_env = allowed_brokered_keys
.iter()
.filter_map(|key| env.get(key).map(|value| (key.clone(), value.clone())))
.collect();
let mut context_credential_keys = HashMap::<String, Vec<String>>::new();
let mut binding_context_credential_keys = HashMap::<String, Vec<String>>::new();
for key in &allowed_brokered_keys {
let mut provider_env = env.clone();
provider_env
.retain(|candidate, _| candidate == key || !allowed_brokered_keys.contains(candidate));
for context_key in brokered_credential_binding_env_keys(&provider_env) {
binding_context_credential_keys
.entry(context_key.to_string())
.or_default()
.push(key.clone());
}
for context_key in brokered_credential_env_keys(&provider_env)
.filter(|context_key| *context_key != key.as_str())
{
context_credential_keys
.entry(context_key.to_string())
.or_default()
.push(key.clone());
}
}
let context_env = context_credential_keys
.keys()
.filter_map(|key| env.get(key).map(|value| (key.clone(), value.clone())))
.collect();
Ok((
snapshot,
Some(SnapshotCredentials {
network_proxy: credential_broker.network_proxy.clone(),
broker_config_revision,
shell_environment_policy: policy.clone(),
credential_env,
context_env,
context_credential_keys,
binding_context_credential_keys,
}),
))
}
codex-rs/core/src/shell_snapshot_sandbox.rs
+58 -14
View File
@@ -10,6 +10,8 @@ use anyhow::anyhow;
use anyhow::bail;
use codex_network_proxy::NetworkProxy;
use codex_protocol::config_types::WindowsSandboxLevel;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::SandboxErr;
use codex_protocol::models::AdditionalPermissionProfile;
use codex_protocol::models::FileSystemPermissions;
use codex_protocol::models::PermissionProfile;
@@ -23,6 +25,11 @@ use codex_utils_path_uri::PathUri;
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Duration;
use tokio_util::sync::CancellationToken;
#[cfg(all(test, unix))]
#[path = "shell_snapshot_sandbox_tests.rs"]
mod tests;
#[derive(Clone)]
pub(crate) struct ShellSnapshotSandbox {
@@ -39,6 +46,7 @@ pub(crate) struct ShellSnapshotSandbox {
network: Option<NetworkProxy>,
environment_id: String,
additional_permissions: Option<AdditionalPermissionProfile>,
network_denial_cancellation_token: Option<CancellationToken>,
}
pub(crate) fn snapshot_read_permissions(
@@ -83,9 +91,31 @@ impl ShellSnapshotSandbox {
network: attempt.network_proxy(network).cloned(),
environment_id: environment_id.to_string(),
additional_permissions: additional_permissions.cloned(),
network_denial_cancellation_token: attempt.network_denial_cancellation_token.clone(),
}
}
pub(crate) fn cache_key(&self) -> Result<String> {
serde_json::to_string(&(
self.sandbox.as_metric_tag(),
self.sandbox_requested,
&self.permissions,
self.enforce_managed_network,
&self.sandbox_policy_cwd,
&self.workspace_roots,
&self.codex_linux_sandbox_exe,
self.use_legacy_landlock,
self.windows_sandbox_level,
self.windows_sandbox_private_desktop,
self.network
.as_ref()
.map(|network| (network.http_addr(), network.socks_addr())),
&self.environment_id,
&self.additional_permissions,
))
.context("failed to serialize shell snapshot sandbox policy")
}
pub(crate) async fn run(
&self,
args: Vec<String>,
@@ -154,29 +184,43 @@ impl ShellSnapshotSandbox {
.map(PathUri::to_abs_path)
.collect::<std::io::Result<Vec<_>>>()
.context("failed to resolve shell snapshot workspace roots")?;
let cancellation = self
.network_denial_cancellation_token
.as_ref()
.map_or_else(CancellationToken::new, CancellationToken::child_token);
let _cancel_on_drop = cancellation.clone().drop_guard();
let expiration = ExecExpiration::TimeoutOrCancellation {
timeout: snapshot_timeout,
cancellation,
};
let request = ExecRequest::from_sandbox_exec_request(
request,
ExecOptions {
expiration: ExecExpiration::Timeout(snapshot_timeout),
capture_policy: ExecCapturePolicy::FullBuffer,
expiration,
capture_policy: ExecCapturePolicy::FullBufferWithExpiration,
},
workspace_roots,
)
.context("failed to prepare shell snapshot execution")?;
let output = tokio::time::timeout(
snapshot_timeout,
execute_env(request, /*stdout_stream*/ None),
)
.await
.map_err(|_| anyhow!("Snapshot command timed out for {shell_name}"))?
.with_context(|| format!("Failed to execute sandboxed {shell_name}"))?;
// Caller cancellation must leave execution alive long enough to clean up its process group.
let output = tokio::spawn(execute_env(request, /*stdout_stream*/ None))
.await
.context("shell snapshot execution task failed")?
.map_err(|err| match err.details() {
// Startup output can contain credentials, even when execution fails.
CodexErrorDetails::Sandbox(SandboxErr::Denied { output, .. }) => anyhow!(
"Snapshot command was denied by sandbox with status {}",
output.exit_code
),
CodexErrorDetails::Sandbox(SandboxErr::Timeout { .. }) => {
anyhow!("Snapshot command timed out")
}
_ => anyhow::Error::new(err)
.context(format!("Failed to execute sandboxed {shell_name}")),
})?;
if output.exit_code != 0 {
bail!(
"Snapshot command exited with status {}: {}",
output.exit_code,
output.stderr.text
);
bail!("Snapshot command exited with status {}", output.exit_code);
}
Ok(output.stdout.text)
codex-rs/core/src/shell_snapshot_sandbox_tests.rs
+144
View File
@@ -0,0 +1,144 @@
use super::*;
use core_test_support::PathExt;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
use test_case::test_case;
#[derive(Clone, Copy)]
enum SnapshotFailure {
Exit,
Timeout,
NetworkDenial,
CallerDrop,
CallerDropWithNetworkToken,
}
#[test_case(SnapshotFailure::Exit; "startup_exit")]
#[test_case(SnapshotFailure::Timeout; "timeout")]
#[test_case(SnapshotFailure::NetworkDenial; "network_denial")]
#[test_case(SnapshotFailure::CallerDrop; "caller_drop")]
#[test_case(SnapshotFailure::CallerDropWithNetworkToken; "caller_drop_preserves_network_token")]
#[tokio::test]
async fn snapshot_failure_omits_credentials_and_stops_descendants(
failure: SnapshotFailure,
) -> Result<()> {
let cancellation = match failure {
SnapshotFailure::NetworkDenial | SnapshotFailure::CallerDropWithNetworkToken => {
Some(CancellationToken::new())
}
SnapshotFailure::Exit | SnapshotFailure::Timeout | SnapshotFailure::CallerDrop => None,
};
let dir = tempdir()?;
let cwd = dir.path().abs();
let cwd_uri = PathUri::from_abs_path(&cwd);
let permissions = PermissionProfile::Disabled;
let manager = SandboxManager::new();
let attempt = SandboxAttempt {
sandbox: SandboxType::None,
sandbox_requested: false,
permissions: &permissions,
exec_server_permissions: &permissions,
enforce_managed_network: false,
manager: &manager,
sandbox_cwd: &cwd_uri,
workspace_roots: std::slice::from_ref(&cwd_uri),
codex_linux_sandbox_exe: None,
use_legacy_landlock: false,
windows_sandbox_level: WindowsSandboxLevel::Disabled,
windows_sandbox_private_desktop: false,
network_denial_cancellation_token: cancellation.clone(),
network_proxy: None,
};
let sandbox = ShellSnapshotSandbox::new(
&attempt, /*network*/ None, "local", /*additional_permissions*/ None,
);
let snapshot_timeout = match failure {
SnapshotFailure::Exit
| SnapshotFailure::NetworkDenial
| SnapshotFailure::CallerDrop
| SnapshotFailure::CallerDropWithNetworkToken => Duration::from_secs(30),
SnapshotFailure::Timeout => Duration::from_millis(250),
};
let script = match failure {
SnapshotFailure::Exit => "exit 23",
SnapshotFailure::Timeout
| SnapshotFailure::NetworkDenial
| SnapshotFailure::CallerDrop
| SnapshotFailure::CallerDropWithNetworkToken => {
"(printf ready > started; sleep 1; printf survived > late_write) & wait"
}
};
let capture = sandbox.run(
vec![
"/bin/sh".to_string(),
"-c".to_string(),
format!("set -x; export GH_TOKEN=ghp_abcdefghijklmnopqrstuvwxyz0123456789; {script}"),
],
&cwd,
std::env::vars().collect(),
snapshot_timeout,
"sh",
/*snapshot_read_path*/ None,
);
let cancel_after_startup = async {
if !matches!(failure, SnapshotFailure::Exit) {
tokio::time::timeout(Duration::from_secs(5), async {
while !dir.path().join("started").exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.context("startup should begin before cancellation")?;
if matches!(failure, SnapshotFailure::NetworkDenial) {
cancellation
.as_ref()
.context("network-denial test must have a cancellation token")?
.cancel();
}
}
Ok::<(), anyhow::Error>(())
};
let result = tokio::time::timeout(Duration::from_secs(6), async {
if matches!(
failure,
SnapshotFailure::CallerDrop | SnapshotFailure::CallerDropWithNetworkToken
) {
let mut capture = Box::pin(capture);
tokio::select! {
result = &mut capture => bail!("capture completed before caller drop: {result:?}"),
result = cancel_after_startup => result?,
}
drop(capture);
assert!(
cancellation
.as_ref()
.is_none_or(|token| !token.is_cancelled())
);
Ok(None)
} else {
let (result, cancellation_result) = tokio::join!(capture, cancel_after_startup);
cancellation_result?;
Ok::<_, anyhow::Error>(Some(result))
}
})
.await
.context("failed snapshot capture should return promptly")??;
if let Some(result) = result {
let error = result.err().context("failed snapshot must not be used")?;
let expected = match failure {
SnapshotFailure::Exit => "Snapshot command exited with status 23",
SnapshotFailure::Timeout => "Snapshot command timed out",
SnapshotFailure::NetworkDenial => "Snapshot command exited with status 1",
SnapshotFailure::CallerDrop | SnapshotFailure::CallerDropWithNetworkToken => {
unreachable!("dropped captures have no result")
}
};
assert_eq!(format!("{error:?}"), expected);
}
tokio::time::sleep(Duration::from_millis(1_100)).await;
assert!(
!dir.path().join("late_write").exists(),
"startup descendant survived snapshot expiration"
);
Ok(())
}
codex-rs/core/src/shell_snapshot_tests.rs
+312 -3
View File
@@ -4,8 +4,6 @@ use crate::config::NetworkProxySpec;
#[cfg(unix)]
use codex_network_proxy::NetworkProxyConfig;
#[cfg(unix)]
use codex_network_proxy::brokered_credential_binding_env_keys;
#[cfg(unix)]
use codex_protocol::config_types::EnvironmentVariablePattern;
#[cfg(unix)]
use codex_protocol::models::PermissionProfile;
@@ -132,6 +130,20 @@ fn snapshot_file_name_parser_supports_legacy_and_suffixed_names() {
#[cfg(unix)]
#[tokio::test]
async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() -> Result<()> {
use crate::config::PermissionProfileSnapshot;
use crate::environment_selection::ThreadEnvironments;
use crate::environment_selection::TurnEnvironmentSnapshot;
use crate::tools::sandboxing::SandboxAttempt;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
use codex_protocol::config_types::WindowsSandboxLevel;
use codex_protocol::protocol::EnvironmentConfig;
use codex_protocol::protocol::EnvironmentConfigState;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_sandboxing::SandboxManager;
use codex_sandboxing::SandboxType;
use tokio_util::sync::CancellationToken;
let dir = tempdir()?;
std::fs::create_dir(dir.path().join(".codex"))?;
std::fs::write(
@@ -156,7 +168,12 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() -
shell_path: PathBuf::from("/bin/bash"),
};
let permission_profile = PermissionProfile::workspace_write();
let mut network_config = NetworkProxyConfig::default();
let mut network_config = NetworkProxyConfig {
enabled: true,
proxy_url: "http://127.0.0.1:0".to_string(),
socks_url: "socks5://127.0.0.1:0".to_string(),
..NetworkProxyConfig::default()
};
network_config.set_credential_broker_enabled(/*enabled*/ true);
let network_spec = NetworkProxySpec::from_config_and_constraints(
network_config,
@@ -212,7 +229,299 @@ async fn inactive_profiles_keep_snapshots_but_active_brokers_require_sandbox() -
assert_eq!(snapshot.is_some(), expect_snapshot);
}
let (sender, _) = watch::channel(SnapshotCredentialBrokerState::Starting);
let snapshot_builder = ShellSnapshot::new(
dir.path().abs(),
session_id,
session_telemetry,
/*state_db*/ None,
Some(sender),
/*prefer_executor_snapshots*/ false,
);
let mut config = EnvironmentConfig {
allow_login_shell: false,
workspace_roots: Vec::new(),
windows_sandbox_level: WindowsSandboxLevel::Disabled,
windows_sandbox_private_desktop: true,
use_legacy_landlock: false,
permission_profile: PermissionProfileSnapshot::legacy(permission_profile),
shell_environment_policy: ShellEnvironmentPolicy::default(),
exec_policy: None,
mcp_policy: None,
network_policy: None,
selected_capability_roots: Vec::new(),
};
let manager = Arc::new(EnvironmentManager::default_for_tests());
let environments = ThreadEnvironments::new(
Arc::clone(&manager),
shell.clone(),
config.clone(),
snapshot_builder.clone(),
TurnEnvironmentSnapshot::default(),
/*non_blocking_snapshots*/ false,
);
environments.update_selections(
&[TurnEnvironmentSelection {
environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
cwd: PathUri::from_abs_path(&dir.path().abs()),
workspace_roots: Vec::new(),
config: EnvironmentConfigState::FromThread,
}],
&config,
);
for (state, expect_snapshot) in [
(SnapshotCredentialBrokerState::Starting, false),
(SnapshotCredentialBrokerState::Inactive, true),
(
SnapshotCredentialBrokerState::Ready(started_proxy.proxy()),
false,
),
(SnapshotCredentialBrokerState::Inactive, true),
] {
environments.set_snapshot_credential_broker(state.clone());
let turn = environments.snapshot().await;
let environment = turn.primary().expect("local environment");
let snapshot = timeout(SNAPSHOT_TIMEOUT, environment.shell_snapshot.clone()).await?;
assert_eq!(snapshot.is_some(), expect_snapshot);
assert_eq!(environment.shell_snapshot_v2_supported, expect_snapshot);
environments.set_snapshot_credential_broker(state.clone());
let next = environments.snapshot().await;
assert!(Arc::ptr_eq(
&environment.shell_snapshot_cache,
&next.primary().expect("next turn").shell_snapshot_cache,
));
if matches!(state, SnapshotCredentialBrokerState::Ready(_)) {
config
.shell_environment_policy
.r#set
.insert("CORP_REGION".into(), "west".into());
environments.update_thread_config(&config);
let updated = environments.snapshot().await;
assert!(!Arc::ptr_eq(
&environment.shell_snapshot_cache,
&updated
.primary()
.expect("updated turn")
.shell_snapshot_cache,
));
let child = ThreadEnvironments::new(
Arc::clone(&manager),
shell.clone(),
config.clone(),
snapshot_builder.clone(),
updated.clone(),
/*non_blocking_snapshots*/ false,
);
let child = child.snapshot().await;
assert!(!Arc::ptr_eq(
&updated.primary().expect("parent turn").shell_snapshot_cache,
&child.primary().expect("child turn").shell_snapshot_cache,
));
}
}
assert!(!dir.path().join("startup-ran").exists());
// Identical concurrent captures must not share one command's cancellation.
fs::write(
dir.path().join(".codex/startup.sh"),
"export OPENAI_API_KEY=sk-snapshot-cache-test\nprintf started > capture-started\nwhile [ ! -f finish-startup ]; do sleep 0.01; done\n",
)
.await?;
config.shell_environment_policy.r#set.insert(
"BASH_ENV".into(),
dir.path().join(".codex/startup.sh").display().to_string(),
);
environments.set_snapshot_credential_broker(SnapshotCredentialBrokerState::Ready(
started_proxy.proxy(),
));
environments.update_thread_config(&config);
let turn = environments.snapshot().await;
let environment = turn.primary().expect("brokered environment");
let mut tool_config = crate::config::ConfigBuilder::without_managed_config_for_tests()
.codex_home(dir.path().to_path_buf())
.build()
.await?;
tool_config
.features
.enable(codex_features::Feature::ShellSnapshot)?;
tool_config.permissions.network = Some(network_spec);
let cwd = dir.path().abs();
let cwd_uri = PathUri::from_abs_path(&cwd);
let manager = SandboxManager::new();
let permissions = PermissionProfile::Disabled;
let cancellation = CancellationToken::new();
let mut attempt = SandboxAttempt {
sandbox: SandboxType::None,
sandbox_requested: false,
permissions: &permissions,
exec_server_permissions: &permissions,
enforce_managed_network: false,
manager: &manager,
sandbox_cwd: &cwd_uri,
workspace_roots: std::slice::from_ref(&cwd_uri),
codex_linux_sandbox_exe: None,
use_legacy_landlock: false,
windows_sandbox_level: WindowsSandboxLevel::Disabled,
windows_sandbox_private_desktop: false,
network_denial_cancellation_token: Some(cancellation.clone()),
network_proxy: None,
};
let first_sandbox = ShellSnapshotSandbox::new(
&attempt,
/*network*/ None,
LOCAL_ENVIRONMENT_ID,
/*additional_permissions*/ None,
);
attempt.network_denial_cancellation_token = Some(CancellationToken::new());
let second_sandbox = ShellSnapshotSandbox::new(
&attempt,
/*network*/ None,
LOCAL_ENVIRONMENT_ID,
/*additional_permissions*/ None,
);
let command = shell.derive_exec_args("true", /*use_login_shell*/ false);
assert!(
environment
.shell_snapshot(&cwd, &command, &shell, &tool_config, /*sandbox*/ None)
.await
.is_none()
);
let first = async {
let snapshot = environment
.shell_snapshot(&cwd, &command, &shell, &tool_config, Some(first_sandbox))
.await;
fs::write(dir.path().join("finish-startup"), "").await?;
Ok::<_, anyhow::Error>(snapshot)
};
let second = environment.shell_snapshot(
&cwd,
&command,
&shell,
&tool_config,
Some(second_sandbox.clone()),
);
let cancel_after_startup = async {
while !dir.path().join("capture-started").exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
started_proxy
.proxy()
.add_allowed_domain("startup.example")
.await
.expect("startup network approval should not invalidate either capture");
cancellation.cancel();
};
let (first, second, ()) = timeout(SNAPSHOT_TIMEOUT * 2, async {
tokio::join!(biased; first, second, cancel_after_startup)
})
.await?;
assert!(first?.is_none(), "the cancelled capture must fail");
let second = second.context("the unrelated capture must succeed")?;
let reused = environment
.shell_snapshot(
&cwd,
&command,
&shell,
&tool_config,
Some(second_sandbox.clone()),
)
.await
.context("successful captures remain reusable")?;
assert!(Arc::ptr_eq(&second, &reused));
fs::remove_file(reused.path()).await?;
let rebuilt = environment
.shell_snapshot(
&cwd,
&command,
&shell,
&tool_config,
Some(second_sandbox.clone()),
)
.await
.context("a deleted cached snapshot must be rebuilt")?;
assert!(!Arc::ptr_eq(&reused, &rebuilt));
assert!(rebuilt.path().as_path().is_file());
let proxy = started_proxy.proxy();
let mut old_env = HashMap::new();
rebuilt.restore_credentials(&mut old_env, environment.shell_environment_policy());
proxy.apply_to_env(&mut old_env);
let original_config = proxy.current_cfg().await?;
let mut changed_config = original_config.clone();
changed_config.set_credential_broker_openai_base_url(Some("https://gateway.example/v1"));
for (config, remains_valid) in [
(original_config.clone(), true),
(changed_config, false),
(original_config, false),
] {
proxy
.replace_config_state(codex_network_proxy::build_config_state(
config,
codex_network_proxy::NetworkProxyConstraints::default(),
)?)
.await?;
assert_eq!(
rebuilt.is_brokered_for(environment.shell_environment_policy()),
remains_valid
);
}
let refreshed = environment
.shell_snapshot(
&cwd,
&command,
&shell,
&tool_config,
Some(second_sandbox.clone()),
)
.await
.context("a changed broker configuration must invalidate cached snapshots")?;
assert!(!Arc::ptr_eq(&rebuilt, &refreshed));
let mut new_env = HashMap::new();
refreshed.restore_credentials(&mut new_env, environment.shell_environment_policy());
proxy.apply_to_env(&mut new_env);
assert_ne!(new_env["OPENAI_API_KEY"], old_env["OPENAI_API_KEY"]);
assert_ne!(new_env["OPENAI_API_KEY"], "sk-snapshot-cache-test");
// A credential change during capture should recapture once under the new configuration.
fs::remove_file(refreshed.path()).await?;
fs::remove_file(dir.path().join("finish-startup")).await?;
fs::remove_file(dir.path().join("capture-started")).await?;
fs::write(
dir.path().join(".codex/startup.sh"),
"export OPENAI_API_KEY=sk-snapshot-cache-test\nprintf x >> capture-started\nwhile [ ! -f finish-startup ]; do sleep 0.01; done\n",
)
.await?;
let capture =
environment.shell_snapshot(&cwd, &command, &shell, &tool_config, Some(second_sandbox));
let reconfigure = async {
while !dir.path().join("capture-started").exists() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
let mut config = proxy.current_cfg().await?;
config.set_credential_broker_openai_base_url(Some("https://new-gateway.example/v1"));
proxy
.replace_config_state(codex_network_proxy::build_config_state(
config,
codex_network_proxy::NetworkProxyConstraints::default(),
)?)
.await?;
fs::write(dir.path().join("finish-startup"), "").await?;
Ok::<_, anyhow::Error>(())
};
let (recaptured, updated) = timeout(SNAPSHOT_TIMEOUT * 2, async {
tokio::join!(capture, reconfigure)
})
.await?;
updated?;
let recaptured = recaptured.context("credential changes during startup must recapture")?;
assert!(recaptured.is_brokered_for(environment.shell_environment_policy()));
assert_eq!(
fs::read_to_string(dir.path().join("capture-started")).await?,
"xx"
);
Ok(())
}
codex-rs/core/tests/suite/approvals.rs
+55
View File
@@ -3141,6 +3141,61 @@ GH_TOKEN = "{REAL_GITHUB_TOKEN}"
output.contains("credential brokerage could not create a protected shell snapshot")
);
assert!(!output.contains(REAL_GITHUB_TOKEN));
if mode == "failed" {
fs::remove_file(test.home.path().join("shell_snapshots"))?;
let mut first_snapshot_paths = None;
for call_id in ["brokered-startup-retry", "brokered-startup-reuse"] {
let event = shell_event(
call_id,
"printenv GH_TOKEN",
/*timeout_ms*/ 5_000,
SandboxPermissions::UseDefault,
)?;
let retry = mount_sse_sequence(
&server,
vec![
sse(vec![
ev_response_created(call_id),
event,
ev_completed(call_id),
]),
sse(vec![
ev_response_created("resp-brokered-startup-done"),
ev_assistant_message("msg-brokered-startup-done", "done"),
ev_completed("resp-brokered-startup-done"),
]),
],
)
.await;
submit_turn_preserving_active_permission_profile(
&test,
"run the command after repairing snapshot storage",
approval_policy,
)
.await?;
wait_for_completion_without_approval(&test).await;
let result = parse_result(&retry.requests()[1].function_call_output(call_id));
assert_eq!(
result.exit_code,
Some(0),
"snapshot retry failed: {result:?}"
);
assert!(result.stdout.starts_with("ghp_"));
assert!(!result.stdout.contains(REAL_GITHUB_TOKEN));
let mut snapshot_paths = fs::read_dir(test.home.path().join("shell_snapshots"))?
.flatten()
.map(|entry| entry.path())
.filter(|path| path.extension().is_some_and(|extension| extension == "sh"))
.collect::<Vec<_>>();
snapshot_paths.sort();
assert_eq!(snapshot_paths.len(), 1);
if let Some(first_snapshot_paths) = &first_snapshot_paths {
assert_eq!(&snapshot_paths, first_snapshot_paths);
} else {
first_snapshot_paths = Some(snapshot_paths);
}
}
}
return Ok(());
}
let result = parse_result(&output);
codex-rs/network-proxy/src/credential_broker.rs
+10
View File
@@ -18,6 +18,7 @@ pub(crate) struct CredentialBroker {
#[derive(Default)]
struct CredentialBrokerState {
config_revision: u64,
enabled: bool,
openai_api_host: Option<String>,
credentials: Vec<CredentialRecord>,
@@ -82,6 +83,11 @@ impl CredentialBroker {
pub(crate) fn configure(&self, config: &NetworkProxyConfig) {
let mut state = self.write_state();
if state.enabled != config.credential_broker
|| state.openai_api_host != config.credential_broker_openai_host
{
state.config_revision += 1;
}
if state.enabled != config.credential_broker {
state.enabled = config.credential_broker;
state.credentials.clear();
@@ -97,6 +103,10 @@ impl CredentialBroker {
}
}
pub(crate) fn config_revision(&self) -> u64 {
self.read_state().config_revision
}
pub(crate) fn discover_parent_credentials(
&self,
parent_env: &HashMap<String, String>,
codex-rs/network-proxy/src/proxy.rs
+5
View File
@@ -882,6 +882,11 @@ impl NetworkProxy {
self.state.current_cfg().await
}
/// Revision of credential configuration, excluding ordinary network policy changes.
pub fn credential_broker_config_revision(&self) -> u64 {
self.state.credential_broker_config_revision()
}
/// Captures the static inputs needed to launch a matching executor-local proxy.
pub async fn remote_launch_config(&self) -> Result<crate::RemoteNetworkProxyLaunchConfig> {
let (mut config, brokerage_created_default_allowlist) =
codex-rs/network-proxy/src/runtime.rs
+12 -22
View File
@@ -495,6 +495,10 @@ impl NetworkProxyState {
.map(|(config, _)| config)
}
pub(crate) fn credential_broker_config_revision(&self) -> u64 {
self.credential_broker.config_revision()
}
pub(crate) async fn current_cfg_with_brokerage_provenance(
&self,
) -> Result<(NetworkProxyConfig, bool)> {
@@ -524,19 +528,13 @@ impl NetworkProxyState {
}
pub async fn force_reload(&self) -> Result<()> {
let previous_cfg = {
let guard = self.state.read().await;
guard.config.clone()
};
match self.reloader.reload_now().await {
Ok(mut new_state) => {
self.credential_broker.configure(&new_state.config);
// Policy changes are operationally sensitive; logging diffs makes changes traceable
// without needing to dump full config blobs (which can include unrelated settings).
log_policy_changes(&previous_cfg, &new_state.config);
{
let mut guard = self.state.write().await;
self.credential_broker.configure(&new_state.config);
// Log policy diffs without dumping potentially sensitive config values.
log_policy_changes(&guard.config, &new_state.config);
new_state.blocked = guard.blocked.clone();
*guard = new_state;
}
@@ -554,8 +552,8 @@ impl NetworkProxyState {
pub async fn replace_config_state(&self, mut new_state: ConfigState) -> Result<()> {
self.reload_if_needed().await?;
self.credential_broker.configure(&new_state.config);
let mut guard = self.state.write().await;
self.credential_broker.configure(&new_state.config);
log_policy_changes(&guard.config, &new_state.config);
new_state.blocked = guard.blocked.clone();
new_state.blocked_total = guard.blocked_total;
@@ -899,20 +897,12 @@ impl NetworkProxyState {
match self.reloader.maybe_reload().await? {
None => Ok(()),
Some(mut new_state) => {
self.credential_broker.configure(&new_state.config);
let (previous_cfg, blocked, blocked_total) = {
let guard = self.state.read().await;
(
guard.config.clone(),
guard.blocked.clone(),
guard.blocked_total,
)
};
log_policy_changes(&previous_cfg, &new_state.config);
new_state.blocked = blocked;
new_state.blocked_total = blocked_total;
{
let mut guard = self.state.write().await;
self.credential_broker.configure(&new_state.config);
log_policy_changes(&guard.config, &new_state.config);
new_state.blocked = guard.blocked.clone();
new_state.blocked_total = guard.blocked_total;
*guard = new_state;
}
let source = self.reloader.source_label();