Let extensions process MCP tool results (#41202)

## What changed

- Add `ToolLifecycleContributor::on_mcp_tool_result` with access to the
  executed MCP tool context, rewritten arguments, extension data stores, and
  mutable server result.
- Run contributors before publishing MCP completion and before preparing the
  result for the model, so extensions can inspect or replace successful and
  error results.

## Testing

- Cover unchanged, replaced, and error results for direct MCP calls, plus
  replaced and error results through Code Mode.
- Verify completion waits for result processing and the processed result flows
  to both completion events and subsequent model input.

GitOrigin-RevId: ece7acfe9c5eed4778cb78a29c3ffd06a74e1bcf

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),

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(),

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<Option<oneshot::Sender<CallToolResult>>>,
+    release: Notify,
+    replacement: Option<CallToolResult>,
+}
+
+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::<Config>::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(()));

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
+/// Implementations can observe tool execution and process MCP responses without
-/// exposed input without rewriting the invocation. Use `ToolContributor` for
+/// rewriting the invocation. Use `ToolContributor` for owning a tool implementation
-/// owning a tool implementation and hooks for policy that changes tool payloads.
+/// 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,

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.

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;