Skip to main content

tycode_core/mcp/
tool.rs

1use serde_json::{json, Value};
2use std::sync::Arc;
3use tokio::sync::RwLock;
4
5use super::{McpModuleInner, McpToolDef};
6use crate::chat::events::{ToolExecutionResult, ToolRequest as ToolRequestEvent, ToolRequestType};
7use crate::tools::r#trait::{
8    ContinuationPreference, ToolCallHandle, ToolCategory, ToolExecutor, ToolOutput, ToolRequest,
9};
10
11fn format_mcp_content(content: &rmcp::model::Content) -> String {
12    match &content.raw {
13        rmcp::model::RawContent::Text(text) => text.text.clone(),
14        rmcp::model::RawContent::Image(img) => {
15            format!("[Image: {} bytes, type: {}]", img.data.len(), img.mime_type)
16        }
17        rmcp::model::RawContent::Resource(_) => "[Resource data]".to_string(),
18        rmcp::model::RawContent::Audio(audio) => {
19            format!(
20                "[Audio: {} bytes, type: {}]",
21                audio.data.len(),
22                audio.mime_type
23            )
24        }
25    }
26}
27
28pub struct McpTool {
29    name: String,
30    description: String,
31    input_schema: Value,
32    server_name: String,
33    mcp_tool_name: String,
34    inner: Arc<RwLock<McpModuleInner>>,
35}
36
37impl McpTool {
38    pub(crate) fn new(
39        def: &McpToolDef,
40        inner: Arc<RwLock<McpModuleInner>>,
41    ) -> anyhow::Result<Self> {
42        let input_schema = serde_json::to_value(def.tool.input_schema.clone())
43            .map_err(|e| anyhow::anyhow!("Failed to serialize MCP tool input schema: {e:?}"))?;
44
45        Ok(Self {
46            name: def.name.clone(),
47            description: def.tool.description.as_deref().unwrap_or("").to_string(),
48            input_schema,
49            server_name: def.server_name.clone(),
50            mcp_tool_name: def.tool.name.to_string(),
51            inner,
52        })
53    }
54
55    pub fn get_server_name(&self) -> &str {
56        &self.server_name
57    }
58}
59
60struct McpToolHandle {
61    server_name: String,
62    tool_name: String,
63    mcp_tool_name: String,
64    arguments: Option<Value>,
65    tool_use_id: String,
66    inner: Arc<RwLock<McpModuleInner>>,
67}
68
69#[async_trait::async_trait(?Send)]
70impl ToolCallHandle for McpToolHandle {
71    fn tool_request(&self) -> ToolRequestEvent {
72        ToolRequestEvent {
73            tool_call_id: self.tool_use_id.clone(),
74            tool_name: self.tool_name.clone(),
75            tool_type: ToolRequestType::Other {
76                args: json!({
77                    "server": self.server_name,
78                    "tool": self.mcp_tool_name,
79                    "arguments": self.arguments
80                }),
81            },
82        }
83    }
84
85    async fn execute(self: Box<Self>) -> ToolOutput {
86        let mut inner = self.inner.write().await;
87        let client = match inner.clients.get_mut(&self.server_name) {
88            Some(c) => c,
89            None => {
90                return ToolOutput::Result {
91                    content: format!("MCP server '{}' not found", self.server_name),
92                    is_error: true,
93                    continuation: ContinuationPreference::Continue,
94                    ui_result: ToolExecutionResult::Error {
95                        short_message: "Server not found".to_string(),
96                        detailed_message: format!("MCP server '{}' not found", self.server_name),
97                    },
98                };
99            }
100        };
101
102        match client.call_tool(&self.mcp_tool_name, self.arguments).await {
103            Ok(result) => {
104                let output = result
105                    .content
106                    .iter()
107                    .map(format_mcp_content)
108                    .collect::<Vec<_>>()
109                    .join("\n");
110
111                ToolOutput::Result {
112                    content: output.clone(),
113                    is_error: false,
114                    continuation: ContinuationPreference::Continue,
115                    ui_result: ToolExecutionResult::Other {
116                        result: json!({ "mcp_result": output }),
117                    },
118                }
119            }
120            Err(e) => ToolOutput::Result {
121                content: format!("MCP tool call failed: {e:?}"),
122                is_error: true,
123                continuation: ContinuationPreference::Continue,
124                ui_result: ToolExecutionResult::Error {
125                    short_message: "MCP call failed".to_string(),
126                    detailed_message: format!("MCP tool call failed: {e:?}"),
127                },
128            },
129        }
130    }
131}
132
133#[async_trait::async_trait(?Send)]
134impl ToolExecutor for McpTool {
135    fn name(&self) -> String {
136        self.name.clone()
137    }
138
139    fn description(&self) -> String {
140        self.description.clone()
141    }
142
143    fn input_schema(&self) -> Value {
144        self.input_schema.clone()
145    }
146
147    fn category(&self) -> ToolCategory {
148        ToolCategory::Execution
149    }
150
151    async fn process(
152        &self,
153        request: &ToolRequest,
154    ) -> Result<Box<dyn ToolCallHandle>, anyhow::Error> {
155        Ok(Box::new(McpToolHandle {
156            server_name: self.server_name.clone(),
157            tool_name: self.name.clone(),
158            mcp_tool_name: self.mcp_tool_name.clone(),
159            arguments: Some(request.arguments.clone()),
160            tool_use_id: request.tool_use_id.clone(),
161            inner: self.inner.clone(),
162        }))
163    }
164}
165
166pub fn mcp_tool_definition(def: &McpToolDef) -> anyhow::Result<crate::ai::ToolDefinition> {
167    let input_schema = serde_json::to_value(def.tool.input_schema.clone())
168        .map_err(|e| anyhow::anyhow!("Failed to serialize MCP tool input schema: {e:?}"))?;
169
170    Ok(crate::ai::ToolDefinition {
171        name: def.name.clone(),
172        description: def.tool.description.as_deref().unwrap_or("").to_string(),
173        input_schema,
174    })
175}