From ab8768306ffdc06f7abf83e76070aaa17dd06bcc Mon Sep 17 00:00:00 2001 From: Krish Chainani Date: Fri, 21 Aug 2026 21:09:22 +0000 Subject: [PATCH] Preserve executor context for MCP stop hooks (#40012) ## What changed - Scope executor-provided stop-hook calls to the MCP server environment that registered the hook, and reject calls when that environment does not match. - Forward turn metadata with executor stop-hook requests while retaining the hook call's `threadId` metadata. ## Testing - Cover environment matching for MCP tool calls and verify executor stop hooks receive their environment ID and request metadata. GitOrigin-RevId: b544ea3a6ca1f61389fb4597978ed7aac27d529a --- codex-rs/codex-mcp/src/connection_manager.rs | 10 ++++ .../codex-mcp/src/connection_manager_tests.rs | 57 +++++++++++++++++++ codex-rs/codex-mcp/src/runtime.rs | 3 + codex-rs/core/src/codex_thread.rs | 4 +- codex-rs/core/src/hook_mcp_executor.rs | 9 ++- codex-rs/core/src/hook_runtime.rs | 14 +++++ codex-rs/core/src/session/mod.rs | 1 + codex-rs/core/tests/suite/mcp_auth_refresh.rs | 1 + codex-rs/hooks/src/engine/dispatcher.rs | 27 ++++++++- codex-rs/hooks/src/engine/mcp_runner.rs | 9 +++ codex-rs/hooks/src/engine/mcp_runner_tests.rs | 3 + codex-rs/hooks/src/engine/mod_tests.rs | 17 ++++++ codex-rs/hooks/src/events/stop.rs | 6 +- codex-rs/hooks/src/mcp.rs | 2 + 14 files changed, 157 insertions(+), 6 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index bc4037384b..eced5326c7 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -854,10 +854,12 @@ impl McpConnectionSet { } /// Invoke the tool indicated by the (server, tool) pair. + #[allow(clippy::too_many_arguments)] pub async fn call_tool( &self, server: &str, tool: &str, + environment_id: Option<&str>, arguments: Option, meta: Option, requested_timeout: Option, @@ -867,6 +869,14 @@ impl McpConnectionSet { .servers .get(server) .ok_or_else(|| anyhow!("unknown MCP server '{server}'"))?; + if let Some(environment_id) = environment_id + && view.metadata.environment_id != environment_id + { + bail!( + "MCP server `{server}` is running in environment `{}`, expected `{environment_id}`", + view.metadata.environment_id + ); + } if !view.tool_filter.allows(tool) { return Err(anyhow!( "tool '{tool}' is disabled for MCP server '{server}'" diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index a33ddf3323..1a9ed07932 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2674,6 +2674,7 @@ async fn call_tool_requires_connection_without_waiting_for_startup() { manager.call_tool( "docs", "search", + /*environment_id*/ None, /*arguments*/ None, /*meta*/ None, Some(Duration::from_secs(5)), @@ -2697,6 +2698,7 @@ async fn call_tool_requires_connection_without_waiting_for_startup() { .call_tool( "docs", "search", + /*environment_id*/ None, /*arguments*/ None, /*meta*/ None, Some(Duration::from_secs(5)), @@ -2731,6 +2733,7 @@ async fn connected_call_respects_server_tool_filters() { .call_tool( "docs", "search", + /*environment_id*/ None, /*arguments*/ None, /*meta*/ None, Some(Duration::from_secs(5)), @@ -2741,6 +2744,60 @@ async fn connected_call_respects_server_tool_filters() { assert!(filtered_error.to_string().contains("disabled")); } +#[tokio::test] +async fn call_tool_validates_environment_without_waiting_for_ready_connections() { + let client = create_test_managed_client(vec![create_test_tool("docs", "search")]).await; + let (client, _, _) = create_gated_async_managed_client(client); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.insert_test_client("docs", client); + manager + .servers + .get_mut("docs") + .expect("test server should exist") + .metadata + .environment_id = "executor-a".to_string(); + + let mismatched_environment = manager + .call_tool( + "docs", + "search", + Some("executor-b"), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + /*wait_for_server*/ false, + ) + .await + .expect_err("calls must reject a server from a different environment"); + assert_eq!( + mismatched_environment.to_string(), + "MCP server `docs` is running in environment `executor-a`, expected `executor-b`" + ); + + let pending_call = tokio::time::timeout( + Duration::from_millis(50), + manager.call_tool( + "docs", + "search", + Some("executor-a"), + /*arguments*/ None, + /*meta*/ None, + Some(Duration::from_secs(5)), + /*wait_for_server*/ false, + ), + ) + .await + .expect("environment-scoped calls must not wait for pending server startup") + .expect_err("pending server must not accept environment-scoped calls"); + assert!(pending_call.to_string().contains("not connected")); +} + #[tokio::test] async fn list_all_tools_resolves_server_catalogs_concurrently() { let first_client = create_test_managed_client(vec![create_test_tool("first", "search")]).await; diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 56bc38df7f..207022c377 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -471,10 +471,12 @@ impl McpRuntime { self.latest_connections().list_all_tools().await } + #[allow(clippy::too_many_arguments)] pub async fn latest_call_tool( &self, server: &str, tool: &str, + environment_id: Option<&str>, arguments: Option, meta: Option, requested_timeout: Option, @@ -484,6 +486,7 @@ impl McpRuntime { .call_tool( server, tool, + environment_id, arguments, meta, requested_timeout, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index dd77ffe382..5c7fc6a453 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -833,8 +833,8 @@ impl CodexThread { .services .mcp_runtime .latest_call_tool( - server, tool, arguments, meta, /*requested_timeout*/ None, - /*wait_for_server*/ true, + server, tool, /*environment_id*/ None, arguments, meta, + /*requested_timeout*/ None, /*wait_for_server*/ true, ) .await } diff --git a/codex-rs/core/src/hook_mcp_executor.rs b/codex-rs/core/src/hook_mcp_executor.rs index ea1252f931..7fa0b78d42 100644 --- a/codex-rs/core/src/hook_mcp_executor.rs +++ b/codex-rs/core/src/hook_mcp_executor.rs @@ -18,13 +18,20 @@ pub(crate) struct CoreHookMcpExecutor { impl HookMcpExecutor for CoreHookMcpExecutor { fn execute(&self, call: HookMcpCall) -> BoxFuture<'_, anyhow::Result> { async move { + let mut metadata = call.metadata.unwrap_or_default(); + metadata.insert( + "threadId".to_string(), + Value::String(self.thread_id.to_string()), + ); + let result = self .runtime .latest_call_tool( &call.server, &call.tool, + call.environment_id.as_deref(), Some(Value::Object(call.input)), - Some(serde_json::json!({ "threadId": self.thread_id.to_string() })), + Some(Value::Object(metadata)), Some(call.timeout), /*wait_for_server*/ false, ) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index c4b9354981..ba8c1227e0 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -57,6 +57,7 @@ use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::PermissionRequestPayload; +use crate::turn_metadata::McpTurnMetadataContext; pub(crate) struct HookRuntimeOutcome { pub should_stop: bool, @@ -356,6 +357,18 @@ pub(crate) async fn run_turn_stop_hooks( SessionSource::SubAgent(_) => return StopOutcome::default(), _ => (StopHookTarget::Stop, sess.hook_transcript_path().await), }; + let request_metadata = turn_context + .turn_metadata_state + .current_meta_value_for_mcp_request(McpTurnMetadataContext { + model: step_context.model_info.slug.as_str(), + reasoning_effort: step_context.reasoning_effort.clone(), + }) + .map(|turn_metadata| { + serde_json::Map::from_iter([( + crate::X_CODEX_TURN_METADATA_HEADER.to_string(), + turn_metadata, + )]) + }); let request = codex_hooks::StopRequest { session_id: sess.session_id().into(), turn_id: turn_context.sub_id.clone(), @@ -364,6 +377,7 @@ pub(crate) async fn run_turn_stop_hooks( transcript_path, model: turn_context.model_info.slug.clone(), permission_mode: hook_permission_mode(turn_context), + request_metadata, stop_hook_active, last_assistant_message, target, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d76155848a..c1bc651b3f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3635,6 +3635,7 @@ impl Session { .latest_call_tool( "notes", "thread_hint", + /*environment_id*/ None, /*arguments*/ None, Some(serde_json::json!({ "threadId": self.thread_id().to_string(), diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index 7b39ff4f56..1552312644 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -133,6 +133,7 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( .latest_call_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", + /*environment_id*/ None, Some(json!({ "title": "Lunch", "starts_at": "2026-06-18T12:00:00Z", diff --git a/codex-rs/hooks/src/engine/dispatcher.rs b/codex-rs/hooks/src/engine/dispatcher.rs index cdcc0cae63..575ee46b43 100644 --- a/codex-rs/hooks/src/engine/dispatcher.rs +++ b/codex-rs/hooks/src/engine/dispatcher.rs @@ -2,6 +2,8 @@ use std::path::Path; use futures::StreamExt; use futures::stream::FuturesUnordered; +use serde_json::Map; +use serde_json::Value; use codex_protocol::protocol::HookCompletedEvent; use codex_protocol::protocol::HookEventName; @@ -100,6 +102,21 @@ pub(crate) async fn execute_handlers( cwd: &Path, turn_id: Option, parse: fn(&ConfiguredHandler, HandlerRunResult, Option) -> ParsedHandler, +) -> Vec> { + execute_handlers_with_metadata( + engine, handlers, input_json, cwd, turn_id, /*metadata*/ None, parse, + ) + .await +} + +pub(crate) async fn execute_handlers_with_metadata( + engine: &ClaudeHooksEngine, + handlers: Vec, + input_json: String, + cwd: &Path, + turn_id: Option, + metadata: Option<&Map>, + parse: fn(&ConfiguredHandler, HandlerRunResult, Option) -> ParsedHandler, ) -> Vec> { let mut executor_handlers = Vec::new(); let mut pending = FuturesUnordered::new(); @@ -124,7 +141,8 @@ pub(crate) async fn execute_handlers( let input_json = input_json.clone(); let turn_id = turn_id.clone(); pending.push(async move { - let result = execute_handler(engine, &handler, &input_json, cwd).await; + let result = + execute_handler(engine, &handler, &input_json, cwd, /*metadata*/ None).await; (configured_order, parse(&handler, result, turn_id)) }); } @@ -147,8 +165,11 @@ pub(crate) async fn execute_handlers( let task_engine = engine.clone(); let input_json = input_json.clone(); let cwd = cwd.to_path_buf(); + let metadata = metadata.cloned(); engine.command_runtime.schedule_async_task(async move { - let result = execute_handler(&task_engine, &handler, &input_json, &cwd).await; + let result = + execute_handler(&task_engine, &handler, &input_json, &cwd, metadata.as_ref()) + .await; if let Some(error) = result.error { tracing::warn!( source_path = %handler.source_path, @@ -168,6 +189,7 @@ async fn execute_handler( handler: &ConfiguredHandler, input_json: &str, cwd: &Path, + metadata: Option<&Map>, ) -> HandlerRunResult { match &handler.kind { ConfiguredHandlerKind::Command { command, env, .. } => { @@ -193,6 +215,7 @@ async fn execute_handler( tool, input, input_json, + metadata, ) .await } diff --git a/codex-rs/hooks/src/engine/mcp_runner.rs b/codex-rs/hooks/src/engine/mcp_runner.rs index ef33a1d871..63befa6bc6 100644 --- a/codex-rs/hooks/src/engine/mcp_runner.rs +++ b/codex-rs/hooks/src/engine/mcp_runner.rs @@ -9,6 +9,7 @@ use serde_json::Value; use super::ConfiguredHandler; use super::HandlerRunResult; +use super::HandlerSourcePath; use crate::mcp::HookMcpCall; use crate::mcp::HookMcpExecutor; @@ -30,6 +31,7 @@ pub(crate) async fn run_mcp_tool( tool: &str, argument_template: &Map, hook_event_json: &str, + metadata: Option<&Map>, ) -> HandlerRunResult { let started_at = chrono::Utc::now().timestamp(); let started = Instant::now(); @@ -41,6 +43,13 @@ pub(crate) async fn run_mcp_tool( .execute(HookMcpCall { server: server.to_string(), tool: tool.to_string(), + environment_id: match &handler.source_path { + HandlerSourcePath::Local(_) => None, + HandlerSourcePath::ExecutorScoped { environment_id, .. } => { + Some(environment_id.clone()) + } + }, + metadata: metadata.cloned(), input, timeout: Duration::from_secs(handler.timeout_sec), }) diff --git a/codex-rs/hooks/src/engine/mcp_runner_tests.rs b/codex-rs/hooks/src/engine/mcp_runner_tests.rs index 67c9353f4e..ceff78a1cf 100644 --- a/codex-rs/hooks/src/engine/mcp_runner_tests.rs +++ b/codex-rs/hooks/src/engine/mcp_runner_tests.rs @@ -118,6 +118,7 @@ async fn mcp_tool_results_use_command_hook_output_contract() { "scan", &configured_input, r#"{"tool_input":{"file_path":"/tmp/example.rs"}}"#, + /*metadata*/ None, ) .await; @@ -129,6 +130,8 @@ async fn mcp_tool_results_use_command_hook_output_contract() { vec![HookMcpCall { server: "security".to_string(), tool: "scan".to_string(), + environment_id: None, + metadata: None, input: serde_json::from_value(json!({ "file_path": "/tmp/example.rs" })) .expect("object input"), timeout: Duration::from_secs(30), diff --git a/codex-rs/hooks/src/engine/mod_tests.rs b/codex-rs/hooks/src/engine/mod_tests.rs index e3ed160811..2534fd25f7 100644 --- a/codex-rs/hooks/src/engine/mod_tests.rs +++ b/codex-rs/hooks/src/engine/mod_tests.rs @@ -2057,6 +2057,10 @@ fn executor_stop_hook_fixture() -> ( HookExecutionMode::Async ); assert!(!engine.handlers[0].can_apply_control_effects()); + let request_metadata = Some(serde_json::Map::from_iter([( + "x-codex-turn-metadata".to_string(), + serde_json::json!({ "turn_id": "turn-1" }), + )])); let request = StopRequest { session_id: ThreadId::new(), turn_id: "turn-1".to_string(), @@ -2064,6 +2068,7 @@ fn executor_stop_hook_fixture() -> ( transcript_path: None, model: "test-model".to_string(), permission_mode: "default".to_string(), + request_metadata: request_metadata.clone(), stop_hook_active: false, last_assistant_message: None, target: StopHookTarget::Stop, @@ -2072,6 +2077,8 @@ fn executor_stop_hook_fixture() -> ( let expected_executor_call = HookMcpCall { server: "node_repl".to_string(), tool: "turn_ended".to_string(), + environment_id: Some("executor-a".to_string()), + metadata: request_metadata, input: serde_json::from_value(serde_json::json!({ "turn_id": "turn-1" })) .expect("expanded executor hook input"), timeout: Duration::from_secs(5), @@ -2131,6 +2138,8 @@ async fn executor_stop_hooks_run_unless_regular_hooks_block_without_stopping() { HookMcpCall { server: "security".to_string(), tool: "check".to_string(), + environment_id: None, + metadata: None, input: Default::default(), timeout: Duration::from_secs(30), }, @@ -2165,18 +2174,24 @@ async fn executor_stop_hooks_run_unless_regular_hooks_block_without_stopping() { HookMcpCall { server: "security".to_string(), tool: "check".to_string(), + environment_id: None, + metadata: None, input: Default::default(), timeout: Duration::from_secs(30), }, HookMcpCall { server: "security".to_string(), tool: "check".to_string(), + environment_id: None, + metadata: None, input: Default::default(), timeout: Duration::from_secs(30), }, HookMcpCall { server: "security".to_string(), tool: "terminate".to_string(), + environment_id: None, + metadata: None, input: Default::default(), timeout: Duration::from_secs(30), }, @@ -2316,6 +2331,8 @@ async fn mcp_tool_hooks_expand_event_input_and_apply_pre_tool_decisions() { vec![HookMcpCall { server: "security".to_string(), tool: "scan".to_string(), + environment_id: None, + metadata: None, input: serde_json::from_value(serde_json::json!({ "command": "rm important.txt", })) diff --git a/codex-rs/hooks/src/events/stop.rs b/codex-rs/hooks/src/events/stop.rs index 6bfee00efa..b29e691911 100644 --- a/codex-rs/hooks/src/events/stop.rs +++ b/codex-rs/hooks/src/events/stop.rs @@ -9,6 +9,8 @@ use codex_protocol::protocol::HookOutputEntryKind; use codex_protocol::protocol::HookRunStatus; use codex_protocol::protocol::HookRunSummary; use codex_utils_absolute_path::AbsolutePathBuf; +use serde_json::Map; +use serde_json::Value; use super::common; use crate::engine::ClaudeHooksEngine; @@ -29,6 +31,7 @@ pub struct StopRequest { pub transcript_path: Option, pub model: String, pub permission_mode: String, + pub request_metadata: Option>, pub stop_hook_active: bool, pub last_assistant_message: Option, pub target: StopHookTarget, @@ -175,12 +178,13 @@ pub(crate) async fn run(engine: &ClaudeHooksEngine, request: StopRequest) -> Sto } }; - let results = dispatcher::execute_handlers( + let results = dispatcher::execute_handlers_with_metadata( engine, matched, input_json, request.cwd.as_path(), Some(request.turn_id), + request.request_metadata.as_ref(), parse_completed, ) .await; diff --git a/codex-rs/hooks/src/mcp.rs b/codex-rs/hooks/src/mcp.rs index 5999db2954..49d907df80 100644 --- a/codex-rs/hooks/src/mcp.rs +++ b/codex-rs/hooks/src/mcp.rs @@ -9,6 +9,8 @@ use serde_json::Value; pub struct HookMcpCall { pub server: String, pub tool: String, + pub environment_id: Option, + pub metadata: Option>, pub input: Map, pub timeout: Duration, }