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.