From 39e013c8b7cb6eec46d16ff2b74fac4e22992e8c Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 5 Oct 2026 10:00:14 +0000 Subject: [PATCH] Use issuing-step context for Guardian MCP elicitation reviews (#51067) ## Why MCP elicitation reviews used turn-start context, which could omit remote environments that became ready before the tool invocation and use stale environment permissions. ## What changed Capture the issuing step's Guardian review context with each live MCP invocation and reuse it for elicitation reviews. Fall back to turn context when no matching live invocation is available, and preserve the captured context for in-flight actions. ## Testing Extend the deferred-executor regression test to cover MCP elicitations, checking that Guardian receives the newly ready remote environment and its denied paths, excludes stale local denied paths, and returns the denial to the parent model. GitOrigin-RevId: 0cc58b8f30543a9af87d53e1a90f574a8e6a2f84 --- codex-rs/core/src/guardian/mod.rs | 2 +- codex-rs/core/src/mcp_tool_call.rs | 29 ++++++--- codex-rs/core/src/mcp_tool_call_tests.rs | 29 ++++++--- codex-rs/core/src/session/mcp.rs | 37 ++++++----- codex-rs/core/src/session/session.rs | 5 +- .../tests/suite/guardian_mcp_elicitation.rs | 2 +- codex-rs/core/tests/suite/remote_env.rs | 61 +++++++++++++------ 7 files changed, 112 insertions(+), 53 deletions(-) diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 8314cd88f7..138e3ef928 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -78,7 +78,7 @@ pub(crate) const GUARDIAN_MAX_NODE_REPL_TOOL_RESULT_TOKENS: usize = 6_000; /// Background network approvals and Unix interception use the active task's resolved settings. /// Startup reviewer prewarming intentionally uses turn-only inputs because it has no issuing step. /// -/// MCP elicitation reviews continue to use turn-only inputs. +/// MCP elicitations use their live invocation's issuing step when available, otherwise turn inputs. #[derive(Clone)] pub(crate) struct GuardianReviewContext { /// The latest response ID received in this turn when review was requested. diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 1fe0016e6e..af9d3d05b6 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -255,8 +255,12 @@ pub(crate) async fn handle_mcp_tool_call( .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), }; } - let _approval_metadata = - sess.register_mcp_tool_approval_metadata(&call_id, &invocation, metadata.clone()); + let _approval_metadata = sess.register_mcp_tool_approval_metadata( + &call_id, + &invocation, + metadata.clone(), + GuardianReviewContext::from(step_context), + ); notify_mcp_tool_call_started( sess.as_ref(), turn_context.as_ref(), @@ -1189,20 +1193,29 @@ pub(crate) struct McpToolApprovalMetadata { openai_file_input_optional_fields: Option>>, } +#[derive(Clone)] +pub(crate) struct McpToolApprovalContext { + pub(crate) invocation: Option, + pub(crate) metadata: McpToolApprovalMetadata, + pub(crate) review_context: GuardianReviewContext, +} + impl Session { fn register_mcp_tool_approval_metadata( &self, call_id: &str, invocation: &McpInvocation, metadata: McpToolApprovalMetadata, - ) -> Arc<(Option, McpToolApprovalMetadata)> { + review_context: GuardianReviewContext, + ) -> Arc { let key = (invocation.server.clone(), call_id.to_string()); - let metadata = Arc::new(( - (invocation.server == CODEX_APPS_MCP_SERVER_NAME + let metadata = Arc::new(McpToolApprovalContext { + invocation: (invocation.server == CODEX_APPS_MCP_SERVER_NAME || is_node_repl_backed_server(&invocation.server)) .then(|| invocation.clone()), metadata, - )); + review_context, + }); let mut registry = self .mcp_tool_approval_metadata .lock() @@ -1218,7 +1231,7 @@ impl Session { &self, server: &str, call_id: &str, - ) -> Option<(Option, McpToolApprovalMetadata)> { + ) -> Option { self.mcp_tool_approval_metadata .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -1627,7 +1640,7 @@ pub(crate) async fn request_mcp_tool_user_approval( let (request_dispatched, decision) = if tool_call_mcp_elicitation_enabled { let link_id = sess .mcp_tool_approval_metadata(server, id) - .and_then(|(_, metadata)| metadata.link_id); + .and_then(|context| context.metadata.link_id); let metadata = McpToolApprovalMetadata { annotations: None, connector_id: connector_id.clone(), diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 8cace5296b..5aec67a193 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -3454,7 +3454,8 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { #[tokio::test] async fn approval_metadata_is_released_when_the_invocation_future_is_dropped() { - let (session, _) = crate::session::tests::make_session_and_context().await; + let (session, turn_context) = crate::session::tests::make_session_and_context().await; + let review_context = GuardianReviewContext::from(Arc::new(turn_context)); let invocation = McpInvocation { server: CODEX_APPS_MCP_SERVER_NAME.to_string(), tool: "write_record".to_string(), @@ -3468,17 +3469,25 @@ async fn approval_metadata_is_released_when_the_invocation_future_is_dropped() { /*tool_description*/ None, ); let mut call = Box::pin(async { - let _approval_metadata = - session.register_mcp_tool_approval_metadata("call", &invocation, metadata.clone()); + let _approval_metadata = session.register_mcp_tool_approval_metadata( + "call", + &invocation, + metadata.clone(), + review_context.clone(), + ); std::future::pending::<()>().await; }); assert!(futures::poll!(call.as_mut()).is_pending()); - let _other_metadata = - session.register_mcp_tool_approval_metadata("other-call", &invocation, metadata.clone()); + let _other_metadata = session.register_mcp_tool_approval_metadata( + "other-call", + &invocation, + metadata.clone(), + review_context.clone(), + ); assert_eq!( session .mcp_tool_approval_metadata(CODEX_APPS_MCP_SERVER_NAME, "call") - .map(|(invocation, metadata)| (invocation, metadata.connector_id)), + .map(|context| (context.invocation, context.metadata.connector_id)), Some((Some(invocation.clone()), Some("connector".to_string()))), ); assert!( @@ -3492,8 +3501,12 @@ async fn approval_metadata_is_released_when_the_invocation_future_is_dropped() { .mcp_tool_approval_metadata(CODEX_APPS_MCP_SERVER_NAME, "call") .is_none() ); - let _next_metadata = - session.register_mcp_tool_approval_metadata("next-call", &invocation, metadata); + let _next_metadata = session.register_mcp_tool_approval_metadata( + "next-call", + &invocation, + metadata, + review_context, + ); let registry = session.mcp_tool_approval_metadata.lock().unwrap(); let mut keys = registry.keys().cloned().collect::>(); keys.sort(); diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index e7c95576db..c62ff6149d 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -2,6 +2,7 @@ use super::mcp_refresh::McpRefreshInvalidationGuard; use super::*; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::environment_selection::combine_selected_capability_roots; +use crate::mcp_tool_call::McpToolApprovalContext; use codex_exec_server::ExecutorCapabilityDiscoveryCache; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; use codex_exec_server::MAX_SELECTED_CAPABILITY_ROOTS; @@ -824,9 +825,14 @@ async fn review_guardian_mcp_elicitation( _ => meta.get("callId"), }) .and_then(Value::as_str); + let approval_context = call_id + .and_then(|call_id| session.mcp_tool_approval_metadata(&request.server_name, call_id)); let (originating_call_id, guardian_scope) = if let Some(call_id) = call_id - && let Some((Some(invocation), metadata)) = - session.mcp_tool_approval_metadata(&request.server_name, call_id) + && let Some(McpToolApprovalContext { + invocation: Some(invocation), + metadata, + .. + }) = approval_context.as_ref() && invocation.server == request.server_name && is_node_repl_backed_connector(&invocation.server, metadata.connector_id.as_deref()) { @@ -843,7 +849,12 @@ async fn review_guardian_mcp_elicitation( Some(Value::Bool(true)) ); - let mut review_context = crate::guardian::GuardianReviewContext::from(&turn_context); + // Bind the review to the issuing step, which may include environments that + // became ready after the turn started. Do not refresh an in-flight action. + let mut review_context = approval_context + .as_ref() + .map(|context| context.review_context.clone()) + .unwrap_or_else(|| crate::guardian::GuardianReviewContext::from(&turn_context)); let strict_auto_review = matches!( request .elicitation @@ -858,18 +869,14 @@ async fn review_guardian_mcp_elicitation( let review_outer_invocation = request.server_name == CODEX_APPS_MCP_SERVER_NAME && originating_call_id.is_none(); let trusted_guardian_request = if review_outer_invocation { - let Some(call_id) = request - .elicitation - .meta() - .and_then(|meta| meta.get(MCP_TOOL_CODEX_APPS_META_KEY)) - .and_then(Value::as_object) - .and_then(|meta| meta.get("call_id")) - .and_then(Value::as_str) - else { + let Some(call_id) = call_id else { return Ok(None); }; - let Some((Some(invocation), metadata)) = - session.mcp_tool_approval_metadata(&request.server_name, call_id) + let Some(McpToolApprovalContext { + invocation: Some(invocation), + metadata, + .. + }) = approval_context.as_ref() else { return Ok(None); }; @@ -886,8 +893,8 @@ async fn review_guardian_mcp_elicitation( Some( crate::mcp_tool_call::build_guardian_mcp_tool_review_request( call_id, - &invocation, - Some(&metadata), + invocation, + Some(metadata), ), ) } else { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2319e3f7ae..8d16c6b9fc 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -14,7 +14,7 @@ use crate::config::ConstraintError; use crate::environment_selection::ThreadEnvironments; use crate::environment_selection::TurnEnvironmentSnapshot; use crate::hook_mcp_executor::CoreHookMcpExecutor; -use crate::mcp_tool_call::McpToolApprovalMetadata; +use crate::mcp_tool_call::McpToolApprovalContext; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_metadata::CompactionTurnMetadata; @@ -39,7 +39,6 @@ use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::protocol::EnvironmentConfig; use codex_protocol::protocol::HookCompletedEvent; -use codex_protocol::protocol::McpInvocation; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; @@ -52,7 +51,7 @@ use std::sync::OnceLock; use tokio::sync::Semaphore; type McpToolApprovalMetadataMap = - HashMap<(String, String), std::sync::Weak<(Option, McpToolApprovalMetadata)>>; + HashMap<(String, String), std::sync::Weak>; /// Context for an initialized model agent /// diff --git a/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs b/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs index 50fbd2091a..aa380b47eb 100644 --- a/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs +++ b/codex-rs/core/tests/suite/guardian_mcp_elicitation.rs @@ -37,7 +37,7 @@ use wiremock::matchers::body_partial_json; // 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#" +pub(super) const ELICITATION_SERVER: &str = r#" import json import sys diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index a60c2e5d6f..90048b9cdb 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -3154,12 +3154,22 @@ async fn deferred_executor_spawn_agent_inherits_ready_step_environments( Ok(()) } +#[test_case("exec_command")] +#[test_case("mcp_elicitation")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn deferred_executor_guardian_uses_newly_ready_step_environment() -> Result<()> { +async fn deferred_executor_guardian_uses_newly_ready_step_environment( + tool: &'static str, +) -> Result<()> { const WAIT_CALL_ID: &str = "wait-for-guardian-environment"; const EXEC_CALL_ID: &str = "guardian-ready-environment-command"; const DENIAL_RATIONALE: &str = "The remote environment policy denies this action."; + if tool == "mcp_elicitation" { + core_test_support::skip_if_wine_exec!( + Ok(()), + "the MCP fixture requires a host Python interpreter" + ); + } let listener = TcpListener::bind("127.0.0.1:0").await?; let server = start_mock_server().await; let completed_response = @@ -3177,17 +3187,21 @@ async fn deferred_executor_guardian_uses_newly_ready_step_environment() -> Resul ), completed_response( "resp-guardian-command", - ev_function_call( - EXEC_CALL_ID, - "exec_command", - &json!({ - "cmd": "printf guardian-should-not-run", - "environment_id": REMOTE_ENVIRONMENT_ID, - "sandbox_permissions": SandboxPermissions::RequireEscalated, - "justification": "Review the newly ready remote environment.", - }) - .to_string(), - ), + if tool == "mcp_elicitation" { + ev_function_call_with_namespace(EXEC_CALL_ID, "mcp__cua_repl", "js", "{}") + } else { + ev_function_call( + EXEC_CALL_ID, + "exec_command", + &json!({ + "cmd": "printf guardian-should-not-run", + "environment_id": REMOTE_ENVIRONMENT_ID, + "sandbox_permissions": SandboxPermissions::RequireEscalated, + "justification": "Review the newly ready remote environment.", + }) + .to_string(), + ) + }, ), completed_response( "resp-guardian-review", @@ -3205,11 +3219,21 @@ async fn deferred_executor_guardian_uses_newly_ready_step_environment() -> Resul .await; let mut builder = test_codex() .with_exec_server_url(format!("ws://{}", listener.local_addr()?)) - .with_config(|config| { + .with_config(move |config| { config.project_doc_max_bytes = 0; config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); config.approvals_reviewer = ApprovalsReviewer::AutoReview; assert!(config.features.enable(Feature::DeferredExecutor).is_ok()); + if tool == "mcp_elicitation" { + config.mcp_servers.set(serde_json::from_value(json!({ + "cua_repl": { + "command": if cfg!(windows) { "python" } else { "python3" }, + "args": ["-u", "-c", super::guardian_mcp_elicitation::ELICITATION_SERVER, + json!([{"codex_sensitive_action": true}]).to_string(), "forward"], + "default_tools_approval_mode": "approve", + } + })).expect("MCP server config")).expect("set MCP fixture"); + } }); let (attach_tx, attach_rx) = tokio::sync::oneshot::channel(); let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); @@ -3220,6 +3244,9 @@ async fn deferred_executor_guardian_uses_newly_ready_step_environment() -> Resul shutdown_rx, )); let test = expect_startup(builder.build_with_remote_and_local_env(&server)).await; + if tool == "mcp_elicitation" { + core_test_support::wait_for_mcp_server(&test.codex, "cua_repl").await?; + } let remote_cwd = test.cwd.path().join("guardian-remote").abs(); let local_cwd = test.cwd.path().abs(); fs::create_dir_all(remote_cwd.as_path())?; @@ -3305,10 +3332,10 @@ async fn deferred_executor_guardian_uses_newly_ready_step_environment() -> Resul "Guardian used the stale local environment's denied-read policy: {guardian_context}" ); let rejection = requests - .iter() - .find_map(|request| request.function_call_output_text(EXEC_CALL_ID)) - .context("Guardian denial should be returned to the parent model")?; - assert!(rejection.contains(DENIAL_RATIONALE)); + .last() + .context("expected parent continuation after review")? + .function_call_output(EXEC_CALL_ID); + assert!(rejection.to_string().contains(DENIAL_RATIONALE)); shutdown_tx .send(())