Skip to main content

a3s_code_core/mcp/
tools.rs

1//! MCP Tools Integration
2//!
3//! Integrates MCP tools with the A3S Code tool system.
4
5use crate::mcp::manager::{tool_result_to_string, McpManager};
6use crate::mcp::protocol::McpTool;
7use crate::tools::{Tool, ToolContext, ToolOutput};
8use anyhow::Result;
9use async_trait::async_trait;
10use std::sync::Arc;
11
12/// MCP tool wrapper that implements the Tool trait
13pub struct McpToolWrapper {
14    /// Full tool name (mcp__server__tool)
15    full_name: String,
16    /// Original MCP tool definition
17    mcp_tool: McpTool,
18    /// Server name
19    server_name: String,
20    /// MCP manager reference
21    manager: Arc<McpManager>,
22}
23
24impl McpToolWrapper {
25    /// Create a new MCP tool wrapper
26    pub fn new(server_name: String, mcp_tool: McpTool, manager: Arc<McpManager>) -> Self {
27        let full_name = format!("mcp__{}_{}", server_name, mcp_tool.name);
28        Self {
29            full_name,
30            mcp_tool,
31            server_name,
32            manager,
33        }
34    }
35
36    /// Get the server name
37    pub fn server_name(&self) -> &str {
38        &self.server_name
39    }
40
41    /// Get the original MCP tool name
42    pub fn mcp_tool_name(&self) -> &str {
43        &self.mcp_tool.name
44    }
45}
46
47#[async_trait]
48impl Tool for McpToolWrapper {
49    fn name(&self) -> &str {
50        &self.full_name
51    }
52
53    fn description(&self) -> &str {
54        self.mcp_tool.description.as_deref().unwrap_or("MCP tool")
55    }
56
57    fn parameters(&self) -> serde_json::Value {
58        self.mcp_tool.input_schema.clone()
59    }
60
61    async fn execute(&self, args: &serde_json::Value, _ctx: &ToolContext) -> Result<ToolOutput> {
62        // Call the MCP tool through the manager
63        let result = self
64            .manager
65            .call_tool(&self.full_name, Some(args.clone()))
66            .await;
67
68        match result {
69            Ok(tool_result) => {
70                let output = tool_result_to_string(&tool_result);
71                if tool_result.is_error {
72                    Ok(ToolOutput::error(output))
73                } else {
74                    Ok(ToolOutput::success(output))
75                }
76            }
77            Err(e) => Ok(ToolOutput::error(format!("MCP tool error: {}", e))),
78        }
79    }
80}
81
82/// Create tool wrappers for all tools from an MCP server
83pub fn create_mcp_tools(
84    server_name: &str,
85    tools: Vec<McpTool>,
86    manager: Arc<McpManager>,
87) -> Vec<Arc<dyn Tool>> {
88    tools
89        .into_iter()
90        .map(|tool| {
91            Arc::new(McpToolWrapper::new(
92                server_name.to_string(),
93                tool,
94                manager.clone(),
95            )) as Arc<dyn Tool>
96        })
97        .collect()
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_mcp_tool_wrapper_name() {
106        let manager = Arc::new(McpManager::new());
107        let mcp_tool = McpTool {
108            name: "create_issue".to_string(),
109            description: Some("Create a GitHub issue".to_string()),
110            input_schema: serde_json::json!({
111                "type": "object",
112                "properties": {
113                    "title": {"type": "string"}
114                }
115            }),
116        };
117
118        let wrapper = McpToolWrapper::new("github".to_string(), mcp_tool, manager);
119
120        assert_eq!(wrapper.name(), "mcp__github_create_issue");
121        assert_eq!(wrapper.server_name(), "github");
122        assert_eq!(wrapper.mcp_tool_name(), "create_issue");
123        assert_eq!(wrapper.description(), "Create a GitHub issue");
124    }
125
126    #[test]
127    fn test_create_mcp_tools() {
128        let manager = Arc::new(McpManager::new());
129        let tools = vec![
130            McpTool {
131                name: "tool1".to_string(),
132                description: Some("Tool 1".to_string()),
133                input_schema: serde_json::json!({}),
134            },
135            McpTool {
136                name: "tool2".to_string(),
137                description: Some("Tool 2".to_string()),
138                input_schema: serde_json::json!({}),
139            },
140        ];
141
142        let wrappers = create_mcp_tools("test", tools, manager);
143
144        assert_eq!(wrappers.len(), 2);
145        assert_eq!(wrappers[0].name(), "mcp__test_tool1");
146        assert_eq!(wrappers[1].name(), "mcp__test_tool2");
147    }
148}