From 7ca5932fd283554bd438659566660801a04d1d99 Mon Sep 17 00:00:00 2001 From: Won Park Date: Fri, 21 Aug 2026 21:41:09 +0000 Subject: [PATCH] Cancel Guardian reviews with their tool calls (#40021) Propagate tool cancellation tokens into Guardian approval reviews so interrupting a tool also aborts its pending review. Apply the same cancellation behavior to server-initiated MCP approval elicitation. Add integration coverage for cancellation from direct tool calls, code-mode turns, background code-mode cells, and MCP elicitation. GitOrigin-RevId: cbee4b422bbcbf6984e415ef6409cee705beebdc --- codex-rs/core/src/guardian/review.rs | 11 +- codex-rs/core/src/mcp_tool_call.rs | 6 + codex-rs/core/src/mcp_tool_call_tests.rs | 10 + codex-rs/core/src/session/mcp.rs | 11 +- codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/tests.rs | 1 + codex-rs/core/src/tools/approvals.rs | 14 +- .../core/src/tools/handlers/apply_patch.rs | 5 + codex-rs/core/src/tools/handlers/mcp.rs | 2 + .../handlers/unified_exec/exec_command.rs | 10 +- codex-rs/core/src/tools/network_approval.rs | 1 + codex-rs/core/src/tools/orchestrator.rs | 4 + .../tools/runtimes/shell/unix_escalation.rs | 2 + codex-rs/core/src/tools/sandboxing.rs | 1 + .../src/unified_exec/async_watcher_tests.rs | 1 + codex-rs/core/src/unified_exec/mod.rs | 10 +- codex-rs/core/src/unified_exec/mod_tests.rs | 1 + .../core/src/unified_exec/process_manager.rs | 1 + .../src/unified_exec/process_manager_tests.rs | 1 + .../tests/suite/guardian_mcp_elicitation.rs | 174 ++++++++++++ .../suite/guardian_review_cancellation.rs | 251 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 3 + 22 files changed, 500 insertions(+), 22 deletions(-) create mode 100644 codex-rs/core/tests/suite/guardian_mcp_elicitation.rs create mode 100644 codex-rs/core/tests/suite/guardian_review_cancellation.rs diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 46b04147e3..143cddd5b3 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -743,7 +743,7 @@ pub(crate) fn spawn_approval_request_review( context: impl Into, review_id: String, request: GuardianApprovalRequest, - retry_reason: Option, + reasons: ApprovalRequestReasons, options: GuardianReviewOptions, ) -> oneshot::Receiver { let context = context.into(); @@ -752,13 +752,8 @@ pub(crate) fn spawn_approval_request_review( let spawn_result = std::thread::Builder::new() .name("codex-approval-review".to_string()) .spawn(move || { - let decision = runtime.block_on(review_approval_request_with_cancel( - &session, - context, - review_id, - request, - retry_reason, - options, + let decision = runtime.block_on(run_guardian_review( + session, context, review_id, request, reasons, options, )); let _ = tx.send(decision); }); diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 8259d623c5..12e51eeb5b 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -82,6 +82,7 @@ use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; use std::sync::Arc; +use tokio_util::sync::CancellationToken; use toml_edit::value; use tracing::Instrument; use tracing::Span; @@ -108,9 +109,11 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; /// Handles the specified tool call and dispatches the appropriate MCP tool-call /// item lifecycle events to the `Session`. +#[expect(clippy::too_many_arguments)] pub(crate) async fn handle_mcp_tool_call( sess: Arc, step_context: &Arc, + cancellation_token: &CancellationToken, call_id: String, tool_info: &ToolInfo, hook_tool_name: HookToolName, @@ -233,6 +236,7 @@ pub(crate) async fn handle_mcp_tool_call( if let Some(decision) = maybe_request_mcp_tool_approval( &sess, step_context, + cancellation_token, &call_id, &invocation, &invocation_tool_name, @@ -1288,6 +1292,7 @@ fn mcp_tool_approval_prompt_options( async fn maybe_request_mcp_tool_approval( sess: &Arc, step_context: &Arc, + cancellation_token: &CancellationToken, call_id: &str, invocation: &McpInvocation, invocation_tool_name: &ToolName, @@ -1375,6 +1380,7 @@ async fn maybe_request_mcp_tool_approval( }; let approval_context = ApprovalContext { review_context: GuardianReviewContext::from(step_context), + cancellation_token: Some(cancellation_token.clone()), call_id: call_id.to_string(), tool_name: invocation_tool_name.clone(), strict_auto_review, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 3fe7e85138..ac9cac8395 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -2322,6 +2322,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-1", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2401,6 +2402,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-guardian", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2462,6 +2464,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-mcp-hook", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2531,6 +2534,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-mcp-hook-no-metadata", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2615,6 +2619,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-mcp-remembered", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2713,6 +2718,7 @@ async fn strict_auto_review_forces_guardian_for_mcp_policy_skip() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-guardian-deny", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2789,6 +2795,7 @@ async fn assert_mcp_user_approval_persistence( maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-mcp-persist", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2873,6 +2880,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-prompt", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -2937,6 +2945,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-2", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), @@ -3029,6 +3038,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { let decision = maybe_request_mcp_tool_approval( &session, &StepContext::for_test(Arc::clone(&turn_context)), + &CancellationToken::new(), "call-3", &invocation, &ToolName::namespaced(&invocation.server, &invocation.tool), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 6b25214762..bf14adae85 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -847,12 +847,17 @@ async fn review_guardian_mcp_elicitation( }; let review_id = crate::guardian::new_guardian_review_id(); - let decision = crate::guardian::review_approval_request( + let decision = crate::guardian::review_approval_request_with_cancel( &session, &turn_context, - review_id.clone(), + review_id, guardian_request, - Default::default(), + /*retry_reason*/ None, + crate::guardian::GuardianReviewOptions { + plugin_attribution_override: None, + approval_request_source: codex_analytics::GuardianApprovalRequestSource::MainTurn, + external_cancel: Some(cancellation_token), + }, ) .await; Ok(Some(mcp_elicitation_response_from_guardian_decision( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 513e7c03eb..195e3cfa0f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2557,6 +2557,7 @@ impl Session { }; let approval_context = ApprovalContext { review_context: review_context.clone(), + cancellation_token: Some(cancellation_token.clone()), call_id, tool_name: ToolName::plain("request_permissions"), strict_auto_review: false, @@ -2570,7 +2571,6 @@ impl Session { decision = self.request_guardian_approval( action, &approval_context, - Some(cancellation_token.clone()), ) => decision, }; let response = match decision { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 4233e5f94d..3a5f2ee5f3 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1339,6 +1339,7 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an let mut orchestrator = crate::tools::orchestrator::ToolOrchestrator::new(); let mut tool = ProbeToolRuntime::default(); let tool_ctx = crate::tools::sandboxing::ToolCtx { + cancellation_token: CancellationToken::new(), session: Arc::clone(&session), step_context: StepContext::for_test(Arc::clone(&turn)), call_id: "probe-call".to_string(), diff --git a/codex-rs/core/src/tools/approvals.rs b/codex-rs/core/src/tools/approvals.rs index 23382bd1fa..1c67322d16 100644 --- a/codex-rs/core/src/tools/approvals.rs +++ b/codex-rs/core/src/tools/approvals.rs @@ -55,6 +55,7 @@ use tracing::warn; #[derive(Clone)] pub(crate) struct ApprovalContext { pub(crate) review_context: GuardianReviewContext, + pub(crate) cancellation_token: Option, pub(crate) call_id: String, pub(crate) tool_name: ToolName, pub(crate) strict_auto_review: bool, @@ -568,10 +569,7 @@ impl Session { }; let decision = match reviewer { - ApprovalReviewer::Guardian => { - self.request_guardian_approval(action, ctx, /*cancellation_token*/ None) - .await - } + ApprovalReviewer::Guardian => self.request_guardian_approval(action, ctx).await, ApprovalReviewer::User => self.request_user_approval(&action, ctx).await, }; let source = match reviewer { @@ -585,7 +583,6 @@ impl Session { self: &Arc, action: ApprovalAction, ctx: &ApprovalContext, - cancellation_token: Option, ) -> ReviewDecision { let is_network_approval = matches!(&action, ApprovalAction::NetworkAccess { .. }); let review_id = new_guardian_review_id(); @@ -599,13 +596,16 @@ impl Session { } }; - if let Some(cancellation_token) = cancellation_token { + if let Some(cancellation_token) = &ctx.cancellation_token { let review = spawn_approval_request_review( Arc::clone(self), ctx.review_context.clone(), review_id, action, - ctx.retry_reason.clone(), + ApprovalRequestReasons { + approval: ctx.approval_reason.clone(), + retry: ctx.retry_reason.clone(), + }, GuardianReviewOptions { plugin_attribution_override: None, approval_request_source: GuardianApprovalRequestSource::MainTurn, diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index c57302c4ce..b474e661ae 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -5,6 +5,7 @@ use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use tokio_util::sync::CancellationToken; use crate::apply_patch; use crate::apply_patch::convert_apply_patch_to_protocol; @@ -370,6 +371,7 @@ impl ApplyPatchHandler { session, turn, step_context, + cancellation_token, tracker, call_id, tool_name, @@ -419,6 +421,7 @@ impl ApplyPatchHandler { let tool_ctx = ToolCtx { session, step_context: Arc::clone(&step_context), + cancellation_token, call_id, tool_name, }; @@ -509,6 +512,7 @@ pub(crate) async fn intercept_apply_patch( turn_environment: TurnEnvironment, session: Arc, step_context: Arc, + cancellation_token: CancellationToken, tracker: Option<&SharedTurnDiffTracker>, call_id: &str, tool_name: &str, @@ -529,6 +533,7 @@ pub(crate) async fn intercept_apply_patch( let tool_ctx = ToolCtx { session, step_context, + cancellation_token, call_id: call_id.to_string(), tool_name: ToolName::plain(tool_name), }; diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 06de617093..1702a4690a 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -173,6 +173,7 @@ impl McpHandler { let ToolInvocation { session, step_context, + cancellation_token, call_id, tool_name, payload, @@ -193,6 +194,7 @@ impl McpHandler { let result = handle_mcp_tool_call( Arc::clone(&session), &step_context, + &cancellation_token, call_id.clone(), &self.tool_info, self.hook_tool_name(), diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index f7a24e86d4..f0a34c6616 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -112,6 +112,7 @@ impl ExecCommandHandler { session, turn, step_context, + cancellation_token, tracker, call_id, payload, @@ -128,8 +129,12 @@ impl ExecCommandHandler { }; let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; - let context = - UnifiedExecContext::new(session.clone(), step_context.clone(), call_id.clone()); + let context = UnifiedExecContext::new( + session.clone(), + step_context.clone(), + cancellation_token, + call_id.clone(), + ); let environment_args: ExecCommandEnvironmentArgs = parse_arguments(&arguments)?; let Some(turn_environment) = resolve_tool_environment( &step_context.environments, @@ -318,6 +323,7 @@ impl ExecCommandHandler { turn_environment.clone(), context.session.clone(), Arc::clone(&context.step_context), + context.cancellation_token.clone(), Some(&tracker), &context.call_id, "exec_command", diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index d32e849bfa..e9d65fb6e7 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -779,6 +779,7 @@ impl NetworkApprovalService { }; let approval_context = ApprovalContext { review_context: GuardianReviewContext::from(&turn_context), + cancellation_token: None, call_id: approval_call_id, tool_name: ToolName::plain(telemetry_tool_name.clone()), strict_auto_review, diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index ad3ab60a9b..4685e0b536 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -77,6 +77,7 @@ impl ToolOrchestrator { let attempt_tool_ctx = ToolCtx { session: tool_ctx.session.clone(), step_context: Arc::clone(&tool_ctx.step_context), + cancellation_token: tool_ctx.cancellation_token.clone(), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), }; @@ -178,6 +179,7 @@ impl ToolOrchestrator { })?; let approval_ctx = ApprovalContext { review_context: GuardianReviewContext::from(&tool_ctx.step_context), + cancellation_token: Some(tool_ctx.cancellation_token.clone()), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), strict_auto_review, @@ -210,6 +212,7 @@ impl ToolOrchestrator { })?; let approval_ctx = ApprovalContext { review_context: GuardianReviewContext::from(&tool_ctx.step_context), + cancellation_token: Some(tool_ctx.cancellation_token.clone()), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), strict_auto_review, @@ -406,6 +409,7 @@ impl ToolOrchestrator { })?; let approval_ctx = ApprovalContext { review_context: GuardianReviewContext::from(&tool_ctx.step_context), + cancellation_token: Some(tool_ctx.cancellation_token.clone()), call_id: tool_ctx.call_id.clone(), tool_name: tool_ctx.tool_name.clone(), strict_auto_review, 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 41cf4020a4..5eb25622a4 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -481,6 +481,8 @@ impl CoreShellActionProvider { })?; let approval_ctx = ApprovalContext { review_context: GuardianReviewContext::from(turn_context), + // The running process can outlive its launching tool or code-mode cell. + cancellation_token: None, call_id: self.call_id.clone(), tool_name: self.tool_name.clone(), strict_auto_review, diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 6d5d89a6cc..e8a0def641 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -349,6 +349,7 @@ pub(crate) trait Sandboxable { pub(crate) struct ToolCtx { pub session: Arc, pub step_context: Arc, + pub cancellation_token: CancellationToken, pub call_id: String, pub tool_name: ToolName, } diff --git a/codex-rs/core/src/unified_exec/async_watcher_tests.rs b/codex-rs/core/src/unified_exec/async_watcher_tests.rs index 7f7ccee1ff..98456bbd72 100644 --- a/codex-rs/core/src/unified_exec/async_watcher_tests.rs +++ b/codex-rs/core/src/unified_exec/async_watcher_tests.rs @@ -52,6 +52,7 @@ async fn streaming_output_harness() -> anyhow::Result { let context = UnifiedExecContext::new( session, crate::session::step_context::StepContext::for_test(turn), + tokio_util::sync::CancellationToken::new(), "streaming-output-test".to_string(), ); let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 77b23625c6..d17671c7ad 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -35,6 +35,7 @@ use codex_utils_path_uri::PathUri; use rand::Rng; use rand::rng; use tokio::sync::Mutex; +use tokio_util::sync::CancellationToken; use crate::sandboxing::SandboxPermissions; use crate::session::session::Session; @@ -77,14 +78,21 @@ pub(crate) const MAX_UNIFIED_EXEC_PROCESSES: usize = 64; pub(crate) struct UnifiedExecContext { pub session: Arc, pub step_context: Arc, + pub cancellation_token: CancellationToken, pub call_id: String, } impl UnifiedExecContext { - pub fn new(session: Arc, step_context: Arc, call_id: String) -> Self { + pub fn new( + session: Arc, + step_context: Arc, + cancellation_token: CancellationToken, + call_id: String, + ) -> Self { Self { session, step_context, + cancellation_token, call_id, } } diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index f0362ebd80..3ba46c2b03 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -126,6 +126,7 @@ async fn exec_command_with_tty( let context = UnifiedExecContext::new( Arc::clone(session), crate::session::step_context::StepContext::for_test(Arc::clone(turn)), + tokio_util::sync::CancellationToken::new(), "call".to_string(), ); let started_at = Instant::now(); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 8edc5a5675..a863277fb8 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -1290,6 +1290,7 @@ impl UnifiedExecProcessManager { let tool_ctx = ToolCtx { session: context.session.clone(), step_context: Arc::clone(&context.step_context), + cancellation_token: context.cancellation_token.clone(), call_id: context.call_id.clone(), tool_name: ToolName::plain("exec_command"), }; 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 5c096ce7e9..56baa81409 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -392,6 +392,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { let context = UnifiedExecContext::new( Arc::clone(&session), crate::session::step_context::StepContext::for_test(Arc::clone(&turn)), + tokio_util::sync::CancellationToken::new(), "call-unified-denied".to_string(), ); let request = ExecCommandRequest { diff --git a/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs b/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs new file mode 100644 index 0000000000..77c631e2d5 --- /dev/null +++ b/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs @@ -0,0 +1,174 @@ +use anyhow::Context; +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_core::config::Constrained; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::GuardianAssessmentStatus; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::local_selections; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_mcp_server; +use serde_json::Value; +use serde_json::json; +use std::time::Duration; + +// The tool itself needs no approval. Its server requests a separate Guardian +// review after tools/call, exercising the ordinary MCP elicitation path. +const ELICITATION_SERVER: &str = r#" +import json +import sys + +def send(message): + print(json.dumps({"jsonrpc": "2.0", **message}), flush=True) + +pending_call = None +for line in sys.stdin: + request = json.loads(line) + method = request.get("method") + if method == "initialize": + result = {"protocolVersion": request["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "guardian-elicitation-test", "version": "1"}} + elif method == "tools/list": + result = {"tools": [{"name": "request_approval", + "inputSchema": {"type": "object", "properties": {}}, + "annotations": {"readOnlyHint": True}}]} + elif method == "tools/call": + pending_call = request["id"] + send({"id": "server-approval", "method": "elicitation/create", "params": { + "message": "Approve the server-side action?", + "requestedSchema": {"type": "object", "properties": {}}, + "_meta": {"codex_request_type": "approval_request", + "codex_approval_kind": "mcp_tool_call", "tool_name": "write_record"}}}) + continue + elif method is None and request.get("id") == "server-approval": + send({"id": pending_call, "result": {"content": [ + {"type": "text", "text": json.dumps(request.get("result"))}]}}) + continue + elif method == "resources/list": + result = {"resources": []} + elif method == "resources/templates/list": + result = {"resourceTemplates": []} + elif "id" not in request: + continue + else: + result = {} + send({"id": request["id"], "result": result}) +"#; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn interrupt_aborts_server_initiated_mcp_guardian_review() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!(Ok(()), "the MCP fixture requires a host Python interpreter"); + + let server = responses::start_mock_server().await; + let mcp_servers = serde_json::from_value(json!({ + "elicitation": { + "command": if cfg!(windows) { "python" } else { "python3" }, + "args": ["-u", "-c", ELICITATION_SERVER], + "default_tools_approval_mode": "approve", + } + }))?; + let mut builder = test_codex().with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config + .mcp_servers + .set(mcp_servers) + .expect("set MCP fixture"); + }); + let test = builder.build_with_auto_env(&server).await?; + wait_for_mcp_server(&test.codex, "elicitation").await?; + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "eliciting-tool", + "mcp__elicitation", + "request_approval", + "{}", + ), + responses::ev_completed("parent-tool"), + ]), + ) + .await; + let pending_guardian = responses::mount_response_once_match( + &server, + |request: &wiremock::Request| { + serde_json::from_slice::(&request.body) + .expect("Responses request body should be valid JSON") + .pointer("/client_metadata/x-openai-subagent") + .and_then(Value::as_str) + == Some("guardian") + }, + responses::sse_response(responses::sse(vec![ + responses::ev_assistant_message( + "review-result", + &json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "This response must be interrupted.", + }) + .to_string(), + ), + responses::ev_completed("pending-review"), + ])) + .set_delay(Duration::from_secs(60)), + ) + .await; + + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Run the tool that requests server-side approval.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + ..Default::default() + }), + ) + .await?; + tokio::time::timeout(Duration::from_secs(10), async { + while pending_guardian.requests().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .context("server-initiated Guardian review did not start")?; + assert!( + pending_guardian + .single_request() + .body_contains_text("write_record") + ); + + test.codex.submit(Op::Interrupt).await?; + tokio::time::timeout(Duration::from_secs(10), async { + let mut guardian_aborted = false; + let mut parent_aborted = false; + while !guardian_aborted || !parent_aborted { + match test.codex.next_event().await?.msg { + EventMsg::GuardianAssessment(assessment) + if assessment.status == GuardianAssessmentStatus::Aborted => + { + guardian_aborted = true; + } + EventMsg::TurnAborted(_) => parent_aborted = true, + _ => {} + } + } + anyhow::Ok(()) + }) + .await + .context("turn interruption did not abort the MCP elicitation review")??; + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/guardian_review_cancellation.rs b/codex-rs/core/tests/suite/guardian_review_cancellation.rs new file mode 100644 index 0000000000..323a627e36 --- /dev/null +++ b/codex-rs/core/tests/suite/guardian_review_cancellation.rs @@ -0,0 +1,251 @@ +use anyhow::Context; +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_core::config::Constrained; +use codex_features::Feature; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::GuardianAssessmentStatus; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_custom_tool_call; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_sandbox; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::local_selections; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use serde_json::Value; +use serde_json::json; +use std::time::Duration; +use test_case::test_case; + +#[derive(Clone, Copy)] +enum Cancellation { + DirectTool, + CodeModeTurn, + CodeModeCell, +} + +#[test_case(Cancellation::DirectTool; "direct tool")] +#[test_case(Cancellation::CodeModeTurn; "code mode turn")] +#[test_case(Cancellation::CodeModeCell; "code mode cell")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelling_tool_aborts_its_guardian_review(cancellation: Cancellation) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_wine_exec!( + Ok(()), + "Guardian approval actions require host-native paths" + ); + + let server = responses::start_mock_server().await; + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config + .set_legacy_sandbox_policy(SandboxPolicy::new_workspace_write_policy()) + .expect("set sandbox policy"); + if !matches!(cancellation, Cancellation::DirectTool) { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::CodeModeInterrupt); + } + }); + if !matches!(cancellation, Cancellation::DirectTool) { + builder = builder + .with_code_mode_host_program(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?); + } + let test = builder.build_with_auto_env(&server).await?; + let output_file = test.cwd.path().join("cancelled-guardian-command.txt"); + let output_path = shlex::try_join([output_file.to_string_lossy().as_ref()])?; + let tool_args = json!({ + "cmd": format!("printf should-not-run > {output_path}"), + "sandbox_permissions": "require_escalated", + "justification": "Exercise tool-owned Guardian cancellation.", + }); + let tool_call = match cancellation { + Cancellation::DirectTool => { + ev_function_call("reviewed-tool", "exec_command", &tool_args.to_string()) + } + Cancellation::CodeModeTurn => ev_custom_tool_call( + "reviewed-tool", + "exec", + &format!("await tools.exec_command({tool_args});"), + ), + Cancellation::CodeModeCell => ev_custom_tool_call( + "reviewed-tool", + "exec", + &format!("yield_control(); await tools.exec_command({tool_args});"), + ), + }; + responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("parent-start"), + tool_call, + ev_completed("parent-start"), + ]), + ) + .await; + let yielded_parent = if matches!(cancellation, Cancellation::CodeModeCell) { + Some( + responses::mount_response_once_match( + &server, + |request: &wiremock::Request| { + let body: Value = serde_json::from_slice(&request.body) + .expect("Responses request body should be valid JSON"); + body["input"].as_array().is_some_and(|items| { + items.iter().any(|item| { + item["type"] == "custom_tool_call_output" + && item["call_id"] == "reviewed-tool" + }) + }) && body + .pointer("/client_metadata/x-openai-subagent") + .and_then(Value::as_str) + != Some("guardian") + }, + responses::sse_response(sse(vec![ + ev_assistant_message("parent-yielded", "cell is running"), + ev_completed("parent-yielded"), + ])), + ) + .await, + ) + } else { + None + }; + let pending_guardian = responses::mount_response_once_match( + &server, + |request: &wiremock::Request| { + serde_json::from_slice::(&request.body) + .expect("Responses request body should be valid JSON") + .pointer("/client_metadata/x-openai-subagent") + .and_then(Value::as_str) + == Some("guardian") + }, + responses::sse_response(sse(vec![ + ev_response_created("pending-review"), + ev_assistant_message( + "review-result", + &json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "The test must cancel this review before it completes.", + }) + .to_string(), + ), + ev_completed("pending-review"), + ])) + .set_delay(Duration::from_secs(60)), + ) + .await; + + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "start a Guardian-reviewed command".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + environments: Some(local_selections(test.config.cwd.clone())), + ..Default::default() + }), + ) + .await?; + tokio::time::timeout(Duration::from_secs(10), async { + while pending_guardian.requests().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .context("Guardian review did not start")?; + + if let Some(yielded_parent) = yielded_parent { + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let output = yielded_parent + .single_request() + .custom_tool_call_output("reviewed-tool"); + let output = &output["output"]; + let header = output + .as_str() + .or_else(|| output[0]["text"].as_str()) + .context("missing code-mode output")?; + let cell_id = header + .strip_prefix("Script running with cell ID ") + .and_then(|rest| rest.lines().next()) + .context("missing cell id")?; + responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_function_call( + "terminate-cell", + "wait", + &json!({ + "cell_id": cell_id, + "terminate": true, + }) + .to_string(), + ), + ev_completed("terminate-cell"), + ]), + sse(vec![ + ev_assistant_message("parent-done", "cell terminated"), + ev_completed("parent-done"), + ]), + ], + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "terminate the background cell".into(), + text_elements: Vec::new(), + }])) + .await?; + } else { + test.codex.submit(Op::Interrupt).await?; + } + + tokio::time::timeout(Duration::from_secs(10), async { + let mut guardian_aborted = false; + let mut parent_finished = false; + while !guardian_aborted || !parent_finished { + let event = test.codex.next_event().await?; + match event.msg { + EventMsg::GuardianAssessment(assessment) + if assessment.status == GuardianAssessmentStatus::Aborted => + { + guardian_aborted = true; + } + EventMsg::TurnComplete(_) if matches!(cancellation, Cancellation::CodeModeCell) => { + parent_finished = true; + } + EventMsg::TurnAborted(_) if !matches!(cancellation, Cancellation::CodeModeCell) => { + parent_finished = true; + } + _ => {} + } + } + anyhow::Ok(()) + }) + .await + .context("tool cancellation did not abort Guardian")??; + assert!(!output_file.exists(), "cancelled command executed"); + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 98032d2087..747680edff 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -67,9 +67,12 @@ mod extension_sandbox; mod external_auth; mod fork_thread; mod git_enrichment; +mod guardian_mcp_elicitation; #[cfg(not(target_os = "windows"))] mod guardian_review; #[cfg(not(target_os = "windows"))] +mod guardian_review_cancellation; +#[cfg(not(target_os = "windows"))] mod hooks; #[cfg(not(target_os = "windows"))] mod hooks_mcp;