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