diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index fd57ee6b3d..1bc9a08750 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -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>, shell_snapshot: ShellSnapshotTask, shell_snapshot_builder: ShellSnapshot, + shell_snapshot_cache: ShellSnapshotCache, shell_snapshot_v2_supported: bool, installed_config: Option, } @@ -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, ) -> 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; diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 46c24421be..fc1b99ffcc 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -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 { 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 { diff --git a/codex-rs/core/src/exec_tests.rs b/codex-rs/core/src/exec_tests.rs index 3454d51db6..b343126dd5 100644 --- a/codex-rs/core/src/exec_tests.rs +++ b/codex-rs/core/src/exec_tests.rs @@ -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, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 84a2815c61..f224c7ed28 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -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>>>; +pub(crate) type ShellSnapshotCache = + Arc>>>; + +#[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, pub(crate) shell_snapshot: ShellSnapshotTask, pub(crate) shell_snapshot_builder: Option>, + 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() } diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index 1bc1aa07db..ad48f88ed7 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -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, + context_env: HashMap, + context_credential_keys: HashMap>, + binding_context_credential_keys: HashMap>, } 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)> { + 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::>::new(); + let mut binding_context_credential_keys = HashMap::>::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, }), )) } diff --git a/codex-rs/core/src/shell_snapshot_sandbox.rs b/codex-rs/core/src/shell_snapshot_sandbox.rs index 2b0dbd1817..3d3b924594 100644 --- a/codex-rs/core/src/shell_snapshot_sandbox.rs +++ b/codex-rs/core/src/shell_snapshot_sandbox.rs @@ -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, environment_id: String, additional_permissions: Option, + network_denial_cancellation_token: Option, } 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 { + 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, @@ -154,29 +184,43 @@ impl ShellSnapshotSandbox { .map(PathUri::to_abs_path) .collect::>>() .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) diff --git a/codex-rs/core/src/shell_snapshot_sandbox_tests.rs b/codex-rs/core/src/shell_snapshot_sandbox_tests.rs new file mode 100644 index 0000000000..2a44248977 --- /dev/null +++ b/codex-rs/core/src/shell_snapshot_sandbox_tests.rs @@ -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(()) +} diff --git a/codex-rs/core/src/shell_snapshot_tests.rs b/codex-rs/core/src/shell_snapshot_tests.rs index 28045a5c87..6ca44a49e9 100644 --- a/codex-rs/core/src/shell_snapshot_tests.rs +++ b/codex-rs/core/src/shell_snapshot_tests.rs @@ -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(()) } diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index 75b753f55c..9a6ffec4a2 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -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::>(); + 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); diff --git a/codex-rs/network-proxy/src/credential_broker.rs b/codex-rs/network-proxy/src/credential_broker.rs index 97e1f5f1d9..04e9b8d5cf 100644 --- a/codex-rs/network-proxy/src/credential_broker.rs +++ b/codex-rs/network-proxy/src/credential_broker.rs @@ -18,6 +18,7 @@ pub(crate) struct CredentialBroker { #[derive(Default)] struct CredentialBrokerState { + config_revision: u64, enabled: bool, openai_api_host: Option, credentials: Vec, @@ -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, diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index a2ad0f95d4..aab1c44a9e 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -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 { let (mut config, brokerage_created_default_allowlist) = diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 68b85be215..3e38030ab3 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -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();