use std::sync::Arc;use std::time::Instant;use crate::function_tool::FunctionCallError;use crate::mcp_tool_call::handle_mcp_tool_call;use crate::original_image_detail::can_request_original_image_detail;use crate::tools::context::McpToolOutput;use crate::tools::context::ToolInvocation;use crate::tools::context::ToolPayload;use crate::tools::context::boxed_tool_output;use crate::tools::flat_tool_name;use crate::tools::hook_names::HookToolName;use crate::tools::registry::CoreToolRuntime;use crate::tools::registry::PostToolUsePayload;use crate::tools::registry::PreToolUsePayload;use crate::tools::registry::ToolExecutor;use crate::tools::registry::ToolTelemetryTags;use codex_mcp::ToolInfo;use codex_tools::ResponsesApiNamespace;use codex_tools::ResponsesApiNamespaceTool;use codex_tools::ToolName;use codex_tools::ToolSearchInfo;use codex_tools::ToolSearchSourceInfo;use codex_tools::ToolSpec;use codex_tools::mcp_tool_to_responses_api_tool;use serde_json::Map;use serde_json::Value;const LEGACY_MCP_TOOL_NAME_PREFIX: &str = "mcp__";const MCP_TOOL_NAME_DELIMITER: &str = "__";pub struct McpHandler {    tool_info: ToolInfo,    spec: ToolSpec,}impl McpHandler {    pub fn new(tool_info: ToolInfo) -> Result<Self, serde_json::Error> {        let spec = create_tool_spec(&tool_info)?;        Ok(Self { tool_info, spec })    }    fn hook_tool_name(&self) -> HookToolName {        HookToolName::new(ensure_mcp_prefix(&join_tool_name(&self.tool_name())))    }}fn join_tool_name(tool_name: &ToolName) -> String {    match tool_name.namespace.as_deref() {        Some(namespace) => {            let namespace = namespace.trim_end_matches('_');            let name = tool_name.name.trim_start_matches('_');            format!("{namespace}{MCP_TOOL_NAME_DELIMITER}{name}")        }        None => tool_name.name.clone(),    }}fn ensure_mcp_prefix(name: &str) -> String {    if name.starts_with(LEGACY_MCP_TOOL_NAME_PREFIX) {        name.to_string()    } else {        format!("{LEGACY_MCP_TOOL_NAME_PREFIX}{name}")    }}impl ToolExecutor<ToolInvocation> for McpHandler {    fn tool_name(&self) -> ToolName {        self.tool_info.canonical_tool_name()    }    fn spec(&self) -> ToolSpec {        self.spec.clone()    }    fn supports_parallel_tool_calls(&self) -> bool {        // Correctly implemented MCP servers should tolerate parallel calls to        // tools that advertise themselves as read-only.        self.tool_info.supports_parallel_tool_calls            || self                .tool_info                .tool                .annotations                .as_ref()                .and_then(|annotations| annotations.read_only_hint)                .unwrap_or(false)    }    fn search_info(&self) -> Option<ToolSearchInfo> {        let source_name = self            .tool_info            .connector_name            .as_deref()            .map(str::trim)            .filter(|connector_name| !connector_name.is_empty())            .unwrap_or_else(|| self.tool_info.server_name.trim());        let source_info = (!source_name.is_empty()).then(|| ToolSearchSourceInfo {            name: source_name.to_string(),            description: self                .tool_info                .namespace_description                .as_deref()                .map(str::trim)                .filter(|description| !description.is_empty())                .map(str::to_string),        });        ToolSearchInfo::from_spec(            build_mcp_search_text(&self.tool_info),            self.spec(),            source_info,        )    }    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {        Box::pin(self.handle_call(invocation))    }}impl McpHandler {    async fn handle_call(        &self,        invocation: ToolInvocation,    ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {        let ToolInvocation {            session,            turn,            call_id,            payload,            ..        } = invocation;        let payload = match payload {            ToolPayload::Function { arguments } => arguments,            _ => {                return Err(FunctionCallError::RespondToModel(                    "mcp handler received unsupported payload".to_string(),                ));            }        };        let started = Instant::now();        let result = handle_mcp_tool_call(            Arc::clone(&session),            &turn,            call_id.clone(),            self.tool_info.server_name.clone(),            self.tool_info.tool.name.to_string(),            self.hook_tool_name(),            payload,        )        .await;        Ok(boxed_tool_output(McpToolOutput {            result: result.result,            tool_input: result.tool_input,            wall_time: started.elapsed(),            original_image_detail_supported: can_request_original_image_detail(&turn.model_info),            truncation_policy: turn.truncation_policy,        }))    }}impl CoreToolRuntime for McpHandler {    fn telemetry_tags<'a>(        &'a self,        _invocation: &'a ToolInvocation,    ) -> futures::future::BoxFuture<'a, ToolTelemetryTags> {        Box::pin(async {            let mut tags = vec![("mcp_server", self.tool_info.server_name.clone())];            if let Some(origin) = &self.tool_info.server_origin {                tags.push(("mcp_server_origin", origin.clone()));            }            tags        })    }    fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option<PreToolUsePayload> {        let ToolPayload::Function { arguments } = &invocation.payload else {            return None;        };        Some(PreToolUsePayload {            tool_name: self.hook_tool_name(),            tool_input: mcp_hook_tool_input(arguments),        })    }    fn with_updated_hook_input(        &self,        mut invocation: ToolInvocation,        updated_input: Value,    ) -> Result<ToolInvocation, FunctionCallError> {        invocation.payload = match invocation.payload {            ToolPayload::Function { .. } => ToolPayload::Function {                arguments: serde_json::to_string(&updated_input).map_err(|err| {                    FunctionCallError::RespondToModel(format!(                        "failed to serialize rewritten MCP arguments: {err}"                    ))                })?,            },            payload => {                return Err(FunctionCallError::RespondToModel(format!(                    "tool {} does not support hook input rewriting for payload {payload:?}",                    self.tool_name()                )));            }        };        Ok(invocation)    }    fn post_tool_use_payload(        &self,        invocation: &ToolInvocation,        result: &dyn crate::tools::context::ToolOutput,    ) -> Option<PostToolUsePayload> {        let ToolPayload::Function { .. } = &invocation.payload else {            return None;        };        let tool_response =            result.post_tool_use_response(&invocation.call_id, &invocation.payload)?;        Some(PostToolUsePayload {            tool_name: self.hook_tool_name(),            tool_use_id: invocation.call_id.clone(),            tool_input: result.post_tool_use_input(&invocation.payload)?,            tool_response,        })    }}fn create_tool_spec(tool_info: &ToolInfo) -> Result<ToolSpec, serde_json::Error> {    let tool_name = tool_info.canonical_tool_name();    let tool = mcp_tool_to_responses_api_tool(&tool_name, &tool_info.tool)?;    let description = tool_info        .namespace_description        .as_deref()        .map(str::trim)        .filter(|description| !description.is_empty())        .map(str::to_string)        .or_else(|| {            tool_info                .connector_name                .as_deref()                .map(str::trim)                .filter(|connector_name| !connector_name.is_empty())                .map(|connector_name| format!("Tools for working with {connector_name}."))        })        .unwrap_or_default();    Ok(ToolSpec::Namespace(ResponsesApiNamespace {        name: tool_info.callable_namespace.clone(),        description,        tools: vec![ResponsesApiNamespaceTool::Function(tool)],    }))}fn mcp_hook_tool_input(raw_arguments: &str) -> Value {    if raw_arguments.trim().is_empty() {        return Value::Object(Map::new());    }    serde_json::from_str(raw_arguments).unwrap_or_else(|_| Value::String(raw_arguments.to_string()))}fn build_mcp_search_text(info: &ToolInfo) -> String {    let tool_name = info.canonical_tool_name();    let mut schema_properties = info        .tool        .input_schema        .get("properties")        .and_then(serde_json::Value::as_object)        .map(|map| map.keys().cloned().collect::<Vec<_>>())        .unwrap_or_default();    schema_properties.sort();    let mut parts = vec![        flat_tool_name(&tool_name).into_owned(),        info.callable_name.clone(),        info.tool.name.to_string(),        info.server_name.clone(),    ];    if let Some(title) = info.tool.title.as_deref().map(str::trim)        && !title.is_empty()    {        parts.push(title.to_string());    }    if let Some(description) = info.tool.description.as_deref().map(str::trim)        && !description.is_empty()    {        parts.push(description.to_string());    }    if let Some(connector_name) = info.connector_name.as_deref().map(str::trim)        && !connector_name.is_empty()    {        parts.push(connector_name.to_string());    }    if let Some(namespace_description) = info.namespace_description.as_deref().map(str::trim)        && !namespace_description.is_empty()    {        parts.push(namespace_description.to_string());    }    parts.extend(        info.plugin_display_names            .iter()            .map(String::as_str)            .map(str::trim)            .filter(|display_name| !display_name.is_empty())            .map(str::to_string),    );    parts.extend(schema_properties);    parts.join(" ")}#[cfg(test)]#[path = "mcp_search_tests.rs"]mod search_tests;#[cfg(test)]mod tests {    use super::*;    use crate::session::tests::make_session_and_context;    use crate::tools::context::ToolCallSource;    use crate::tools::hook_names::HookToolName;    use crate::tools::registry::PostToolUsePayload;    use crate::tools::registry::PreToolUsePayload;    use crate::turn_diff_tracker::TurnDiffTracker;    use pretty_assertions::assert_eq;    use serde_json::json;    use std::time::Duration;    use tokio::sync::Mutex;    #[tokio::test]    async fn mcp_pre_tool_use_payload_uses_prefixed_tool_name_and_raw_args() {        let payload = ToolPayload::Function {            arguments: json!({                "entities": [{                    "name": "Ada",                    "entityType": "person"                }]            })            .to_string(),        };        let (session, turn) = make_session_and_context().await;        let handler = McpHandler::new(tool_info("memory", "memory", "create_entities"))            .expect("MCP tool spec should build");        assert_eq!(            handler.pre_tool_use_payload(&ToolInvocation {                session: session.into(),                turn: turn.into(),                cancellation_token: tokio_util::sync::CancellationToken::new(),                tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),                call_id: "call-mcp-pre".to_string(),                tool_name: codex_tools::ToolName::namespaced("memory", "create_entities"),                source: ToolCallSource::Direct,                payload,            }),            Some(PreToolUsePayload {                tool_name: HookToolName::new("mcp__memory__create_entities"),                tool_input: json!({                    "entities": [{                        "name": "Ada",                        "entityType": "person"                    }]                }),            })        );    }    #[tokio::test]    async fn mcp_pre_tool_use_payload_keeps_builtin_like_tool_names_namespaced() {        let payload = ToolPayload::Function {            arguments: json!({ "message": "hello" }).to_string(),        };        let (session, turn) = make_session_and_context().await;        let handler = McpHandler::new(tool_info("foo", "mcp__foo", "exec_command"))            .expect("MCP tool spec should build");        assert_eq!(            handler.pre_tool_use_payload(&ToolInvocation {                session: session.into(),                turn: turn.into(),                cancellation_token: tokio_util::sync::CancellationToken::new(),                tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),                call_id: "call-mcp-pre-builtin-like".to_string(),                tool_name: codex_tools::ToolName::namespaced("mcp__foo", "exec_command"),                source: ToolCallSource::Direct,                payload,            }),            Some(PreToolUsePayload {                tool_name: HookToolName::new("mcp__foo__exec_command"),                tool_input: json!({ "message": "hello" }),            })        );    }    #[tokio::test]    async fn mcp_updated_input_rewrites_builtin_like_tool_names_as_mcp() {        let payload = ToolPayload::Function {            arguments: json!({ "message": "hello" }).to_string(),        };        let (session, turn) = make_session_and_context().await;        let handler = McpHandler::new(tool_info("foo", "mcp__foo", "exec_command"))            .expect("MCP tool spec should build");        let invocation = handler            .with_updated_hook_input(                ToolInvocation {                    session: session.into(),                    turn: turn.into(),                    cancellation_token: tokio_util::sync::CancellationToken::new(),                    tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),                    call_id: "call-mcp-rewrite-builtin-like".to_string(),                    tool_name: codex_tools::ToolName::namespaced("mcp__foo", "exec_command"),                    source: ToolCallSource::Direct,                    payload,                },                json!({ "message": "rewritten" }),            )            .expect("MCP rewrite should succeed");        let ToolPayload::Function { arguments } = invocation.payload else {            panic!("builtin-like MCP tool should stay function-shaped");        };        assert_eq!(arguments, json!({ "message": "rewritten" }).to_string());    }    #[tokio::test]    async fn mcp_post_tool_use_payload_uses_prefixed_tool_name_args_and_result() {        let payload = ToolPayload::Function {            arguments: json!({ "path": "/tmp/notes.txt" }).to_string(),        };        let output = McpToolOutput {            result: codex_protocol::mcp::CallToolResult {                content: vec![json!({                    "type": "text",                    "text": "notes"                })],                structured_content: Some(json!({ "bytes": 5 })),                is_error: None,                meta: None,            },            tool_input: json!({                "path": {                    "file_id": "file_123"                }            }),            wall_time: Duration::from_millis(42),            original_image_detail_supported: true,            truncation_policy: codex_utils_output_truncation::TruncationPolicy::Bytes(1024),        };        let (session, turn) = make_session_and_context().await;        let handler = McpHandler::new(tool_info("filesystem", "filesystem", "read_file"))            .expect("MCP tool spec should build");        let invocation = ToolInvocation {            session: session.into(),            turn: turn.into(),            cancellation_token: tokio_util::sync::CancellationToken::new(),            tracker: Arc::new(Mutex::new(TurnDiffTracker::new())),            call_id: "call-mcp-post".to_string(),            tool_name: codex_tools::ToolName::namespaced("filesystem", "read_file"),            source: ToolCallSource::Direct,            payload,        };        assert_eq!(            handler.post_tool_use_payload(&invocation, &output),            Some(PostToolUsePayload {                tool_name: HookToolName::new("mcp__filesystem__read_file"),                tool_use_id: "call-mcp-post".to_string(),                tool_input: json!({                    "path": {                        "file_id": "file_123"                    }                }),                tool_response: json!({                    "content": [{                        "type": "text",                        "text": "notes"                    }],                    "structuredContent": { "bytes": 5 }                }),            })        );    }    #[test]    fn mcp_read_only_hint_supports_parallel_calls_without_server_opt_in() {        let mut read_only_info = tool_info("foo", "mcp__foo__", "read");        read_only_info.tool.annotations = Some(rmcp::model::ToolAnnotations::new().read_only(true));        assert!(            McpHandler::new(read_only_info)                .expect("MCP tool spec should build")                .supports_parallel_tool_calls()        );    }    #[test]    fn mcp_parallel_calls_require_read_only_hint_or_server_opt_in() {        let missing_hint_info = tool_info("foo", "mcp__foo__", "unannotated");        assert!(            !McpHandler::new(missing_hint_info)                .expect("MCP tool spec should build")                .supports_parallel_tool_calls()        );        let mut writable_info = tool_info("foo", "mcp__foo__", "write");        writable_info.tool.annotations = Some(rmcp::model::ToolAnnotations::new().read_only(false));        assert!(            !McpHandler::new(writable_info)                .expect("MCP tool spec should build")                .supports_parallel_tool_calls()        );        let mut server_opt_in_info = tool_info("foo", "mcp__foo__", "server_opt_in");        server_opt_in_info.supports_parallel_tool_calls = true;        assert!(            McpHandler::new(server_opt_in_info)                .expect("MCP tool spec should build")                .supports_parallel_tool_calls()        );    }    fn tool_info(server_name: &str, callable_namespace: &str, tool_name: &str) -> ToolInfo {        ToolInfo {            server_name: server_name.to_string(),            supports_parallel_tool_calls: false,            server_origin: None,            callable_name: tool_name.to_string(),            callable_namespace: callable_namespace.to_string(),            namespace_description: None,            tool: rmcp::model::Tool::new_with_raw(                tool_name.to_string(),                None,                Arc::new(rmcp::model::object(serde_json::json!({                    "type": "object",                }))),            ),            connector_id: None,            connector_name: None,            plugin_display_names: Vec::new(),        }    }}