diff --git a/codex-rs/core/src/tools/handlers/shell_tests.rs b/codex-rs/core/src/tools/handlers/shell_tests.rs index b69f3be230..d0d3d7d278 100644 --- a/codex-rs/core/src/tools/handlers/shell_tests.rs +++ b/codex-rs/core/src/tools/handlers/shell_tests.rs @@ -1,10 +1,17 @@ use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use codex_protocol::models::ShellCommandToolCallParams; +use codex_protocol::models::function_call_output_content_items_to_text; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::ExecCommandStatus; +use codex_protocol::protocol::ReviewDecision; use pretty_assertions::assert_eq; use crate::codex::make_session_and_context; +use crate::codex::make_session_and_context_with_rx; use crate::exec_env::create_env; use crate::is_safe_command::is_known_safe_command; use crate::powershell::try_find_powershell_executable_blocking; @@ -13,8 +20,15 @@ use crate::sandboxing::SandboxPermissions; use crate::shell::Shell; use crate::shell::ShellType; use crate::shell_snapshot::ShellSnapshot; +use crate::state::ActiveTurn; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; use crate::tools::handlers::ShellCommandHandler; +use crate::tools::handlers::ShellHandler; +use crate::tools::registry::ToolHandler; +use crate::turn_diff_tracker::TurnDiffTracker; use tokio::sync::watch; +use tokio::time::timeout; /// The logic for is_known_safe_command() has heuristics for known shells, /// so we must ensure the commands generated by [ShellCommandHandler] can be @@ -63,6 +77,10 @@ fn assert_safe(shell: &Shell, command: &str) { )); } +fn tool_output_text(output: &crate::tools::context::FunctionToolOutput) -> String { + function_call_output_content_items_to_text(&output.body).expect("tool output text") +} + #[tokio::test] async fn shell_command_handler_to_exec_params_uses_session_shell_and_turn_context() { let (session, turn_context) = make_session_and_context().await; @@ -178,3 +196,303 @@ fn shell_command_handler_rejects_login_when_disallowed() { "unexpected error: {err}" ); } + +#[tokio::test] +async fn shell_handler_retry_emits_single_begin_before_approval_request() { + let (session, mut turn_context, rx) = make_session_and_context_with_rx().await; + *session.active_turn.lock().await = Some(ActiveTurn::default()); + Arc::get_mut(&mut turn_context) + .expect("single turn context ref") + .approval_policy + .set(AskForApproval::OnFailure) + .expect("test setup should allow updating approval policy"); + + let temp_dir = tempfile::tempdir_in( + turn_context + .cwd + .parent() + .expect("workspace cwd should have a parent"), + ) + .expect("create temp dir outside workspace"); + let file_path = temp_dir.path().join("outside-sandbox.txt"); + let expected_stdout = "hello from outside sandbox\n".to_string(); + std::fs::write(&file_path, expected_stdout.as_bytes()).expect("write readable file"); + + let command = if cfg!(windows) { + vec![ + "cmd.exe".to_string(), + "/C".to_string(), + format!("type \"{}\"", file_path.display()), + ] + } else { + vec!["/bin/cat".to_string(), file_path.display().to_string()] + }; + + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let call_id = "call-1".to_string(); + let handler = ShellHandler; + let handle: tokio::task::JoinHandle< + Result, + > = tokio::spawn({ + let session = Arc::clone(&session); + let turn_context = Arc::clone(&turn_context); + let tracker = Arc::clone(&tracker); + let command = command.clone(); + let call_id = call_id.clone(); + async move { + handler + .handle(ToolInvocation { + session, + turn: turn_context.clone(), + tracker, + call_id: call_id.clone(), + tool_name: "shell".to_string(), + tool_namespace: None, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "command": command, + "workdir": Some(turn_context.cwd.to_string_lossy().to_string()), + "timeout_ms": 1_000u64, + "sandbox_permissions": SandboxPermissions::UseDefault, + }) + .to_string(), + }, + }) + .await + } + }); + + let mut begin_count = 0usize; + let mut begin_idx = None; + let mut event_idx = 0usize; + let approval_idx = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for approval flow event") + .expect("approval flow event missing"); + match event.msg { + EventMsg::ExecCommandBegin(begin) if begin.call_id == call_id => { + begin_count += 1; + if begin_idx.is_none() { + begin_idx = Some(event_idx); + } + } + EventMsg::ExecApprovalRequest(request) if request.call_id == call_id => { + break event_idx; + } + _ => {} + } + event_idx += 1; + }; + + assert_eq!(begin_count, 1); + assert_eq!(begin_idx, Some(0)); + assert!( + approval_idx > begin_idx.expect("begin event should exist"), + "approval request should follow the begin event" + ); + + session + .notify_approval(&call_id, ReviewDecision::Approved) + .await; + + let end_event = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for exec end event") + .expect("exec end event missing"); + match event.msg { + EventMsg::ExecCommandBegin(begin) if begin.call_id == call_id => { + begin_count += 1; + } + EventMsg::ExecCommandEnd(end) if end.call_id == call_id => { + break end; + } + _ => {} + } + }; + + let _response = timeout(Duration::from_secs(5), handle) + .await + .expect("timed out waiting for shell handler") + .expect("shell handler join error") + .expect("shell handler should succeed"); + + assert_eq!(begin_count, 1); + assert_eq!(end_event.status, ExecCommandStatus::Completed); + assert_eq!(end_event.command, command); + assert_eq!(end_event.stdout, expected_stdout); + assert_eq!(end_event.exit_code, 0); +} + +#[tokio::test] +async fn shell_handler_retry_override_keeps_single_begin_and_reports_override() { + let (session, mut turn_context, rx) = make_session_and_context_with_rx().await; + *session.active_turn.lock().await = Some(ActiveTurn::default()); + Arc::get_mut(&mut turn_context) + .expect("single turn context ref") + .approval_policy + .set(AskForApproval::OnFailure) + .expect("test setup should allow updating approval policy"); + + let temp_dir = tempfile::tempdir_in( + turn_context + .cwd + .parent() + .expect("workspace cwd should have a parent"), + ) + .expect("create temp dir outside workspace"); + let file_path = temp_dir.path().join("outside-sandbox.txt"); + std::fs::write(&file_path, b"should not run\n").expect("write readable file"); + + let requested_command = if cfg!(windows) { + vec![ + "cmd.exe".to_string(), + "/C".to_string(), + format!("type \"{}\"", file_path.display()), + ] + } else { + vec!["/bin/cat".to_string(), file_path.display().to_string()] + }; + let override_command = if cfg!(windows) { + vec![ + "cmd.exe".to_string(), + "/C".to_string(), + "echo override-after-retry".to_string(), + ] + } else { + vec![ + "/bin/sh".to_string(), + "-c".to_string(), + "printf 'override-after-retry\\n'".to_string(), + ] + }; + let expected_stdout = "override-after-retry\n".to_string(); + + let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let call_id = "call-retry-override".to_string(); + let handler = ShellHandler; + let handle: tokio::task::JoinHandle< + Result, + > = tokio::spawn({ + let session = Arc::clone(&session); + let turn_context = Arc::clone(&turn_context); + let tracker = Arc::clone(&tracker); + let requested_command = requested_command.clone(); + let call_id = call_id.clone(); + async move { + handler + .handle(ToolInvocation { + session, + turn: turn_context.clone(), + tracker, + call_id: call_id.clone(), + tool_name: "shell".to_string(), + tool_namespace: None, + payload: ToolPayload::Function { + arguments: serde_json::json!({ + "command": requested_command, + "workdir": Some(turn_context.cwd.to_string_lossy().to_string()), + "timeout_ms": 1_000u64, + "sandbox_permissions": SandboxPermissions::UseDefault, + }) + .to_string(), + }, + }) + .await + } + }); + + let mut begin_count = 0usize; + let begin_event = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for begin event") + .expect("begin event missing"); + if let EventMsg::ExecCommandBegin(begin) = event.msg + && begin.call_id == call_id + { + begin_count += 1; + break begin; + } + }; + + let approval_event = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for retry approval request") + .expect("retry approval request missing"); + match event.msg { + EventMsg::ExecCommandBegin(begin) if begin.call_id == call_id => { + begin_count += 1; + } + EventMsg::ExecApprovalRequest(request) if request.call_id == call_id => { + break request; + } + _ => {} + } + }; + + assert_eq!(begin_count, 1); + assert_eq!(begin_event.command, requested_command); + assert_eq!(approval_event.command, requested_command); + + session + .notify_approval( + &call_id, + ReviewDecision::ApprovedOverrideCommand { + command: override_command.clone(), + }, + ) + .await; + + let end_event = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for exec end event") + .expect("exec end event missing"); + match event.msg { + EventMsg::ExecCommandBegin(begin) if begin.call_id == call_id => { + begin_count += 1; + } + EventMsg::ExecCommandEnd(end) if end.call_id == call_id => { + break end; + } + _ => {} + } + }; + + let response = timeout(Duration::from_secs(5), handle) + .await + .expect("timed out waiting for shell handler") + .expect("shell handler join error") + .expect("shell handler should succeed"); + + let duration_seconds = (end_event.duration.as_secs_f32() * 10.0).round() / 10.0; + let expected_body = serde_json::json!({ + "output": expected_stdout, + "metadata": { + "exit_code": 0, + "duration_seconds": duration_seconds, + }, + "executed_command": override_command, + }); + + assert_eq!(begin_count, 1); + assert_eq!(end_event.status, ExecCommandStatus::Completed); + assert_eq!( + end_event.command, + expected_body["executed_command"] + .as_array() + .expect("array") + .iter() + .map(|value| value.as_str().expect("string").to_string()) + .collect::>() + ); + assert_eq!(end_event.stdout, expected_stdout); + assert_eq!( + serde_json::from_str::(&tool_output_text(&response)) + .expect("valid json"), + expected_body + ); +} diff --git a/codex-rs/core/src/tools/network_approval_tests.rs b/codex-rs/core/src/tools/network_approval_tests.rs index 44fef07f1f..e633013a08 100644 --- a/codex-rs/core/src/tools/network_approval_tests.rs +++ b/codex-rs/core/src/tools/network_approval_tests.rs @@ -1,9 +1,47 @@ use super::*; +use crate::codex::make_session_and_context_with_rx; +use crate::protocol::EventMsg; +use crate::protocol::TurnAbortReason; +use crate::state::TaskKind; +use crate::tasks::SessionTask; +use crate::tasks::SessionTaskContext; +use async_trait::async_trait; use codex_network_proxy::BlockedRequestArgs; +use codex_network_proxy::NetworkPolicyRequest; +use codex_network_proxy::NetworkPolicyRequestArgs; +use codex_network_proxy::NetworkProtocol; use codex_protocol::approvals::ExecPolicyAmendment; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::ReviewDecision; +use codex_protocol::user_input::UserInput; use pretty_assertions::assert_eq; +use std::time::Duration; +use tokio::time::timeout; +use tokio_util::sync::CancellationToken; + +struct WaitForCancellationTask; + +#[async_trait] +impl SessionTask for WaitForCancellationTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "network-approval-test" + } + + async fn run( + self: Arc, + _session: Arc, + _ctx: Arc, + _input: Vec, + cancellation_token: CancellationToken, + ) -> Option { + cancellation_token.cancelled().await; + None + } +} #[tokio::test] async fn pending_approvals_are_deduped_per_host_protocol_and_port() { @@ -158,6 +196,79 @@ fn network_review_rejects_command_and_execpolicy_overrides() { ); } +#[tokio::test] +async fn inline_network_review_rejects_command_override_at_runtime() { + let service = Arc::new(NetworkApprovalService::default()); + let (session, turn_context, rx) = make_session_and_context_with_rx().await; + service.register_call("registration-1".to_string()).await; + session + .spawn_task( + Arc::clone(&turn_context), + vec![UserInput::Text { + text: "need network".to_string(), + text_elements: Vec::new(), + }], + WaitForCancellationTask, + ) + .await; + + let request = NetworkPolicyRequest::new(NetworkPolicyRequestArgs { + protocol: NetworkProtocol::Http, + host: "example.com".to_string(), + port: 80, + client_addr: None, + method: Some("GET".to_string()), + command: None, + exec_policy_hint: None, + }); + + let handle = tokio::spawn({ + let service = Arc::clone(&service); + let session = Arc::clone(&session); + async move { service.handle_inline_policy_request(session, request).await } + }); + + let approval_request = loop { + let event = timeout(Duration::from_secs(5), rx.recv()) + .await + .expect("timed out waiting for approval request") + .expect("approval request missing"); + if let EventMsg::ExecApprovalRequest(request) = event.msg { + break request; + } + }; + + assert_eq!( + approval_request.command, + vec![ + "network-access".to_string(), + "http://example.com:80".to_string(), + ] + ); + + session + .notify_approval( + &approval_request.call_id, + ReviewDecision::ApprovedOverrideCommand { + command: vec!["echo".to_string(), "override".to_string()], + }, + ) + .await; + + let decision = timeout(Duration::from_secs(5), handle) + .await + .expect("timed out waiting for inline network decision") + .expect("inline network decision join error"); + + assert_eq!(decision, NetworkDecision::deny("not_allowed")); + assert_eq!( + service.take_call_outcome("registration-1").await, + Some(NetworkApprovalOutcome::DeniedByUser) + ); + + session.abort_all_tasks(TurnAbortReason::Interrupted).await; +} + fn denied_blocked_request(host: &str) -> BlockedRequest { BlockedRequest::new(BlockedRequestArgs { host: host.to_string(), 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 95d31a2135..d02ac3dbd2 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -549,7 +549,7 @@ impl CoreShellActionProvider { } ReviewDecision::ApprovedOverrideCommand { .. } => { EscalationDecision::deny(Some( - "command override is not supported for execve approvals" + "Command override is not supported for execve approvals" .to_string(), )) }