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}