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
This commit is contained in:
jif
2026-10-05 10:08:59 +00:00
committed by copyberry
parent 315f0efb34
commit 39e013c8b7
7 changed files with 112 additions and 53 deletions
+1 -1
View File
@@ -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.
+21 -8
View File
@@ -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<HashMap<String, Vec<String>>>,
}
#[derive(Clone)]
pub(crate) struct McpToolApprovalContext {
pub(crate) invocation: Option<McpInvocation>,
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<McpInvocation>, McpToolApprovalMetadata)> {
review_context: GuardianReviewContext,
) -> Arc<McpToolApprovalContext> {
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<McpInvocation>, McpToolApprovalMetadata)> {
) -> Option<McpToolApprovalContext> {
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(),
+21 -8
View File
@@ -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::<Vec<_>>();
keys.sort();
+22 -15
View File
@@ -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 {
+2 -3
View File
@@ -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<McpInvocation>, McpToolApprovalMetadata)>>;
HashMap<(String, String), std::sync::Weak<McpToolApprovalContext>>;
/// Context for an initialized model agent
///
@@ -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
+44 -17
View File
@@ -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(())