add test coverage

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<crate::tools::context::FunctionToolOutput, crate::function_tool::FunctionCallError>,
+    > = 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<crate::tools::context::FunctionToolOutput, crate::function_tool::FunctionCallError>,
+    > = 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::<Vec<_>>()
+    );
+    assert_eq!(end_event.stdout, expected_stdout);
+    assert_eq!(
+        serde_json::from_str::<serde_json::Value>(&tool_output_text(&response))
+            .expect("valid json"),
+        expected_body
+    );
+}

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<Self>,
+        _session: Arc<SessionTaskContext>,
+        _ctx: Arc<crate::codex::TurnContext>,
+        _input: Vec<UserInput>,
+        cancellation_token: CancellationToken,
+    ) -> Option<String> {
+        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(),

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(),
                             ))
                         }