From 07c8b8c77c8cb50c71e7e9d98919c90c23226ed0 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 29 Apr 2026 12:13:57 -0700 Subject: [PATCH] fix: handle deferred network proxy denials (#19184) ## Why This bug is exposed by Guardian/auto-review approvals. With the managed network proxy enabled, a blocked network request can be reported back through the network approval service as an approval denial after the command has already started. Before this change, the shell and unified exec runtimes registered those network approval calls, but did not have a way to observe an async proxy denial as a cancellation/failure signal for the running process. The result was confusing: Guardian/auto-review could correctly deny network access, but the command path could keep running or unregister the approval without surfacing the denial as the command failure. ## What Changed - `NetworkApprovalService` now attaches a cancellation token to active and deferred network approvals. - Proxy-denial outcomes are recorded only for active registrations, cancel the owning token, and are consumed when the approval is finalized. - The shell runtime combines the normal command timeout with the network-denial cancellation token. - Unified exec stores the deferred network approval object, terminates tracked processes when the proxy denial arrives, and returns the denial as a process failure while polling or completing the process. - Tool orchestration passes the active network approval cancellation token into the sandbox attempt and preserves deferred approval errors instead of silently unregistering them. - App-server `command/exec` now handles the combined timeout-or-cancellation expiration variant used by the runtime. ## Verification - `cargo test -p codex-core network_approval --lib` - `cargo clippy -p codex-app-server --all-targets -- -D warnings` - `cargo clippy -p codex-core --all-targets -- -D warnings` --------- Co-authored-by: Codex --- codex-rs/app-server/src/command_exec.rs | 89 ++++- codex-rs/core/src/exec.rs | 102 +++++- codex-rs/core/src/exec_tests.rs | 32 +- codex-rs/core/src/session/tests.rs | 8 +- codex-rs/core/src/tools/network_approval.rs | 146 +++++--- .../core/src/tools/network_approval_tests.rs | 81 ++++- codex-rs/core/src/tools/orchestrator.rs | 27 +- .../src/tools/runtimes/apply_patch_tests.rs | 2 + codex-rs/core/src/tools/runtimes/mod_tests.rs | 11 +- codex-rs/core/src/tools/runtimes/shell.rs | 6 +- .../tools/runtimes/shell/unix_escalation.rs | 6 +- .../core/src/tools/runtimes/unified_exec.rs | 53 ++- codex-rs/core/src/tools/sandboxing.rs | 2 + .../core/src/unified_exec/async_watcher.rs | 8 +- codex-rs/core/src/unified_exec/mod.rs | 3 +- codex-rs/core/src/unified_exec/mod_tests.rs | 2 +- codex-rs/core/src/unified_exec/process.rs | 8 + .../core/src/unified_exec/process_manager.rs | 315 +++++++++++++++--- .../src/unified_exec/process_manager_tests.rs | 88 +++++ .../core/src/unified_exec/process_tests.rs | 14 + codex-rs/core/tests/suite/unified_exec.rs | 234 ++++++++++++- .../write/templates/memories/consolidation.md | 2 +- 22 files changed, 1078 insertions(+), 161 deletions(-) diff --git a/codex-rs/app-server/src/command_exec.rs b/codex-rs/app-server/src/command_exec.rs index 18077a994..48acf28bc 100644 --- a/codex-rs/app-server/src/command_exec.rs +++ b/codex-rs/app-server/src/command_exec.rs @@ -19,8 +19,8 @@ use codex_app_server_protocol::CommandExecWriteResponse; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::ServerNotification; use codex_core::config::StartedNetworkProxy; -use codex_core::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS; use codex_core::exec::ExecExpiration; +use codex_core::exec::ExecExpirationOutcome; use codex_core::exec::IO_DRAIN_TIMEOUT_MS; use codex_core::sandboxing::ExecRequest; use codex_protocol::exec_output::bytes_to_string_smart; @@ -453,17 +453,7 @@ async fn run_command(params: RunCommandParams) { } = params; let mut control_rx = control_rx; let mut control_open = true; - let expiration = async { - match expiration { - ExecExpiration::Timeout(duration) => tokio::time::sleep(duration).await, - ExecExpiration::DefaultTimeout => { - tokio::time::sleep(Duration::from_millis(DEFAULT_EXEC_COMMAND_TIMEOUT_MS)).await; - } - ExecExpiration::Cancellation(cancel) => { - cancel.cancelled().await; - } - } - }; + let expiration = expiration.wait_with_outcome(); tokio::pin!(expiration); let SpawnedProcess { session, @@ -472,7 +462,7 @@ async fn run_command(params: RunCommandParams) { exit_rx, } = spawned; tokio::pin!(exit_rx); - let mut timed_out = false; + let mut expiration_outcome = None; let (stdio_timeout_tx, stdio_timeout_rx) = watch::channel(false); let stdout_handle = spawn_process_output(SpawnProcessOutputParams { @@ -528,12 +518,12 @@ async fn run_command(params: RunCommandParams) { } } } - _ = &mut expiration, if !timed_out => { - timed_out = true; + outcome = &mut expiration, if expiration_outcome.is_none() => { + expiration_outcome = Some(outcome); session.request_terminate(); } exit = &mut exit_rx => { - if timed_out { + if matches!(expiration_outcome, Some(ExecExpirationOutcome::TimedOut)) { break EXEC_TIMEOUT_EXIT_CODE; } else { break exit.unwrap_or(-1); @@ -877,6 +867,73 @@ mod tests { // replying, so shell startup noise is allowed here. } + #[cfg(not(target_os = "windows"))] + #[tokio::test] + async fn timeout_or_cancellation_reports_cancellation_without_timeout_exit_code() { + let (tx, mut rx) = mpsc::channel(4); + let manager = CommandExecManager::default(); + let request_id = ConnectionRequestId { + connection_id: ConnectionId(9), + request_id: codex_app_server_protocol::RequestId::Integer(101), + }; + let cancellation = CancellationToken::new(); + let cancel = cancellation.clone(); + + manager + .start(StartCommandExecParams { + outgoing: Arc::new(OutgoingMessageSender::new(tx)), + request_id: request_id.clone(), + process_id: Some("proc-101".to_string()), + exec_request: ExecRequest::new( + vec!["sh".to_string(), "-lc".to_string(), "sleep 30".to_string()], + AbsolutePathBuf::current_dir().expect("current dir"), + HashMap::new(), + /*network*/ None, + ExecExpiration::TimeoutOrCancellation { + timeout: Duration::from_secs(30), + cancellation, + }, + codex_core::exec::ExecCapturePolicy::ShellTool, + SandboxType::None, + WindowsSandboxLevel::Disabled, + /*windows_sandbox_private_desktop*/ false, + PermissionProfile::read_only(), + /*arg0*/ None, + ), + started_network_proxy: None, + tty: false, + stream_stdin: false, + stream_stdout_stderr: false, + output_bytes_cap: Some(DEFAULT_OUTPUT_BYTES_CAP), + size: None, + }) + .await + .expect("timeout-or-cancellation exec should start"); + + cancel.cancel(); + + let envelope = timeout(Duration::from_secs(1), rx.recv()) + .await + .expect("timed out waiting for outgoing message") + .expect("channel closed before outgoing message"); + let OutgoingEnvelope::ToConnection { + connection_id, + message, + .. + } = envelope + else { + panic!("expected connection-scoped outgoing message"); + }; + assert_eq!(connection_id, request_id.connection_id); + let OutgoingMessage::Response(response) = message else { + panic!("expected execution response after cancellation"); + }; + assert_eq!(response.id, request_id.request_id); + let response: CommandExecResponse = + serde_json::from_value(response.result).expect("deserialize command/exec response"); + assert_ne!(response.exit_code, EXEC_TIMEOUT_EXIT_CODE); + } + #[tokio::test] async fn windows_sandbox_process_ids_reject_write_requests() { let manager = CommandExecManager::default(); diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index c261fd335..fd5cd7bcd 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -152,6 +152,19 @@ pub enum ExecExpiration { Timeout(Duration), DefaultTimeout, Cancellation(CancellationToken), + TimeoutOrCancellation { + timeout: Duration, + cancellation: CancellationToken, + }, +} + +/// Why an `ExecExpiration` completed. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExecExpirationOutcome { + /// The configured timeout elapsed. + TimedOut, + /// The cancellation token was cancelled. + Cancelled, } impl From> for ExecExpiration { @@ -169,14 +182,30 @@ impl From for ExecExpiration { } impl ExecExpiration { - pub(crate) async fn wait(self) { + /// Waits for this expiration and reports whether it timed out or was cancelled. + pub async fn wait_with_outcome(self) -> ExecExpirationOutcome { match self { - ExecExpiration::Timeout(duration) => tokio::time::sleep(duration).await, + ExecExpiration::Timeout(duration) => { + tokio::time::sleep(duration).await; + ExecExpirationOutcome::TimedOut + } ExecExpiration::DefaultTimeout => { - tokio::time::sleep(Duration::from_millis(DEFAULT_EXEC_COMMAND_TIMEOUT_MS)).await + tokio::time::sleep(Duration::from_millis(DEFAULT_EXEC_COMMAND_TIMEOUT_MS)).await; + ExecExpirationOutcome::TimedOut } ExecExpiration::Cancellation(cancel) => { cancel.cancelled().await; + ExecExpirationOutcome::Cancelled + } + ExecExpiration::TimeoutOrCancellation { + timeout, + cancellation, + } => { + tokio::select! { + biased; + _ = cancellation.cancelled() => ExecExpirationOutcome::Cancelled, + _ = tokio::time::sleep(timeout) => ExecExpirationOutcome::TimedOut, + } } } } @@ -187,8 +216,50 @@ impl ExecExpiration { ExecExpiration::Timeout(duration) => Some(duration.as_millis() as u64), ExecExpiration::DefaultTimeout => Some(DEFAULT_EXEC_COMMAND_TIMEOUT_MS), ExecExpiration::Cancellation(_) => None, + ExecExpiration::TimeoutOrCancellation { timeout, .. } => { + Some(timeout.as_millis() as u64) + } } } + + pub(crate) fn with_cancellation(self, cancellation: CancellationToken) -> Self { + match self { + ExecExpiration::Timeout(timeout) => ExecExpiration::TimeoutOrCancellation { + timeout, + cancellation, + }, + ExecExpiration::DefaultTimeout => ExecExpiration::TimeoutOrCancellation { + timeout: Duration::from_millis(DEFAULT_EXEC_COMMAND_TIMEOUT_MS), + cancellation, + }, + ExecExpiration::Cancellation(existing) => { + ExecExpiration::Cancellation(cancel_when_either(existing, cancellation)) + } + ExecExpiration::TimeoutOrCancellation { + timeout, + cancellation: existing, + } => ExecExpiration::TimeoutOrCancellation { + timeout, + cancellation: cancel_when_either(existing, cancellation), + }, + } + } +} + +pub(crate) fn cancel_when_either( + first: CancellationToken, + second: CancellationToken, +) -> CancellationToken { + let combined = CancellationToken::new(); + let cancel = combined.clone(); + tokio::spawn(async move { + tokio::select! { + _ = first.cancelled() => {} + _ = second.cancelled() => {} + } + cancel.cancel(); + }); + combined } impl ExecCapturePolicy { @@ -1266,9 +1337,9 @@ async fn consume_output( let expiration_wait = async { if capture_policy.uses_expiration() { - expiration.wait().await; + Some(expiration.wait_with_outcome().await) } else { - std::future::pending::<()>().await; + std::future::pending::>().await } }; tokio::pin!(expiration_wait); @@ -1277,10 +1348,16 @@ async fn consume_output( let exit_status = status_result?; (exit_status, false) } - _ = &mut expiration_wait => { + outcome = &mut expiration_wait => { kill_child_process_group(&mut child)?; child.start_kill()?; - (synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE), true) + let timed_out = matches!(outcome, Some(ExecExpirationOutcome::TimedOut)); + let exit_status = if timed_out { + synthetic_exit_status(EXIT_CODE_SIGNAL_BASE + TIMEOUT_CODE) + } else { + synthetic_exit_status_for_code(/*code*/ 1) + }; + (exit_status, timed_out) } _ = tokio::signal::ctrl_c() => { kill_child_process_group(&mut child)?; @@ -1390,6 +1467,12 @@ fn synthetic_exit_status(code: i32) -> ExitStatus { std::process::ExitStatus::from_raw(code) } +#[cfg(unix)] +fn synthetic_exit_status_for_code(code: i32) -> ExitStatus { + use std::os::unix::process::ExitStatusExt; + std::process::ExitStatus::from_raw(code << 8) +} + #[cfg(windows)] fn synthetic_exit_status(code: i32) -> ExitStatus { use std::os::windows::process::ExitStatusExt; @@ -1398,6 +1481,11 @@ fn synthetic_exit_status(code: i32) -> ExitStatus { std::process::ExitStatus::from_raw(code as u32) } +#[cfg(windows)] +fn synthetic_exit_status_for_code(code: i32) -> ExitStatus { + synthetic_exit_status(code) +} + #[cfg(test)] #[path = "exec_tests.rs"] mod tests; diff --git a/codex-rs/core/src/exec_tests.rs b/codex-rs/core/src/exec_tests.rs index a9304e75e..9d335a81c 100644 --- a/codex-rs/core/src/exec_tests.rs +++ b/codex-rs/core/src/exec_tests.rs @@ -8,6 +8,7 @@ use pretty_assertions::assert_eq; use std::collections::HashMap; use std::time::Duration; use tokio::io::AsyncWriteExt; +use tokio::time::timeout; fn make_exec_output( exit_code: i32, @@ -1033,22 +1034,23 @@ async fn process_exec_tool_call_respects_cancellation_token() -> Result<()> { tokio::time::sleep(Duration::from_millis(1_000)).await; cancel_tx.cancel(); }); - let permission_profile = PermissionProfile::Disabled; - let result = process_exec_tool_call( - params, - &permission_profile, - &cwd, - &None, - /*use_legacy_landlock*/ false, - /*stdout_stream*/ None, + let result = timeout( + Duration::from_secs(5), + process_exec_tool_call( + params, + &PermissionProfile::Disabled, + &cwd, + &None, + /*use_legacy_landlock*/ false, + /*stdout_stream*/ None, + ), ) - .await; - let output = match result { - Err(CodexErr::Sandbox(SandboxErr::Timeout { output })) => output, - other => panic!("expected timeout error, got {other:?}"), - }; - assert!(output.timed_out); - assert_eq!(output.exit_code, EXEC_TIMEOUT_EXIT_CODE); + .await + .expect("cancellation should stop the process promptly"); + let output = result.expect("cancellation should return a non-timeout exec result"); + assert!(!output.timed_out); + assert_ne!(output.exit_code, 0); + assert_ne!(output.exit_code, EXEC_TIMEOUT_EXIT_CODE); Ok(()) } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 1c2e87605..e7484e43b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1566,14 +1566,12 @@ async fn session_configured_reports_permission_profile_for_external_sandbox() -> }; let expected_sandbox_policy = sandbox_policy.clone(); let mut builder = test_codex().with_config(move |config| { + config.permissions.permission_profile = codex_config::Constrained::allow_any( + PermissionProfile::from_legacy_sandbox_policy(&sandbox_policy), + ); config .set_legacy_sandbox_policy(sandbox_policy) .expect("set sandbox policy"); - config.permissions.permission_profile = - codex_config::Constrained::allow_any(PermissionProfile::from_runtime_permissions( - &FileSystemSandboxPolicy::external_sandbox(), - NetworkSandboxPolicy::Restricted, - )); }); let test = builder.build(&server).await?; diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 1264e809b..14af2c9c5 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -33,7 +33,9 @@ use std::collections::HashSet; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::Notify; +use tokio::sync::OnceCell; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; use tracing::warn; use uuid::Uuid; @@ -54,18 +56,38 @@ pub(crate) struct NetworkApprovalSpec { #[derive(Clone, Debug)] pub(crate) struct DeferredNetworkApproval { registration_id: String, + cancellation_token: CancellationToken, + finish_outcome: Arc>>, } impl DeferredNetworkApproval { pub(crate) fn registration_id(&self) -> &str { &self.registration_id } + + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancellation_token.clone() + } + + pub(crate) fn is_cancelled(&self) -> bool { + self.cancellation_token.is_cancelled() + } + + async fn finish(&self, service: &NetworkApprovalService) -> Result<(), ToolError> { + let outcome = self + .finish_outcome + .get_or_init(|| async { service.finish_call_outcome(&self.registration_id).await }) + .await + .clone(); + network_approval_outcome_to_result(outcome) + } } #[derive(Debug)] pub(crate) struct ActiveNetworkApproval { registration_id: Option, mode: NetworkApprovalMode, + cancellation_token: CancellationToken, } impl ActiveNetworkApproval { @@ -73,10 +95,23 @@ impl ActiveNetworkApproval { self.mode } + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancellation_token.clone() + } + pub(crate) fn into_deferred(self) -> Option { - match (self.mode, self.registration_id) { + let ActiveNetworkApproval { + registration_id, + mode, + cancellation_token, + } = self; + match (mode, registration_id) { (NetworkApprovalMode::Deferred, Some(registration_id)) => { - Some(DeferredNetworkApproval { registration_id }) + Some(DeferredNetworkApproval { + registration_id, + cancellation_token, + finish_outcome: Arc::new(OnceCell::new()), + }) } _ => None, } @@ -122,6 +157,18 @@ enum NetworkApprovalOutcome { DeniedByPolicy(String), } +fn network_approval_outcome_to_result( + outcome: Option, +) -> Result<(), ToolError> { + match outcome { + Some(NetworkApprovalOutcome::DeniedByUser) => { + Err(ToolError::Rejected("rejected by user".to_string())) + } + Some(NetworkApprovalOutcome::DeniedByPolicy(message)) => Err(ToolError::Rejected(message)), + None => Ok(()), + } +} + /// Whether an allowlist miss may be reviewed instead of hard-denied. fn allows_network_approval_flow(policy: AskForApproval) -> bool { !matches!(policy, AskForApproval::Never) @@ -177,11 +224,17 @@ struct ActiveNetworkApprovalCall { turn_id: String, trigger: GuardianNetworkAccessTrigger, command: String, + cancellation_token: CancellationToken, +} + +#[derive(Default)] +struct NetworkApprovalCallState { + active_calls: IndexMap>, + call_outcomes: HashMap, } pub(crate) struct NetworkApprovalService { - active_calls: Mutex>>, - call_outcomes: Mutex>, + calls: Mutex, pending_host_approvals: Mutex>>, session_approved_hosts: Mutex>, session_denied_hosts: Mutex>, @@ -190,8 +243,7 @@ pub(crate) struct NetworkApprovalService { impl Default for NetworkApprovalService { fn default() -> Self { Self { - active_calls: Mutex::new(IndexMap::new()), - call_outcomes: Mutex::new(HashMap::new()), + calls: Mutex::new(NetworkApprovalCallState::default()), pending_host_approvals: Mutex::new(HashMap::new()), session_approved_hosts: Mutex::new(HashSet::new()), session_denied_hosts: Mutex::new(HashSet::new()), @@ -215,29 +267,33 @@ impl NetworkApprovalService { turn_id: String, trigger: GuardianNetworkAccessTrigger, command: String, + cancellation_token: CancellationToken, ) { - let mut active_calls = self.active_calls.lock().await; + let mut calls = self.calls.lock().await; let key = registration_id.clone(); - active_calls.insert( + calls.active_calls.insert( key, Arc::new(ActiveNetworkApprovalCall { registration_id, turn_id, trigger, command, + cancellation_token, }), ); } pub(crate) async fn unregister_call(&self, registration_id: &str) { - self.active_calls.lock().await.shift_remove(registration_id); - self.call_outcomes.lock().await.remove(registration_id); + self.remove_call(registration_id).await; } async fn resolve_single_active_call(&self) -> Option> { - let active_calls = self.active_calls.lock().await; - if active_calls.len() == 1 { - return active_calls.values().next().cloned(); + let calls = self.calls.lock().await; + // Blocked proxy requests are not attributed to a specific tool call. Only pick an owner + // when there is exactly one candidate; with concurrent calls, canceling one would be a guess. + // TODO: Carry blocked-request attribution so concurrent active calls can be handled safely. + if calls.active_calls.len() == 1 { + return calls.active_calls.values().next().cloned(); } None @@ -265,20 +321,43 @@ impl NetworkApprovalService { .await; } + #[cfg(test)] async fn take_call_outcome(&self, registration_id: &str) -> Option { - let mut call_outcomes = self.call_outcomes.lock().await; - call_outcomes.remove(registration_id) + let mut calls = self.calls.lock().await; + calls.call_outcomes.remove(registration_id) } async fn record_call_outcome(&self, registration_id: &str, outcome: NetworkApprovalOutcome) { - let mut call_outcomes = self.call_outcomes.lock().await; + let mut calls = self.calls.lock().await; + let Some(call) = calls.active_calls.get(registration_id).cloned() else { + return; + }; if matches!( - call_outcomes.get(registration_id), + calls.call_outcomes.get(registration_id), Some(NetworkApprovalOutcome::DeniedByUser) ) { return; } - call_outcomes.insert(registration_id.to_string(), outcome); + calls + .call_outcomes + .insert(registration_id.to_string(), outcome); + + drop(calls); + call.cancellation_token.cancel(); + } + + async fn remove_call(&self, registration_id: &str) -> Option { + let mut calls = self.calls.lock().await; + calls.active_calls.shift_remove(registration_id); + calls.call_outcomes.remove(registration_id) + } + + async fn finish_call_outcome(&self, registration_id: &str) -> Option { + self.remove_call(registration_id).await + } + + async fn finish_call(&self, registration_id: &str) -> Result<(), ToolError> { + network_approval_outcome_to_result(self.finish_call_outcome(registration_id).await) } pub(crate) async fn record_blocked_request(&self, blocked: BlockedRequest) { @@ -637,6 +716,7 @@ pub(crate) async fn begin_network_approval( } let registration_id = Uuid::new_v4().to_string(); + let cancellation_token = CancellationToken::new(); session .services .network_approval @@ -645,12 +725,14 @@ pub(crate) async fn begin_network_approval( turn_id.to_string(), trigger, command, + cancellation_token.clone(), ) .await; Some(ActiveNetworkApproval { registration_id: Some(registration_id), mode, + cancellation_token, }) } @@ -662,39 +744,21 @@ pub(crate) async fn finish_immediate_network_approval( return Ok(()); }; - let approval_outcome = session - .services - .network_approval - .take_call_outcome(registration_id) - .await; - session .services .network_approval - .unregister_call(registration_id) - .await; - - match approval_outcome { - Some(NetworkApprovalOutcome::DeniedByUser) => { - Err(ToolError::Rejected("rejected by user".to_string())) - } - Some(NetworkApprovalOutcome::DeniedByPolicy(message)) => Err(ToolError::Rejected(message)), - None => Ok(()), - } + .finish_call(registration_id) + .await } pub(crate) async fn finish_deferred_network_approval( session: &Session, deferred: Option, -) { +) -> Result<(), ToolError> { let Some(deferred) = deferred else { - return; + return Ok(()); }; - session - .services - .network_approval - .unregister_call(deferred.registration_id()) - .await; + deferred.finish(&session.services.network_approval).await } #[cfg(test)] diff --git a/codex-rs/core/src/tools/network_approval_tests.rs b/codex-rs/core/src/tools/network_approval_tests.rs index 683c8c753..37a28b125 100644 --- a/codex-rs/core/src/tools/network_approval_tests.rs +++ b/codex-rs/core/src/tools/network_approval_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::protocol::SandboxPolicy; use core_test_support::PathBufExt; use core_test_support::test_path_buf; use pretty_assertions::assert_eq; +use tokio_util::sync::CancellationToken; #[tokio::test] async fn pending_approvals_are_deduped_per_host_protocol_and_port() { @@ -220,7 +221,8 @@ fn denied_blocked_request(host: &str) -> BlockedRequest { async fn register_call_with_default_shell_trigger( service: &NetworkApprovalService, registration_id: &str, -) { +) -> CancellationToken { + let cancellation_token = CancellationToken::new(); service .register_call( registration_id.to_string(), @@ -236,8 +238,10 @@ async fn register_call_with_default_shell_trigger( tty: None, }, "curl https://example.com".to_string(), + cancellation_token.clone(), ) .await; + cancellation_token } #[tokio::test] @@ -260,6 +264,7 @@ async fn active_call_preserves_triggering_command_context() { "turn-1".to_string(), expected.clone(), "curl https://example.com".to_string(), + CancellationToken::new(), ) .await; @@ -275,12 +280,14 @@ async fn active_call_preserves_triggering_command_context() { #[tokio::test] async fn record_blocked_request_sets_policy_outcome_for_owner_call() { let service = NetworkApprovalService::default(); - register_call_with_default_shell_trigger(&service, "registration-1").await; + let cancellation_token = + register_call_with_default_shell_trigger(&service, "registration-1").await; service .record_blocked_request(denied_blocked_request("example.com")) .await; + assert!(cancellation_token.is_cancelled()); assert_eq!( service.take_call_outcome("registration-1").await, Some(NetworkApprovalOutcome::DeniedByPolicy( @@ -307,6 +314,76 @@ async fn blocked_request_policy_does_not_override_user_denial_outcome() { ); } +#[tokio::test] +async fn finish_call_returns_denial_and_unregisters_active_call() { + let service = NetworkApprovalService::default(); + register_call_with_default_shell_trigger(&service, "registration-1").await; + + service + .record_call_outcome( + "registration-1", + NetworkApprovalOutcome::DeniedByPolicy("network denied".to_string()), + ) + .await; + + let err = service + .finish_call("registration-1") + .await + .expect_err("denial should be returned"); + + assert!(matches!(err, ToolError::Rejected(message) if message == "network denied")); + assert!(service.resolve_single_active_call().await.is_none()); + assert_eq!(service.take_call_outcome("registration-1").await, None); +} + +#[tokio::test] +async fn deferred_finish_reuses_denial_result_after_first_consumer() { + let service = NetworkApprovalService::default(); + let cancellation_token = + register_call_with_default_shell_trigger(&service, "registration-1").await; + let deferred = DeferredNetworkApproval { + registration_id: "registration-1".to_string(), + cancellation_token, + finish_outcome: Arc::new(OnceCell::new()), + }; + service + .record_call_outcome( + "registration-1", + NetworkApprovalOutcome::DeniedByPolicy("network denied".to_string()), + ) + .await; + + let first = deferred + .finish(&service) + .await + .expect_err("first consumer should see denial"); + let second = deferred + .finish(&service) + .await + .expect_err("second consumer should reuse denial"); + + assert!(matches!(first, ToolError::Rejected(message) if message == "network denied")); + assert!(matches!(second, ToolError::Rejected(message) if message == "network denied")); +} + +#[tokio::test] +async fn record_call_outcome_ignores_inactive_call() { + let service = NetworkApprovalService::default(); + let cancellation_token = + register_call_with_default_shell_trigger(&service, "registration-1").await; + service.unregister_call("registration-1").await; + + service + .record_call_outcome( + "registration-1", + NetworkApprovalOutcome::DeniedByPolicy("network denied".to_string()), + ) + .await; + + assert!(!cancellation_token.is_cancelled()); + assert_eq!(service.take_call_outcome("registration-1").await, None); +} + #[tokio::test] async fn record_blocked_request_ignores_ambiguous_unattributed_blocked_requests() { let service = NetworkApprovalService::default(); diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index 0b59215f2..dcb42c36c 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -12,6 +12,7 @@ use crate::guardian::new_guardian_review_id; use crate::guardian::routes_approval_to_guardian; use crate::hook_runtime::run_permission_request_hooks; use crate::network_policy_decision::network_approval_context_from_payload; +use crate::tools::network_approval::ActiveNetworkApproval; use crate::tools::network_approval::DeferredNetworkApproval; use crate::tools::network_approval::NetworkApprovalMode; use crate::tools::network_approval::begin_network_approval; @@ -76,7 +77,23 @@ impl ToolOrchestrator { call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), }; - let run_result = tool.run(req, attempt, &attempt_tool_ctx).await; + let attempt_with_network_approval = SandboxAttempt { + sandbox: attempt.sandbox, + permissions: attempt.permissions, + enforce_managed_network: attempt.enforce_managed_network, + manager: attempt.manager, + sandbox_cwd: attempt.sandbox_cwd, + codex_linux_sandbox_exe: attempt.codex_linux_sandbox_exe, + use_legacy_landlock: attempt.use_legacy_landlock, + windows_sandbox_level: attempt.windows_sandbox_level, + windows_sandbox_private_desktop: attempt.windows_sandbox_private_desktop, + network_denial_cancellation_token: network_approval + .as_ref() + .map(ActiveNetworkApproval::cancellation_token), + }; + let run_result = tool + .run(req, &attempt_with_network_approval, &attempt_tool_ctx) + .await; let Some(network_approval) = network_approval else { return (run_result, None); @@ -94,7 +111,11 @@ impl ToolOrchestrator { NetworkApprovalMode::Deferred => { let deferred = network_approval.into_deferred(); if run_result.is_err() { - finish_deferred_network_approval(&tool_ctx.session, deferred).await; + let finalize_result = + finish_deferred_network_approval(&tool_ctx.session, deferred).await; + if let Err(err) = finalize_result { + return (Err(err), None); + } return (run_result, None); } (run_result, deferred) @@ -219,6 +240,7 @@ impl ToolOrchestrator { .config .permissions .windows_sandbox_private_desktop, + network_denial_cancellation_token: None, }; let (first_result, first_deferred_network_approval) = Self::run_attempt( @@ -336,6 +358,7 @@ impl ToolOrchestrator { .config .permissions .windows_sandbox_private_desktop, + network_denial_cancellation_token: None, }; // Second attempt. diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index c4f9b44ae..173fa3e2a 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -153,6 +153,7 @@ fn file_system_sandbox_context_uses_active_attempt() { use_legacy_landlock: true, windows_sandbox_level: WindowsSandboxLevel::RestrictedToken, windows_sandbox_private_desktop: true, + network_denial_cancellation_token: None, }; let sandbox = ApplyPatchRuntime::file_system_sandbox_context_for_attempt(&req, &attempt) @@ -204,6 +205,7 @@ fn no_sandbox_attempt_has_no_file_system_context() { use_legacy_landlock: false, windows_sandbox_level: WindowsSandboxLevel::Disabled, windows_sandbox_private_desktop: false, + network_denial_cancellation_token: None, }; assert_eq!( diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index 6aa64d1e3..fd10f2224 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -115,6 +115,7 @@ async fn explicit_escalation_prepares_exec_without_managed_network() -> anyhow:: use_legacy_landlock: false, windows_sandbox_level: WindowsSandboxLevel::Disabled, windows_sandbox_private_desktop: false, + network_denial_cancellation_token: None, }; let exec_request = attempt @@ -675,10 +676,12 @@ fn maybe_wrap_shell_lc_with_snapshot_keeps_user_proxy_env_when_proxy_inactive() &HashMap::new(), &HashMap::new(), ); - let output = Command::new(&rewritten[0]) - .args(&rewritten[1..]) - .output() - .expect("run rewritten command"); + let mut command = Command::new(&rewritten[0]); + command.args(&rewritten[1..]); + for key in PROXY_ENV_KEYS { + command.env_remove(key); + } + let output = command.output().expect("run rewritten command"); assert!(output.status.success(), "command failed: {output:?}"); assert_eq!( diff --git a/codex-rs/core/src/tools/runtimes/shell.rs b/codex-rs/core/src/tools/runtimes/shell.rs index b5e9c9805..7f17285db 100644 --- a/codex-rs/core/src/tools/runtimes/shell.rs +++ b/codex-rs/core/src/tools/runtimes/shell.rs @@ -275,8 +275,12 @@ impl ToolRuntime for ShellRuntime { let command = build_sandbox_command(&command, &req.cwd, &env, req.additional_permissions.clone())?; + let mut expiration: crate::exec::ExecExpiration = req.timeout_ms.into(); + if let Some(cancellation) = attempt.network_denial_cancellation_token.clone() { + expiration = expiration.with_cancellation(cancellation); + } let options = ExecOptions { - expiration: req.timeout_ms.into(), + expiration, capture_policy: ExecCapturePolicy::ShellTool, }; let env = attempt diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index e61c78359..41aa552d8 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -1,6 +1,7 @@ use super::ShellRequest; use crate::exec::ExecCapturePolicy; use crate::exec::ExecExpiration; +use crate::exec::cancel_when_either; use crate::exec::is_likely_sandbox_denied; use crate::guardian::GuardianApprovalRequest; use crate::guardian::guardian_rejection_message; @@ -191,7 +192,10 @@ pub(super) async fn try_run_zsh_fork( // to minimize the time between creating the Stopwatch and starting the // escalation server. let stopwatch = Stopwatch::new(effective_timeout); - let cancel_token = stopwatch.cancellation_token(); + let mut cancel_token = stopwatch.cancellation_token(); + if let Some(cancellation) = attempt.network_denial_cancellation_token.clone() { + cancel_token = cancel_when_either(cancel_token, cancellation); + } let approval_sandbox_permissions = approval_sandbox_permissions( req.sandbox_permissions, req.additional_permissions_preapproved, diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index f8ce24831..dbdd6efb5 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -48,6 +48,7 @@ use codex_tools::UnifiedExecShellMode; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::collections::HashMap; +use tokio_util::sync::CancellationToken; /// Request payload used by the unified-exec runtime after approvals and /// sandbox preferences have been resolved for the current turn. @@ -88,6 +89,19 @@ pub struct UnifiedExecRuntime<'a> { shell_mode: UnifiedExecShellMode, } +fn unified_exec_options( + network_denial_cancellation_token: Option, +) -> ExecOptions { + let mut expiration = ExecExpiration::DefaultTimeout; + if let Some(cancellation) = network_denial_cancellation_token { + expiration = expiration.with_cancellation(cancellation); + } + ExecOptions { + expiration, + capture_policy: ExecCapturePolicy::ShellTool, + } +} + impl<'a> UnifiedExecRuntime<'a> { /// Creates a runtime bound to the shared unified-exec process manager. pub fn new(manager: &'a UnifiedExecProcessManager, shell_mode: UnifiedExecShellMode) -> Self { @@ -264,10 +278,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt let command = build_sandbox_command(&command, &req.cwd, &env, req.additional_permissions.clone()) .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; - let options = ExecOptions { - expiration: ExecExpiration::DefaultTimeout, - capture_policy: ExecCapturePolicy::ShellTool, - }; + let options = unified_exec_options(attempt.network_denial_cancellation_token.clone()); let mut exec_env = attempt .env_for(command, options, managed_network) .map_err(|err| ToolError::Codex(err.into()))?; @@ -322,10 +333,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt let command = build_sandbox_command(&command, &req.cwd, &env, req.additional_permissions.clone()) .map_err(|_| ToolError::Rejected("missing command line for PTY".to_string()))?; - let options = ExecOptions { - expiration: ExecExpiration::DefaultTimeout, - capture_policy: ExecCapturePolicy::ShellTool, - }; + let options = unified_exec_options(attempt.network_denial_cancellation_token.clone()); let mut exec_env = attempt .env_for(command, options, managed_network) .map_err(|err| ToolError::Codex(err.into()))?; @@ -355,3 +363,32 @@ impl<'a> ToolRuntime for UnifiedExecRunt }) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS; + use std::time::Duration; + + #[test] + fn unified_exec_options_combines_default_timeout_with_network_denial_cancellation() { + let cancellation = CancellationToken::new(); + let options = unified_exec_options(Some(cancellation.clone())); + + assert_eq!(options.capture_policy, ExecCapturePolicy::ShellTool); + match options.expiration { + ExecExpiration::TimeoutOrCancellation { + timeout, + cancellation: actual, + } => { + assert_eq!( + timeout, + Duration::from_millis(DEFAULT_EXEC_COMMAND_TIMEOUT_MS) + ); + cancellation.cancel(); + assert!(actual.is_cancelled()); + } + other => panic!("expected timeout-or-cancellation expiration, got {other:?}"), + } + } +} diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index f6b960ee1..122cd00fa 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -35,6 +35,7 @@ use std::collections::HashMap; use std::fmt::Debug; use std::hash::Hash; use std::sync::Arc; +use tokio_util::sync::CancellationToken; #[derive(Clone, Default, Debug)] pub(crate) struct ApprovalStore { @@ -375,6 +376,7 @@ pub(crate) struct SandboxAttempt<'a> { pub use_legacy_landlock: bool, pub windows_sandbox_level: codex_protocol::config_types::WindowsSandboxLevel, pub windows_sandbox_private_desktop: bool, + pub network_denial_cancellation_token: Option, } impl<'a> SandboxAttempt<'a> { diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 6b112f9ec..b4c7c9c8b 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -132,6 +132,7 @@ pub(crate) fn spawn_exit_watcher( cwd, Some(process_id.to_string()), transcript, + String::new(), message, duration, ) @@ -238,10 +239,15 @@ pub(crate) async fn emit_failed_exec_end_for_unified_exec( cwd: AbsolutePathBuf, process_id: Option, transcript: Arc>, + fallback_output: String, message: String, duration: Duration, ) { - let stdout = resolve_aggregated_output(&transcript, String::new()).await; + let stdout = if fallback_output.is_empty() { + resolve_aggregated_output(&transcript, fallback_output).await + } else { + fallback_output + }; let aggregated_output = if stdout.is_empty() { message.clone() } else { diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index acf361f94..97b37e8d8 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -37,6 +37,7 @@ use tokio::sync::Mutex; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; use crate::session::turn_context::TurnContext; +use crate::tools::network_approval::DeferredNetworkApproval; mod async_watcher; mod errors; @@ -150,7 +151,7 @@ struct ProcessEntry { process_id: i32, hook_command: String, tty: bool, - network_approval_id: Option, + network_approval: Option, session: Weak, last_used: tokio::time::Instant, } diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 0730335c8..fe87c6261 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -111,7 +111,7 @@ async fn exec_command_with_tty( process_id, hook_command: cmd.to_string(), tty, - network_approval_id: None, + network_approval: None, session: Arc::downgrade(session), last_used: started_at, }; diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 62a60b0a8..9671429d0 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -212,6 +212,14 @@ impl UnifiedExecProcess { } } + pub(super) fn fail_and_terminate(&self, message: String) { + let state = self.state_rx.borrow().clone(); + if state.failure_message.is_none() { + let _ = self.state_tx.send_replace(state.failed(message)); + } + self.terminate(); + } + async fn snapshot_output(&self) -> Vec> { let guard = self.output_buffer.lock().await; guard.snapshot_chunks() diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 76d8021d3..c67abc48d 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -69,6 +69,9 @@ const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [ ("GH_PAGER", "cat"), ("CODEX_CI", "1"), ]; +const NETWORK_ACCESS_DENIED_MESSAGE: &str = + "Network access was denied by the Codex sandbox network proxy."; +const LATE_NETWORK_DENIAL_GRACE_PERIOD: Duration = Duration::from_millis(100); /// Test-only override for deterministic unified exec process IDs. /// @@ -169,6 +172,8 @@ struct PreparedProcessHandles { output_closed_notify: Arc, cancellation_token: CancellationToken, pause_state: Option>, + session: Option>, + network_approval: Option, hook_command: String, process_id: i32, tty: bool, @@ -178,6 +183,151 @@ fn exec_server_process_id(process_id: i32) -> String { process_id.to_string() } +async fn unregister_network_approval_for_entry(entry: &ProcessEntry) { + if let Some(network_approval) = entry.network_approval.as_ref() + && let Some(session) = entry.session.upgrade() + { + session + .services + .network_approval + .unregister_call(network_approval.registration_id()) + .await; + } +} + +async fn finish_network_approval_after_process_exit_for_entry( + entry: &ProcessEntry, +) -> Result<(), String> { + let session = entry.session.upgrade(); + finish_deferred_network_approval_after_process_exit_for_session( + session.as_ref(), + entry.network_approval.clone(), + ) + .await +} + +async fn finish_deferred_network_approval_for_session( + session: Option<&Arc>, + deferred: Option, +) -> Result<(), String> { + let Some(session) = session else { + return Ok(()); + }; + finish_deferred_network_approval(session.as_ref(), deferred) + .await + .map_err(network_approval_error_message) +} + +fn network_approval_error_message(err: ToolError) -> String { + match err { + ToolError::Rejected(message) => message, + ToolError::Codex(err) => err.to_string(), + } +} + +async fn network_denial_message_for_session( + session: Option<&Arc>, + deferred: Option, +) -> String { + let Some(session) = session else { + return NETWORK_ACCESS_DENIED_MESSAGE.to_string(); + }; + match finish_deferred_network_approval(session.as_ref(), deferred).await { + Ok(()) => NETWORK_ACCESS_DENIED_MESSAGE.to_string(), + Err(err) => network_approval_error_message(err), + } +} + +async fn wait_for_late_network_denial(network_cancelled: Option) -> bool { + let Some(network_cancelled) = network_cancelled else { + return false; + }; + if network_cancelled.is_cancelled() { + return true; + } + + tokio::select! { + _ = network_cancelled.cancelled() => true, + _ = tokio::time::sleep(LATE_NETWORK_DENIAL_GRACE_PERIOD) => false, + } +} + +async fn finish_deferred_network_approval_after_process_exit_for_session( + session: Option<&Arc>, + deferred: Option, +) -> Result<(), String> { + wait_for_late_network_denial( + deferred + .as_ref() + .map(DeferredNetworkApproval::cancellation_token), + ) + .await; + finish_deferred_network_approval_for_session(session, deferred).await +} + +fn fail_process_with_message(process: &UnifiedExecProcess, message: String) -> UnifiedExecError { + if let Some(message) = process.failure_message() { + process.terminate(); + return UnifiedExecError::process_failed(message); + } + + process.fail_and_terminate(message.clone()); + UnifiedExecError::process_failed(process.failure_message().unwrap_or(message)) +} + +#[allow(clippy::too_many_arguments)] +async fn emit_failed_initial_exec_end_if_unstored( + process_started_alive: bool, + context: &UnifiedExecContext, + request: &ExecCommandRequest, + cwd: AbsolutePathBuf, + transcript: Arc>, + fallback_output: String, + message: String, + wall_time: Duration, +) { + if process_started_alive { + return; + } + + emit_failed_exec_end_for_unified_exec( + Arc::clone(&context.session), + Arc::clone(&context.turn), + context.call_id.clone(), + request.command.clone(), + cwd, + Some(request.process_id.to_string()), + transcript, + fallback_output, + message, + wall_time, + ) + .await; +} + +fn terminate_process_on_network_denial( + process: Arc, + session: std::sync::Weak, + deferred: DeferredNetworkApproval, +) { + let network_cancelled = deferred.cancellation_token(); + let process_exited = process.cancellation_token(); + tokio::spawn(async move { + let denied = tokio::select! { + _ = network_cancelled.cancelled() => true, + _ = process_exited.cancelled() => { + wait_for_late_network_denial(Some(network_cancelled.clone())).await + } + }; + if !denied { + return; + } + let session = session.upgrade(); + let message = network_denial_message_for_session(session.as_ref(), Some(deferred)).await; + process.fail_and_terminate(message); + }); +} + impl UnifiedExecProcessManager { pub(crate) async fn allocate_process_id(&self) -> i32 { loop { @@ -212,19 +362,7 @@ impl UnifiedExecProcessManager { store.remove(process_id) }; if let Some(entry) = removed { - Self::unregister_network_approval_for_entry(&entry).await; - } - } - - async fn unregister_network_approval_for_entry(entry: &ProcessEntry) { - if let Some(network_approval_id) = entry.network_approval_id.as_deref() - && let Some(session) = entry.session.upgrade() - { - session - .services - .network_approval - .unregister_call(network_approval_id) - .await; + unregister_network_approval_for_entry(&entry).await; } } @@ -250,6 +388,13 @@ impl UnifiedExecProcessManager { return Err(err); } }; + if let Some(deferred) = deferred_network_approval.as_ref() { + terminate_process_on_network_denial( + Arc::clone(&process), + Arc::downgrade(&context.session), + deferred.clone(), + ); + } let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); let event_ctx = ToolEventCtx::new( @@ -272,9 +417,6 @@ impl UnifiedExecProcessManager { // turn cannot drop the last Arc and terminate the background process. let process_started_alive = !process.has_exited() && process.exit_code().is_none(); if process_started_alive { - let network_approval_id = deferred_network_approval - .as_ref() - .map(|deferred| deferred.registration_id().to_string()); self.store_process( Arc::clone(&process), context, @@ -284,7 +426,7 @@ impl UnifiedExecProcessManager { start, request.process_id, request.tty, - network_approval_id, + deferred_network_approval.clone(), Arc::clone(&transcript), ) .await; @@ -320,27 +462,50 @@ impl UnifiedExecProcessManager { let text = String::from_utf8_lossy(&collected).to_string(); let chunk_id = generate_chunk_id(); - if let Some(message) = process.failure_message() { - if !process_started_alive { - emit_failed_exec_end_for_unified_exec( - Arc::clone(&context.session), - Arc::clone(&context.turn), - context.call_id.clone(), - request.command.clone(), - cwd.clone(), - Some(request.process_id.to_string()), - Arc::clone(&transcript), - message.clone(), - wall_time, - ) - .await; - } - self.release_process_id(request.process_id).await; - finish_deferred_network_approval( - context.session.as_ref(), + if deferred_network_approval + .as_ref() + .is_some_and(DeferredNetworkApproval::is_cancelled) + { + let message = network_denial_message_for_session( + Some(&context.session), deferred_network_approval.take(), ) .await; + emit_failed_initial_exec_end_if_unstored( + process_started_alive, + context, + &request, + cwd.clone(), + Arc::clone(&transcript), + text.clone(), + message.clone(), + wall_time, + ) + .await; + self.release_process_id(request.process_id).await; + return Err(fail_process_with_message(process.as_ref(), message)); + } + if let Some(message) = process.failure_message() { + let finish_result = finish_deferred_network_approval_for_session( + Some(&context.session), + deferred_network_approval.take(), + ) + .await; + emit_failed_initial_exec_end_if_unstored( + process_started_alive, + context, + &request, + cwd.clone(), + Arc::clone(&transcript), + text.clone(), + message.clone(), + wall_time, + ) + .await; + self.release_process_id(request.process_id).await; + if let Err(message) = finish_result { + return Err(fail_process_with_message(process.as_ref(), message)); + } return Err(UnifiedExecError::process_failed(message)); } let process_id = request.process_id; @@ -351,7 +516,16 @@ impl UnifiedExecProcessManager { process_id, .. } => (Some(process_id), exit_code), - ProcessStatus::Exited { exit_code, .. } => { + ProcessStatus::Exited { exit_code, entry } => { + if let Err(message) = + finish_deferred_network_approval_after_process_exit_for_session( + Some(&context.session), + deferred_network_approval.take(), + ) + .await + { + return Err(fail_process_with_message(entry.process.as_ref(), message)); + } process.check_for_sandbox_denial_with_text(&text).await?; (None, exit_code) } @@ -363,6 +537,26 @@ impl UnifiedExecProcessManager { // Short‑lived command: emit ExecCommandEnd immediately using the // same helper as the background watcher, so all end events share // one implementation. + let finish_result = finish_deferred_network_approval_after_process_exit_for_session( + Some(&context.session), + deferred_network_approval.take(), + ) + .await; + if let Err(message) = finish_result { + emit_failed_initial_exec_end_if_unstored( + process_started_alive, + context, + &request, + cwd.clone(), + Arc::clone(&transcript), + text.clone(), + message.clone(), + wall_time, + ) + .await; + self.release_process_id(request.process_id).await; + return Err(fail_process_with_message(process.as_ref(), message)); + } let exit_code = process.exit_code(); let exit = exit_code.unwrap_or(-1); emit_exec_end_for_unified_exec( @@ -380,11 +574,6 @@ impl UnifiedExecProcessManager { .await; self.release_process_id(request.process_id).await; - finish_deferred_network_approval( - context.session.as_ref(), - deferred_network_approval.take(), - ) - .await; process.check_for_sandbox_denial_with_text(&text).await?; (None, exit_code) }; @@ -419,6 +608,8 @@ impl UnifiedExecProcessManager { output_closed_notify, cancellation_token, pause_state, + session, + network_approval, hook_command, process_id, tty, @@ -478,8 +669,26 @@ impl UnifiedExecProcessManager { let text = String::from_utf8_lossy(&collected).to_string(); let original_token_count = approx_token_count(&text); let chunk_id = generate_chunk_id(); - if let Some(message) = process.failure_message() { + if network_approval + .as_ref() + .is_some_and(DeferredNetworkApproval::is_cancelled) + { + let message = + network_denial_message_for_session(session.as_ref(), network_approval.clone()) + .await; self.release_process_id(process_id).await; + return Err(fail_process_with_message(process.as_ref(), message)); + } + if let Some(message) = process.failure_message() { + let finish_result = finish_deferred_network_approval_for_session( + session.as_ref(), + network_approval.clone(), + ) + .await; + self.release_process_id(process_id).await; + if let Err(message) = finish_result { + return Err(fail_process_with_message(process.as_ref(), message)); + } return Err(UnifiedExecError::process_failed(message)); } @@ -500,6 +709,11 @@ impl UnifiedExecProcessManager { } => (Some(process_id), exit_code, call_id), ProcessStatus::Exited { exit_code, entry } => { let call_id = entry.call_id.clone(); + if let Err(message) = + finish_network_approval_after_process_exit_for_entry(&entry).await + { + return Err(fail_process_with_message(entry.process.as_ref(), message)); + } (None, exit_code, call_id) } ProcessStatus::Unknown => { @@ -525,7 +739,7 @@ impl UnifiedExecProcessManager { } async fn refresh_process_state(&self, process_id: i32) -> ProcessStatus { - let status = { + { let mut store = self.process_store.lock().await; let Some(entry) = store.processes.get(&process_id) else { return ProcessStatus::Unknown; @@ -549,11 +763,7 @@ impl UnifiedExecProcessManager { process_id, } } - }; - if let ProcessStatus::Exited { entry, .. } = &status { - Self::unregister_network_approval_for_entry(entry).await; } - status } async fn prepare_process_handles( @@ -577,6 +787,7 @@ impl UnifiedExecProcessManager { .session .upgrade() .map(|session| session.subscribe_out_of_band_elicitation_pause_state()); + let session = entry.session.upgrade(); Ok(PreparedProcessHandles { process: Arc::clone(&entry.process), @@ -586,6 +797,8 @@ impl UnifiedExecProcessManager { output_closed_notify, cancellation_token, pause_state, + session, + network_approval: entry.network_approval.clone(), hook_command: entry.hook_command.clone(), process_id: entry.process_id, tty: entry.tty, @@ -603,7 +816,7 @@ impl UnifiedExecProcessManager { started_at: Instant, process_id: i32, tty: bool, - network_approval_id: Option, + network_approval: Option, transcript: Arc>, ) { let entry = ProcessEntry { @@ -612,7 +825,7 @@ impl UnifiedExecProcessManager { process_id, hook_command, tty, - network_approval_id, + network_approval, session: Arc::downgrade(&context.session), last_used: started_at, }; @@ -625,7 +838,7 @@ impl UnifiedExecProcessManager { // prune_processes_if_needed runs while holding process_store; do async // network-approval cleanup only after dropping that lock. if let Some(pruned_entry) = pruned_entry { - Self::unregister_network_approval_for_entry(&pruned_entry).await; + unregister_network_approval_for_entry(&pruned_entry).await; pruned_entry.process.terminate(); } @@ -1041,7 +1254,7 @@ impl UnifiedExecProcessManager { }; for entry in entries { - Self::unregister_network_approval_for_entry(&entry).await; + unregister_network_approval_for_entry(&entry).await; entry.process.terminate(); } } diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 18930afb6..0c5b71416 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -135,6 +135,94 @@ fn exec_server_process_id_matches_unified_exec_process_id() { assert_eq!(exec_server_process_id(/*process_id*/ 4321), "4321"); } +#[tokio::test] +async fn network_denial_fallback_message_names_sandbox_network_proxy() { + let message = network_denial_message_for_session(/*session*/ None, /*deferred*/ None).await; + + assert_eq!( + message, + "Network access was denied by the Codex sandbox network proxy." + ); +} + +#[tokio::test] +async fn late_network_denial_grace_observes_cancellation_after_exit() { + let cancellation = CancellationToken::new(); + let cancellation_for_task = cancellation.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(10)).await; + cancellation_for_task.cancel(); + }); + + assert!(wait_for_late_network_denial(Some(cancellation)).await); +} + +#[tokio::test] +async fn failed_initial_end_for_unstored_process_uses_fallback_output() { + let (session, turn, rx_event) = crate::session::tests::make_session_and_context_with_rx().await; + let context = UnifiedExecContext::new( + Arc::clone(&session), + Arc::clone(&turn), + "call-unified-denied".to_string(), + ); + let request = ExecCommandRequest { + command: vec![ + "sh".to_string(), + "-lc".to_string(), + "echo before".to_string(), + ], + hook_command: "echo before".to_string(), + process_id: 123, + yield_time_ms: 1000, + max_output_tokens: None, + workdir: None, + network: None, + tty: true, + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + additional_permissions_preapproved: false, + justification: None, + prefix_rule: None, + }; + + let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); + transcript + .lock() + .await + .push_chunk(b"PARTIAL_TRANSCRIPT".to_vec()); + + emit_failed_initial_exec_end_if_unstored( + /*process_started_alive*/ false, + &context, + &request, + turn.cwd.clone(), + transcript, + "PRE_DENIAL_MARKER".to_string(), + "Network access denied".to_string(), + Duration::from_millis(7), + ) + .await; + + let event = tokio::time::timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("timed out waiting for failed exec end event") + .expect("event channel closed"); + let codex_protocol::protocol::EventMsg::ExecCommandEnd(end_event) = event.msg else { + panic!("expected ExecCommandEnd event"); + }; + assert_eq!(end_event.call_id, "call-unified-denied"); + assert_eq!( + end_event.status, + codex_protocol::protocol::ExecCommandStatus::Failed + ); + assert_eq!(end_event.exit_code, -1); + assert_eq!(end_event.process_id.as_deref(), Some("123")); + assert_eq!( + end_event.aggregated_output, + "PRE_DENIAL_MARKER\nNetwork access denied" + ); +} + #[test] fn pruning_prefers_exited_processes_outside_recently_used() { let now = Instant::now(); diff --git a/codex-rs/core/src/unified_exec/process_tests.rs b/codex-rs/core/src/unified_exec/process_tests.rs index f9e8a6caa..ee0ea1cba 100644 --- a/codex-rs/core/src/unified_exec/process_tests.rs +++ b/codex-rs/core/src/unified_exec/process_tests.rs @@ -112,6 +112,20 @@ async fn remote_write_closed_stdin_marks_process_exited() { assert!(process.has_exited()); } +#[tokio::test] +async fn fail_and_terminate_preserves_failure_message() { + let process = remote_process(WriteStatus::Accepted).await; + + process.fail_and_terminate("network denied".to_string()); + process.fail_and_terminate("second failure".to_string()); + + assert!(process.has_exited()); + assert_eq!( + process.failure_message(), + Some("network denied".to_string()) + ); +} + #[tokio::test] async fn remote_process_waits_for_early_exit_event() { let (wake_tx, _wake_rx) = watch::channel(0); diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 9c3272882..c9ad84649 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -7,10 +7,12 @@ use anyhow::Context; use anyhow::Result; use codex_exec_server::CreateDirectoryOptions; use codex_features::Feature; +use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecCommandSource; +use codex_protocol::protocol::ExecCommandStatus; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::UserInput; @@ -234,10 +236,9 @@ async fn unified_exec_intercepts_apply_patch_exec_command() -> Result<()> { let builder = test_codex().with_config(|config| { config.include_apply_patch_tool = true; config.use_experimental_unified_exec_tool = true; - config - .features - .enable(Feature::UnifiedExec) - .expect("test config should allow feature update"); + if let Err(err) = config.features.enable(Feature::UnifiedExec) { + panic!("test config should allow feature update: {err}"); + } }); let harness = TestCodexHarness::with_builder(builder).await?; @@ -781,6 +782,231 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_network_denial_emits_failed_background_end_event() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_windows!(Ok(())); + + let server = start_mock_server().await; + let (test, sandbox_policy) = unified_exec_network_denial_test(&server).await?; + + let call_id = "uexec-network-denied"; + let args = json!({ + "cmd": "python3 -c \"import os, socket, time, urllib.parse; time.sleep(0.3); proxy = urllib.parse.urlparse(os.environ['HTTP_PROXY']); sock = socket.create_connection((proxy.hostname, proxy.port), timeout=2); sock.sendall(b'GET http://codex-network-denied.invalid/ HTTP/1.1\\r\\nHost: codex-network-denied.invalid\\r\\n\\r\\n'); sock.recv(1024); time.sleep(5)\"", + "yield_time_ms": 50, + }); + let response_mock = + mount_unified_exec_network_denial_responses(&server, call_id, &args).await?; + + submit_unified_exec_turn(&test, "exercise network denial", sandbox_policy).await?; + + let (end_event, turn_completed) = + wait_for_unified_exec_end(&test, call_id, &response_mock).await; + + assert_eq!(end_event.status, ExecCommandStatus::Failed); + assert_eq!(end_event.exit_code, -1); + assert!( + end_event.aggregated_output.contains("Network access"), + "expected network denial message in aggregated output: {:?}", + end_event.aggregated_output + ); + assert!( + end_event.process_id.is_some(), + "background denial should end the stored unified exec process" + ); + + if !turn_completed { + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_short_lived_network_denial_emits_failed_end_event() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_windows!(Ok(())); + + let server = start_mock_server().await; + let (test, sandbox_policy) = unified_exec_network_denial_test(&server).await?; + + let call_id = "uexec-short-network-denied"; + let args = json!({ + "cmd": "python3 -c \"import os, socket, urllib.parse; proxy = urllib.parse.urlparse(os.environ['HTTP_PROXY']); sock = socket.create_connection((proxy.hostname, proxy.port), timeout=2); sock.sendall(b'GET http://codex-short-network-denied.invalid/ HTTP/1.1\\r\\nHost: codex-short-network-denied.invalid\\r\\n\\r\\n'); sock.recv(1024)\"", + "yield_time_ms": 1000, + }); + let response_mock = + mount_unified_exec_network_denial_responses(&server, call_id, &args).await?; + + submit_unified_exec_turn(&test, "exercise short network denial", sandbox_policy).await?; + + let (end_event, turn_completed) = + wait_for_unified_exec_end(&test, call_id, &response_mock).await; + + assert_eq!(end_event.status, ExecCommandStatus::Failed); + assert_eq!(end_event.exit_code, -1); + assert!( + end_event.aggregated_output.contains("Network access"), + "expected network denial message in aggregated output: {:?}", + end_event.aggregated_output + ); + assert!( + end_event.process_id.is_some(), + "short-lived denial should still emit an end event for the command" + ); + + if !turn_completed { + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + } + Ok(()) +} + +#[allow(clippy::expect_used)] +async fn unified_exec_network_denial_test( + server: &wiremock::MockServer, +) -> Result<(TestCodex, SandboxPolicy)> { + use codex_config::ConfigLayerStack; + use codex_config::ConfigLayerStackOrdering; + use codex_config::Constrained; + use codex_config::NetworkConstraints; + use codex_config::NetworkRequirementsToml; + use codex_config::RequirementSource; + use codex_config::Sourced; + use std::sync::Arc; + use tempfile::TempDir; + + let home = Arc::new(TempDir::new()?); + fs::write( + home.path().join("config.toml"), + r#"default_permissions = "workspace" + +[permissions.workspace.filesystem] +":minimal" = "read" + +[permissions.workspace.network] +enabled = true +mode = "limited" +allow_local_binding = true +"#, + )?; + let mut sandbox_policy = SandboxPolicy::new_workspace_write_policy(); + if let SandboxPolicy::WorkspaceWrite { network_access, .. } = &mut sandbox_policy { + *network_access = true; + } + let sandbox_policy_for_config = sandbox_policy.clone(); + let mut builder = test_codex().with_home(home).with_config(move |config| { + config.use_experimental_unified_exec_tool = true; + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow feature update"); + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never); + config.permissions.permission_profile = Constrained::allow_any( + PermissionProfile::from_legacy_sandbox_policy(&sandbox_policy_for_config), + ); + let layers = config + .config_layer_stack + .get_layers( + ConfigLayerStackOrdering::LowestPrecedenceFirst, + /*include_disabled*/ true, + ) + .into_iter() + .cloned() + .collect(); + let mut requirements = config.config_layer_stack.requirements().clone(); + requirements.network = Some(Sourced::new( + NetworkConstraints { + enabled: Some(true), + allow_local_binding: Some(true), + ..Default::default() + }, + RequirementSource::CloudRequirements, + )); + let mut requirements_toml = config.config_layer_stack.requirements_toml().clone(); + requirements_toml.network = Some(NetworkRequirementsToml { + enabled: Some(true), + allow_local_binding: Some(true), + ..Default::default() + }); + config.config_layer_stack = + match ConfigLayerStack::new(layers, requirements, requirements_toml) { + Ok(stack) => stack, + Err(err) => panic!("rebuild config layer stack with network requirements: {err}"), + }; + }); + let test = builder.build_remote_aware(server).await?; + assert!( + test.config.permissions.network.is_some(), + "expected managed network proxy config to be present" + ); + + Ok((test, sandbox_policy)) +} + +async fn mount_unified_exec_network_denial_responses( + server: &wiremock::MockServer, + call_id: &str, + args: &Value, +) -> Result { + let responses = vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "exec_command", &serde_json::to_string(&args)?), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "finished"), + ev_completed("resp-2"), + ]), + ]; + Ok(mount_sse_sequence(server, responses).await) +} + +async fn wait_for_unified_exec_end( + test: &TestCodex, + call_id: &str, + response_mock: &core_test_support::responses::ResponseMock, +) -> (codex_protocol::protocol::ExecCommandEndEvent, bool) { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(15); + let mut observed_events = Vec::new(); + let mut turn_completed = false; + let end_event = loop { + let remaining = deadline + .checked_duration_since(std::time::Instant::now()) + .unwrap_or_default(); + if remaining.is_zero() { + panic!( + "timed out waiting for network denial end event; observed {observed_events:?}; response requests: {}", + response_mock.requests().len() + ); + } + let event = match tokio::time::timeout(remaining, test.codex.next_event()).await { + Ok(Ok(event)) => event.msg, + Ok(Err(err)) => panic!("event stream ended unexpectedly: {err}"), + Err(_) => panic!( + "timed out waiting for network denial end event; observed {observed_events:?}; response requests: {}", + response_mock.requests().len() + ), + }; + turn_completed |= matches!(event, EventMsg::TurnComplete(_)); + observed_events.push(format!("{event:?}")); + if let EventMsg::ExecCommandEnd(ev) = event + && ev.call_id == call_id + { + break ev; + } + }; + (end_event, turn_completed) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn unified_exec_emits_terminal_interaction_for_write_stdin() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/memories/write/templates/memories/consolidation.md b/codex-rs/memories/write/templates/memories/consolidation.md index 96e9b5ddc..cc586bd4c 100644 --- a/codex-rs/memories/write/templates/memories/consolidation.md +++ b/codex-rs/memories/write/templates/memories/consolidation.md @@ -155,7 +155,7 @@ Incremental update and forgetting mechanism: - Use the git-style diff in `{{ phase2_workspace_diff_file }}` to identify relevant changed sections and deleted inputs. -- Every changes in `{{ phase2_workspace_diff_file }}` are authoritative and must propagated and consolidated. If a +- Every changes in `{{ phase2_workspace_diff_file }}` are authoritative and must propagated and consolidated. If a changes appears to be randomly placed in the files, it is probably a user change and you shouldn't just drop it. Make sure to add it to the overall memories consolidation - Do not open raw sessions / original rollout transcripts.