use std::collections::HashMap;use std::fs;use std::path::Path;use std::time::Duration;use anyhow::Context;use anyhow::Result;use codex_config::types::AppToolApproval;use codex_config::types::McpServerConfig;use codex_config::types::McpServerTransportConfig;use codex_core::config::Config;use codex_features::Feature;use core_test_support::hooks::trust_discovered_hooks;use core_test_support::responses::ev_assistant_message;use core_test_support::responses::ev_completed;use core_test_support::responses::ev_function_call_with_namespace;use core_test_support::responses::ev_response_created;use core_test_support::responses::mount_sse_once;use core_test_support::responses::mount_sse_sequence;use core_test_support::responses::sse;use core_test_support::responses::start_mock_server;use core_test_support::skip_if_no_network;use core_test_support::stdio_server_bin;use core_test_support::test_codex::test_codex;use core_test_support::wait_for_mcp_server;use pretty_assertions::assert_eq;use serde_json::Value;use serde_json::json;const RMCP_SERVER: &str = "rmcp";const RMCP_PREFIXED_NAMESPACE: &str = "mcp__rmcp";const RMCP_UNPREFIXED_NAMESPACE: &str = "rmcp";const RMCP_ECHO_TOOL_NAME: &str = "mcp__rmcp__echo";const RMCP_HOOK_MATCHER: &str = RMCP_ECHO_TOOL_NAME;const RMCP_ECHO_MESSAGE: &str = "hook e2e ping";fn enable_mcp_tool_name_features(config: &mut Config, prefix_mcp_tool_names: bool) { if !prefix_mcp_tool_names { let _ = config.features.enable(Feature::NonPrefixedMcpToolNames); }}fn write_pre_tool_use_hook(home: &Path, reason: &str) -> Result<()> { let script_path = home.join("pre_tool_use_hook.py"); let log_path = home.join("pre_tool_use_hook_log.jsonl"); let reason_json = serde_json::to_string(reason).context("serialize pre tool use reason")?; let script = format!( r#"import jsonfrom pathlib import Pathimport syspayload = json.load(sys.stdin)with Path(r"{log_path}").open("a", encoding="utf-8") as handle: handle.write(json.dumps(payload) + "\n")print(json.dumps({{ "hookSpecificOutput": {{ "hookEventName": "PreToolUse", "permissionDecision": "deny", "permissionDecisionReason": {reason_json} }}}}))"#, log_path = log_path.display(), reason_json = reason_json, ); let hooks = serde_json::json!({ "hooks": { "PreToolUse": [{ "matcher": RMCP_HOOK_MATCHER, "hooks": [{ "type": "command", "command": format!("python3 {}", script_path.display()), "statusMessage": "running MCP pre tool use hook", }] }] } }); fs::write(&script_path, script).context("write pre tool use hook script")?; fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; Ok(())}fn write_updating_pre_tool_use_hook(home: &Path, updated_message: &str) -> Result<()> { let script_path = home.join("pre_tool_use_hook.py"); let log_path = home.join("pre_tool_use_hook_log.jsonl"); let updated_message_json = serde_json::to_string(updated_message).context("serialize updated MCP message")?; let script = format!( r#"import jsonfrom pathlib import Pathimport syspayload = json.load(sys.stdin)with Path(r"{log_path}").open("a", encoding="utf-8") as handle: handle.write(json.dumps(payload) + "\n")print(json.dumps({{ "hookSpecificOutput": {{ "hookEventName": "PreToolUse", "permissionDecision": "allow", "updatedInput": {{ "message": {updated_message_json} }} }}}}))"#, log_path = log_path.display(), updated_message_json = updated_message_json, ); let hooks = serde_json::json!({ "hooks": { "PreToolUse": [{ "matcher": RMCP_HOOK_MATCHER, "hooks": [{ "type": "command", "command": format!("python3 {}", script_path.display()), "statusMessage": "rewriting MCP pre tool input", }] }] } }); fs::write(&script_path, script).context("write updating pre tool use hook script")?; fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; Ok(())}fn write_post_tool_use_hook(home: &Path, additional_context: &str) -> Result<()> { let script_path = home.join("post_tool_use_hook.py"); let log_path = home.join("post_tool_use_hook_log.jsonl"); let additional_context_json = serde_json::to_string(additional_context).context("serialize post tool use context")?; let script = format!( r#"import jsonfrom pathlib import Pathimport syspayload = json.load(sys.stdin)with Path(r"{log_path}").open("a", encoding="utf-8") as handle: handle.write(json.dumps(payload) + "\n")print(json.dumps({{ "hookSpecificOutput": {{ "hookEventName": "PostToolUse", "additionalContext": {additional_context_json} }}}}))"#, log_path = log_path.display(), additional_context_json = additional_context_json, ); let hooks = serde_json::json!({ "hooks": { "PostToolUse": [{ "matcher": RMCP_HOOK_MATCHER, "hooks": [{ "type": "command", "command": format!("python3 {}", script_path.display()), "statusMessage": "running MCP post tool use hook", }] }] } }); fs::write(&script_path, script).context("write post tool use hook script")?; fs::write(home.join("hooks.json"), hooks.to_string()).context("write hooks.json")?; Ok(())}fn read_hook_inputs(home: &Path, log_name: &str) -> Result<Vec<Value>> { fs::read_to_string(home.join(log_name)) .with_context(|| format!("read {log_name}"))? .lines() .filter(|line| !line.trim().is_empty()) .map(|line| serde_json::from_str(line).with_context(|| format!("parse {log_name} line"))) .collect()}fn insert_rmcp_test_server(config: &mut Config, command: String, approval_mode: AppToolApproval) { let mut servers = config.mcp_servers.get().clone(); servers.insert( RMCP_SERVER.to_string(), McpServerConfig { transport: McpServerTransportConfig::Stdio { command, args: Vec::new(), env: None, env_vars: Vec::new(), cwd: None, }, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, required: false, supports_parallel_tool_calls: false, disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, default_tools_approval_mode: Some(approval_mode), enabled_tools: None, disabled_tools: None, scopes: None, oauth: None, oauth_resource: None, tools: HashMap::new(), }, ); config .mcp_servers .set(servers) .expect("test mcp servers should accept any configuration");}fn enable_hooks_and_rmcp_server( config: &mut Config, rmcp_test_server_bin: String, approval_mode: AppToolApproval, prefix_mcp_tool_names: bool,) { trust_discovered_hooks(config); enable_mcp_tool_name_features(config, prefix_mcp_tool_names); insert_rmcp_test_server(config, rmcp_test_server_bin, approval_mode);}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn pre_tool_use_blocks_mcp_tool_before_execution_with_legacy_prefixed_names() -> Result<()> { pre_tool_use_blocks_mcp_tool_before_execution( /*prefix_mcp_tool_names*/ true, RMCP_PREFIXED_NAMESPACE, ) .await}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn pre_tool_use_blocks_mcp_tool_before_execution_with_non_prefixed_names() -> Result<()> { pre_tool_use_blocks_mcp_tool_before_execution( /*prefix_mcp_tool_names*/ false, RMCP_UNPREFIXED_NAMESPACE, ) .await}async fn pre_tool_use_blocks_mcp_tool_before_execution( prefix_mcp_tool_names: bool, mcp_namespace: &'static str,) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; let call_id = "pretooluse-rmcp-echo"; let arguments = json!({ "message": RMCP_ECHO_MESSAGE }).to_string(); let responses = mount_sse_sequence( &server, vec![ sse(vec![ ev_response_created("resp-1"), ev_function_call_with_namespace(call_id, mcp_namespace, "echo", &arguments), ev_completed("resp-1"), ]), sse(vec![ ev_response_created("resp-2"), ev_assistant_message("msg-1", "mcp hook blocked it"), ev_completed("resp-2"), ]), ], ) .await; let block_reason = "blocked mcp pre hook"; let rmcp_test_server_bin = stdio_server_bin()?; let test = test_codex() .with_pre_build_hook(move |home| { write_pre_tool_use_hook(home, block_reason) .expect("failed to write MCP pre tool use hook fixture"); }) .with_config(move |config| { enable_hooks_and_rmcp_server( config, rmcp_test_server_bin, AppToolApproval::Approve, prefix_mcp_tool_names, ); }) .build(&server) .await?; wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP pre hook") .await?; let requests = responses.requests(); assert_eq!(requests.len(), 2); let output_item = requests[1].function_call_output(call_id); let output = output_item .get("output") .and_then(Value::as_str) .expect("blocked MCP tool output should be a string"); assert!( output.contains(&format!( "Tool call blocked by PreToolUse hook: {block_reason}. Tool: {RMCP_ECHO_TOOL_NAME}" )), "blocked MCP tool output should surface the hook reason and tool name", ); let hook_inputs = read_hook_inputs(test.codex_home_path(), "pre_tool_use_hook_log.jsonl")?; assert_eq!(hook_inputs.len(), 1); assert_eq!( json!({ "hook_event_name": hook_inputs[0]["hook_event_name"], "tool_name": hook_inputs[0]["tool_name"], "tool_use_id": hook_inputs[0]["tool_use_id"], "tool_input": hook_inputs[0]["tool_input"], }), json!({ "hook_event_name": "PreToolUse", "tool_name": RMCP_ECHO_TOOL_NAME, "tool_use_id": call_id, "tool_input": { "message": RMCP_ECHO_MESSAGE }, }) ); let transcript_path = hook_inputs[0]["transcript_path"] .as_str() .expect("pre tool use hook transcript_path should be a string"); assert!( Path::new(transcript_path).exists(), "pre tool use hook transcript_path should be materialized on disk", ); Ok(())}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn pre_tool_use_rewrites_mcp_tool_before_execution() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; let call_id = "pretooluse-rmcp-echo-rewrite"; let rewritten_message = "rewritten mcp hook input"; let arguments = json!({ "message": RMCP_ECHO_MESSAGE }).to_string(); let call_mock = mount_sse_once( &server, sse(vec![ ev_response_created("resp-1"), ev_function_call_with_namespace(call_id, RMCP_PREFIXED_NAMESPACE, "echo", &arguments), ev_completed("resp-1"), ]), ) .await; let final_mock = mount_sse_once( &server, sse(vec![ ev_response_created("resp-2"), ev_assistant_message("msg-1", "mcp pre hook rewrote it"), ev_completed("resp-2"), ]), ) .await; let rmcp_test_server_bin = stdio_server_bin()?; let test = test_codex() .with_pre_build_hook(move |home| { write_updating_pre_tool_use_hook(home, rewritten_message) .expect("failed to write MCP updating pre tool use hook fixture"); }) .with_config(move |config| { enable_hooks_and_rmcp_server( config, rmcp_test_server_bin, AppToolApproval::Approve, /*prefix_mcp_tool_names*/ true, ); }) .build(&server) .await?; wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP pre hook rewrite") .await?; let final_request = final_mock.single_request(); let output_item = final_request.function_call_output(call_id); let output = output_item .get("output") .and_then(Value::as_str) .expect("MCP tool output should be a string"); assert!( output.contains(&format!("ECHOING: {rewritten_message}")), "MCP tool should execute the rewritten input", ); assert!( !output.contains(RMCP_ECHO_MESSAGE), "MCP tool should not execute the original input", ); let hook_inputs = read_hook_inputs(test.codex_home_path(), "pre_tool_use_hook_log.jsonl")?; assert_eq!(hook_inputs.len(), 1); assert_eq!( hook_inputs[0]["tool_input"], json!({ "message": RMCP_ECHO_MESSAGE }), ); call_mock.single_request(); Ok(())}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn post_tool_use_records_mcp_tool_payload_and_context_with_legacy_prefixed_names()-> Result<()> { post_tool_use_records_mcp_tool_payload_and_context( /*prefix_mcp_tool_names*/ true, RMCP_PREFIXED_NAMESPACE, ) .await}#[tokio::test(flavor = "multi_thread", worker_threads = 2)]async fn post_tool_use_records_mcp_tool_payload_and_context_with_non_prefixed_names() -> Result<()>{ post_tool_use_records_mcp_tool_payload_and_context( /*prefix_mcp_tool_names*/ false, RMCP_UNPREFIXED_NAMESPACE, ) .await}async fn post_tool_use_records_mcp_tool_payload_and_context( prefix_mcp_tool_names: bool, mcp_namespace: &'static str,) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; let call_id = "posttooluse-rmcp-echo"; let arguments = json!({ "message": RMCP_ECHO_MESSAGE }).to_string(); let call_mock = mount_sse_once( &server, sse(vec![ ev_response_created("resp-1"), ev_function_call_with_namespace(call_id, mcp_namespace, "echo", &arguments), ev_completed("resp-1"), ]), ) .await; let final_mock = mount_sse_once( &server, sse(vec![ ev_response_created("resp-2"), ev_assistant_message("msg-1", "mcp post hook context observed"), ev_completed("resp-2"), ]), ) .await; let post_context = "Remember the MCP post-tool note."; let rmcp_test_server_bin = stdio_server_bin()?; let test = test_codex() .with_pre_build_hook(move |home| { write_post_tool_use_hook(home, post_context) .expect("failed to write MCP post tool use hook fixture"); }) .with_config(move |config| { enable_hooks_and_rmcp_server( config, rmcp_test_server_bin, AppToolApproval::Approve, prefix_mcp_tool_names, ); }) .build(&server) .await?; wait_for_mcp_server(&test.codex, RMCP_SERVER).await?; test.submit_turn("call the rmcp echo tool with the MCP post hook") .await?; let final_request = final_mock.single_request(); assert!( final_request .message_input_texts("developer") .contains(&post_context.to_string()), "follow-up request should include MCP post tool use additional context", ); let output_item = final_request.function_call_output(call_id); let output = output_item .get("output") .and_then(Value::as_str) .expect("MCP tool output should be a string"); assert!( output.contains(&format!("ECHOING: {RMCP_ECHO_MESSAGE}")), "MCP tool output should still reach the model", ); let hook_inputs = read_hook_inputs(test.codex_home_path(), "post_tool_use_hook_log.jsonl")?; assert_eq!(hook_inputs.len(), 1); assert_eq!( json!({ "hook_event_name": hook_inputs[0]["hook_event_name"], "tool_name": hook_inputs[0]["tool_name"], "tool_use_id": hook_inputs[0]["tool_use_id"], "tool_input": hook_inputs[0]["tool_input"], "tool_response": hook_inputs[0]["tool_response"], }), json!({ "hook_event_name": "PostToolUse", "tool_name": RMCP_ECHO_TOOL_NAME, "tool_use_id": call_id, "tool_input": { "message": RMCP_ECHO_MESSAGE }, "tool_response": { "content": [], "structuredContent": { "echo": format!("ECHOING: {RMCP_ECHO_MESSAGE}"), "env": null, }, "isError": false, }, }) ); let transcript_path = hook_inputs[0]["transcript_path"] .as_str() .expect("post tool use hook transcript_path should be a string"); assert!( Path::new(transcript_path).exists(), "post tool use hook transcript_path should be materialized on disk", ); call_mock.single_request(); Ok(())}