diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 454e1b6bce..a21fa49be1 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -17,6 +17,7 @@ use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::ApprovalContext; use crate::tools::hook_names::HookToolName; +use crate::tools::lifecycle::process_mcp_tool_result; use crate::tools::sandboxing::ApprovalAction; use crate::tools::sandboxing::ToolError; use crate::turn_metadata::McpTurnMetadataContext; @@ -29,6 +30,7 @@ use codex_config::types::AppToolApproval; use codex_connectors::AppToolPolicy; use codex_connectors::AppToolPolicyEvaluator; use codex_connectors::AppToolPolicyInput; +use codex_extension_api::McpToolContext; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; @@ -422,7 +424,7 @@ async fn handle_approved_mcp_tool_call( .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())); let result = async { let result = async { - let result = prepared_call + let mut result = prepared_call .call_with_preparation(/*requested_timeout*/ None, || async { if let McpToolApprovalApplication::Apply { decision, policy } = &approval_application @@ -501,6 +503,19 @@ async fn handle_approved_mcp_tool_call( }) .await .map_err(|error| format!("tool call error: {error:?}"))?; + let mcp_tool = McpToolContext::from_prepared_call( + &prepared_call, + turn_context.config.mcp_servers.get().get(&server), + ); + process_mcp_tool_result( + sess, + turn_context, + call_id, + &mcp_tool, + &tool_input, + &mut result, + ) + .await; let result = sanitize_mcp_tool_result_for_model( &turn_context.model_info().input_modalities, Ok(result), diff --git a/codex-rs/core/src/tools/lifecycle.rs b/codex-rs/core/src/tools/lifecycle.rs index ddaed3b900..17108f4bb3 100644 --- a/codex-rs/core/src/tools/lifecycle.rs +++ b/codex-rs/core/src/tools/lifecycle.rs @@ -1,10 +1,12 @@ use std::sync::Arc; use codex_extension_api::McpToolContext; +use codex_extension_api::McpToolResultInput; use codex_extension_api::ToolCallOutcome; use codex_extension_api::ToolCallSource as ExtensionToolCallSource; use codex_extension_api::ToolFinishInput; use codex_extension_api::ToolStartInput; +use codex_protocol::mcp::CallToolResult; use codex_tools::ToolName; use crate::session::session::Session; @@ -45,6 +47,30 @@ pub(crate) async fn notify_tool_start( } } +pub(crate) async fn process_mcp_tool_result( + session: &Session, + turn: &TurnContext, + call_id: &str, + mcp_tool: &McpToolContext, + arguments: &serde_json::Value, + result: &mut CallToolResult, +) { + for contributor in session.services.extensions.tool_lifecycle_contributors() { + contributor + .on_mcp_tool_result(McpToolResultInput { + session_store: &session.services.session_extension_data, + thread_store: &session.services.thread_extension_data, + turn_store: turn.extension_data.as_ref(), + turn_id: &turn.sub_id, + call_id, + mcp_tool, + arguments, + result, + }) + .await; + } +} + pub(crate) async fn notify_tool_finish(invocation: &ToolInvocation, outcome: ToolCallOutcome) { notify_tool_finish_parts( invocation.session.as_ref(), diff --git a/codex-rs/core/tests/suite/tool_lifecycle.rs b/codex-rs/core/tests/suite/tool_lifecycle.rs index 82007cf657..396d90bf18 100644 --- a/codex-rs/core/tests/suite/tool_lifecycle.rs +++ b/codex-rs/core/tests/suite/tool_lifecycle.rs @@ -2,22 +2,29 @@ use std::fs; use std::path::Path; use std::sync::Arc; use std::sync::Mutex; +use std::time::Duration; use anyhow::Context; use anyhow::Result; +use codex_core::TurnInputRequest; use codex_core::config::Config; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; +use codex_extension_api::McpToolResultInput; use codex_extension_api::McpToolSource; use codex_extension_api::ResponseItem; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; use codex_extension_api::ToolStartInput; +use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::mcp::CallToolResult; use codex_protocol::models::ContentItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::UserInput; use codex_utils_path_uri::PathUri; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::SEARCH_CALENDAR_LIST_TOOL; @@ -29,11 +36,19 @@ use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::skip_if_wine_exec; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; use test_case::test_case; +use tokio::sync::Notify; +use tokio::sync::oneshot; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; struct RecordedHistory { call_id: String, @@ -286,6 +301,194 @@ async fn tool_start_receives_executed_mcp_call_for_connector( Ok(()) } +#[derive(Clone, Copy)] +enum McpCallMode { + Direct, + CodeMode, +} + +#[derive(Clone, Copy)] +enum McpResultProcessing { + Unchanged, + Replace, + Fail, +} + +struct BlockingMcpResultContributor { + entered: Mutex>>, + release: Notify, + replacement: Option, +} + +impl ToolLifecycleContributor for BlockingMcpResultContributor { + fn on_mcp_tool_result<'a>(&'a self, input: McpToolResultInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + assert_eq!(input.mcp_tool.source(), &McpToolSource::Connector); + assert_eq!(input.mcp_tool.tool_info().tool.name, "calendar_list_events"); + assert_eq!(input.arguments, &json!({"query": "callback test"})); + self.entered + .lock() + .expect("callback gate lock") + .take() + .expect("callback should run once") + .send(input.result.clone()) + .expect("test should be waiting for the result"); + self.release.notified().await; + if let Some(replacement) = &self.replacement { + *input.result = replacement.clone(); + } + }) + } +} + +#[test_case(McpCallMode::Direct, McpResultProcessing::Unchanged; "direct_unchanged")] +#[test_case(McpCallMode::Direct, McpResultProcessing::Replace; "direct_replaced")] +#[test_case(McpCallMode::Direct, McpResultProcessing::Fail; "direct_error")] +#[test_case(McpCallMode::CodeMode, McpResultProcessing::Replace; "code_mode_replaced")] +#[test_case(McpCallMode::CodeMode, McpResultProcessing::Fail; "code_mode_error")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_result_processing_precedes_completion( + mode: McpCallMode, + processing: McpResultProcessing, +) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let apps_server = AppsTestServer::mount(&server).await?; + let server_result: CallToolResult = serde_json::from_value(json!({ + "content": [{"type": "text", "text": "server result"}], + "structuredContent": {"text": "server result"}, + "isError": false, + "_meta": {"private-mcp-metadata": "server"}, + }))?; + let response_result = server_result.clone(); + Mock::given(body_partial_json(json!({"method": "tools/call"}))) + .respond_with(move |request: &Request| { + let body: Value = serde_json::from_slice(&request.body).expect("MCP request JSON"); + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body["id"], + "result": response_result, + })) + }) + .with_priority(1) + .mount(&server) + .await; + + let replacement = match processing { + McpResultProcessing::Unchanged => None, + McpResultProcessing::Replace | McpResultProcessing::Fail => { + Some(serde_json::from_value(json!({ + "content": [{"type": "text", "text": "extension result"}], + "structuredContent": {"text": "extension result"}, + "isError": matches!(processing, McpResultProcessing::Fail), + "_meta": {"private-mcp-metadata": "extension"}, + }))?) + } + }; + let expected_result = replacement.clone().unwrap_or_else(|| server_result.clone()); + let (entered_tx, entered_rx) = oneshot::channel(); + let contributor = Arc::new(BlockingMcpResultContributor { + entered: Mutex::new(Some(entered_tx)), + release: Notify::new(), + replacement, + }); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.tool_lifecycle_contributor(contributor.clone()); + let test = apps_enabled_builder(apps_server.chatgpt_base_url) + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + if matches!(mode, McpCallMode::CodeMode) { + let _ = config.features.enable(Feature::CodeMode); + } + }) + .build_with_auto_env(&server) + .await?; + wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + + let call_id = "mcp-result-call"; + let call = match mode { + McpCallMode::Direct => responses::ev_function_call_with_namespace( + call_id, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_LIST_TOOL, + r#"{"query":"callback test"}"#, + ), + McpCallMode::CodeMode => responses::ev_custom_tool_call( + call_id, + "exec", + r#"text(await tools.mcp__codex_apps__calendar_list_events({query: "callback test"}));"#, + ), + }; + responses::mount_sse_once( + &server, + responses::sse(vec![call, responses::ev_completed("first-response")]), + ) + .await; + let follow_up = responses::mount_sse_once( + &server, + responses::sse(vec![responses::ev_completed("second-response")]), + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "List my calendar events.".to_string(), + text_elements: Vec::new(), + }])) + .await?; + + assert_eq!( + timeout(Duration::from_secs(30), entered_rx).await??, + server_result + ); + assert!( + timeout( + Duration::from_millis(100), + wait_for_event(&test.codex, |event| matches!( + event, + EventMsg::McpToolCallEnd(_) | EventMsg::TurnComplete(_) + )), + ) + .await + .is_err(), + "completion must wait for the extension" + ); + assert!(follow_up.requests().is_empty()); + + contributor.release.notify_one(); + let EventMsg::McpToolCallEnd(end) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::McpToolCallEnd(_)) + }) + .await + else { + unreachable!(); + }; + assert_eq!( + end.result.expect("MCP server returned a result"), + expected_result + ); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let request = follow_up.single_request(); + let output = match mode { + McpCallMode::Direct => request.function_call_output(call_id), + McpCallMode::CodeMode => request.custom_tool_call_output(call_id), + }; + let output = output["output"].to_string(); + assert!( + output.contains( + expected_result.content[0]["text"] + .as_str() + .expect("fixture should contain text") + ) + ); + assert!(!output.contains("private-mcp-metadata")); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tool_start_receives_frozen_host_plugin_root() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 8e9fd0a1f9..218b5eef79 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -43,6 +43,7 @@ pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; pub use tool_lifecycle::McpToolContext; +pub use tool_lifecycle::McpToolResultInput; pub use tool_lifecycle::McpToolSource; pub use tool_lifecycle::ToolCallOutcome; pub use tool_lifecycle::ToolFinishInput; @@ -318,9 +319,9 @@ pub trait ToolContributor: Send + Sync { /// Contributor for host-owned tool lifecycle gates. /// -/// Implementations should use these callbacks to observe tool execution and its -/// exposed input without rewriting the invocation. Use `ToolContributor` for -/// owning a tool implementation and hooks for policy that changes tool payloads. +/// Implementations can observe tool execution and process MCP responses without +/// rewriting the invocation. Use `ToolContributor` for owning a tool implementation +/// and hooks for policy that changes tool payloads. pub trait ToolLifecycleContributor: Send + Sync { /// Called after pre-tool hooks finalize an invocation and before execution. /// @@ -330,6 +331,11 @@ pub trait ToolLifecycleContributor: Send + Sync { Box::pin(std::future::ready(())) } + /// Runs before the MCP result is sent to the client and model. + fn on_mcp_tool_result<'a>(&'a self, _input: McpToolResultInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(std::future::ready(())) + } + /// Called after the tool call returns, is blocked, fails, or is cancelled. /// /// A matching start callback does not exist when execution is blocked, diff --git a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs index 484202d7f9..cdf42a920d 100644 --- a/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/tool_lifecycle.rs @@ -6,6 +6,7 @@ use codex_config::McpServerConfig; use codex_mcp::McpServerSource; use codex_mcp::PreparedMcpCall; use codex_mcp::ResolvedMcpServer; +use codex_protocol::mcp::CallToolResult; use codex_tools::ToolCallSource; use codex_tools::ToolName; use codex_tools::ToolPayload; @@ -138,6 +139,28 @@ pub struct ToolStartInput<'a> { pub source: ToolCallSource, } +/// Input supplied after an MCP server responds, before the host reports completion. +pub struct McpToolResultInput<'a> { + /// Store scoped to the host session runtime. + pub session_store: &'a ExtensionData, + /// Store scoped to this thread runtime. + pub thread_store: &'a ExtensionData, + /// Store scoped to this turn runtime. + pub turn_store: &'a ExtensionData, + /// Current turn submission id. + pub turn_id: &'a str, + /// Host tool call id, also used in the MCP completion notification. + pub call_id: &'a str, + /// Read-only metadata and provenance from the exact MCP call that executed. + pub mcp_tool: &'a McpToolContext, + /// Tool arguments after host-side rewriting, including file uploads. + pub arguments: &'a serde_json::Value, + /// Server response, including `_meta`. Changes feed the normal client and model output paths. + /// + /// Arguments and results can contain sensitive plaintext and must not be logged. + pub result: &'a mut CallToolResult, +} + /// Input supplied when the host finishes executing one tool call. pub struct ToolFinishInput<'a> { /// Store scoped to the host session runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 209c10e10e..d927ab65a1 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -51,6 +51,7 @@ pub use contributors::McpServerContribution; pub use contributors::McpServerContributionContext; pub use contributors::McpServerContributor; pub use contributors::McpToolContext; +pub use contributors::McpToolResultInput; pub use contributors::McpToolSource; pub use contributors::PreviousWorldStateSection; pub use contributors::PromptFragment;