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
This commit is contained in:
Krish Chainani
2026-08-21 21:17:42 +00:00
committed by copyberry
parent 8b5beea5c0
commit ab8768306f
14 changed files with 157 additions and 6 deletions
+25 -2
View File
@@ -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<T: 'static>(
cwd: &Path,
turn_id: Option<String>,
parse: fn(&ConfiguredHandler, HandlerRunResult, Option<String>) -> ParsedHandler<T>,
) -> Vec<ParsedHandler<T>> {
execute_handlers_with_metadata(
engine, handlers, input_json, cwd, turn_id, /*metadata*/ None, parse,
)
.await
}
pub(crate) async fn execute_handlers_with_metadata<T: 'static>(
engine: &ClaudeHooksEngine,
handlers: Vec<ConfiguredHandler>,
input_json: String,
cwd: &Path,
turn_id: Option<String>,
metadata: Option<&Map<String, Value>>,
parse: fn(&ConfiguredHandler, HandlerRunResult, Option<String>) -> ParsedHandler<T>,
) -> Vec<ParsedHandler<T>> {
let mut executor_handlers = Vec::new();
let mut pending = FuturesUnordered::new();
@@ -124,7 +141,8 @@ pub(crate) async fn execute_handlers<T: 'static>(
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<T: 'static>(
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<String, Value>>,
) -> HandlerRunResult {
match &handler.kind {
ConfiguredHandlerKind::Command { command, env, .. } => {
@@ -193,6 +215,7 @@ async fn execute_handler(
tool,
input,
input_json,
metadata,
)
.await
}
+9
View File
@@ -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<String, Value>,
hook_event_json: &str,
metadata: Option<&Map<String, Value>>,
) -> 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),
})
@@ -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),
+17
View File
@@ -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",
}))
+5 -1
View File
@@ -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<PathBuf>,
pub model: String,
pub permission_mode: String,
pub request_metadata: Option<Map<String, Value>>,
pub stop_hook_active: bool,
pub last_assistant_message: Option<String>,
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;
+2
View File
@@ -9,6 +9,8 @@ use serde_json::Value;
pub struct HookMcpCall {
pub server: String,
pub tool: String,
pub environment_id: Option<String>,
pub metadata: Option<Map<String, Value>>,
pub input: Map<String, Value>,
pub timeout: Duration,
}