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::McpManager;
6use crate::mcp::protocol::McpTool;
7use crate::mcp::result::project_tool_result;
8use crate::tools::{Tool, ToolContext, ToolOutput};
9use anyhow::Result;
10use async_trait::async_trait;
11use std::sync::Arc;
12
13/// MCP tool wrapper that implements the Tool trait
14pub struct McpToolWrapper {
15    /// Full tool name (mcp__server__tool)
16    full_name: String,
17    /// Original MCP tool definition
18    mcp_tool: McpTool,
19    /// Server name
20    server_name: String,
21    /// MCP manager reference
22    manager: Arc<McpManager>,
23}
24
25impl McpToolWrapper {
26    /// Create a new MCP tool wrapper
27    pub fn new(server_name: String, mcp_tool: McpTool, manager: Arc<McpManager>) -> Self {
28        let full_name = format!("mcp__{}__{}", server_name, mcp_tool.name);
29        Self {
30            full_name,
31            mcp_tool,
32            server_name,
33            manager,
34        }
35    }
36
37    /// Get the server name
38    pub fn server_name(&self) -> &str {
39        &self.server_name
40    }
41
42    /// Get the original MCP tool name
43    pub fn mcp_tool_name(&self) -> &str {
44        &self.mcp_tool.name
45    }
46}
47
48fn annotation_requires_confirmation(tool: &McpTool) -> bool {
49    let Some(annotations) = tool.annotations.as_ref() else {
50        // Missing behavior metadata is unknown, not read-only.
51        return true;
52    };
53
54    if annotations.destructive_hint == Some(true)
55        || annotations.read_only_hint != Some(true)
56        || annotations.open_world_hint != Some(false)
57    {
58        return true;
59    }
60
61    annotations
62        .additional
63        .iter()
64        .any(|(key, value)| custom_risk_requires_confirmation(key, value))
65        || tool.meta.as_ref().is_some_and(|meta| {
66            meta.as_object().is_some_and(|fields| {
67                fields
68                    .iter()
69                    .any(|(key, value)| custom_risk_requires_confirmation(key, value))
70            })
71        })
72}
73
74fn custom_risk_requires_confirmation(key: &str, value: &serde_json::Value) -> bool {
75    let key = key.to_ascii_lowercase();
76    if !matches!(
77        key.as_str(),
78        "x-a3s-risk" | "a3s/risk" | "a3s.risk" | "risk"
79    ) {
80        return false;
81    }
82
83    match value {
84        serde_json::Value::String(value) => !matches!(
85            value.trim().to_ascii_lowercase().as_str(),
86            "read" | "read_only" | "read-only" | "routine" | "closed_world_read"
87        ),
88        serde_json::Value::Array(values) => values
89            .iter()
90            .any(|value| custom_risk_requires_confirmation(key.as_str(), value)),
91        // A declared but malformed risk value cannot reduce confirmation.
92        _ => true,
93    }
94}
95
96#[async_trait]
97impl Tool for McpToolWrapper {
98    fn name(&self) -> &str {
99        &self.full_name
100    }
101
102    fn description(&self) -> &str {
103        self.mcp_tool.description.as_deref().unwrap_or("MCP tool")
104    }
105
106    fn parameters(&self) -> serde_json::Value {
107        self.mcp_tool.input_schema.clone()
108    }
109
110    fn requires_confirmation(&self, _args: &serde_json::Value) -> bool {
111        annotation_requires_confirmation(&self.mcp_tool)
112    }
113
114    async fn execute(&self, args: &serde_json::Value, ctx: &ToolContext) -> Result<ToolOutput> {
115        // Call the MCP tool through the manager
116        let result = self
117            .manager
118            .call_tool(&self.full_name, Some(args.clone()))
119            .await;
120
121        match result {
122            Ok(tool_result) => project_tool_result(&self.full_name, &tool_result, ctx).await,
123            Err(e) => Ok(ToolOutput::error(format!("MCP tool error: {}", e))),
124        }
125    }
126}
127
128/// Create tool wrappers for all tools from an MCP server
129pub fn create_mcp_tools(
130    server_name: &str,
131    tools: Vec<McpTool>,
132    manager: Arc<McpManager>,
133) -> Vec<Arc<dyn Tool>> {
134    tools
135        .into_iter()
136        .map(|tool| {
137            Arc::new(McpToolWrapper::new(
138                server_name.to_string(),
139                tool,
140                manager.clone(),
141            )) as Arc<dyn Tool>
142        })
143        .collect()
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149    use crate::mcp::protocol::McpToolAnnotations;
150    use crate::tools::Tool;
151    use std::collections::HashMap;
152
153    #[test]
154    fn test_mcp_tool_wrapper_name() {
155        let manager = Arc::new(McpManager::new());
156        let mcp_tool = McpTool {
157            name: "create_issue".to_string(),
158            title: None,
159            description: Some("Create a GitHub issue".to_string()),
160            input_schema: serde_json::json!({
161                "type": "object",
162                "properties": {
163                    "title": {"type": "string"}
164                }
165            }),
166            output_schema: None,
167            annotations: None,
168            icons: Vec::new(),
169            meta: None,
170        };
171
172        let wrapper = McpToolWrapper::new("github".to_string(), mcp_tool, manager);
173
174        assert_eq!(wrapper.name(), "mcp__github__create_issue");
175        assert_eq!(wrapper.server_name(), "github");
176        assert_eq!(wrapper.mcp_tool_name(), "create_issue");
177        assert_eq!(wrapper.description(), "Create a GitHub issue");
178    }
179
180    #[test]
181    fn test_create_mcp_tools() {
182        let manager = Arc::new(McpManager::new());
183        let tools = vec![
184            McpTool {
185                name: "tool1".to_string(),
186                title: None,
187                description: Some("Tool 1".to_string()),
188                input_schema: serde_json::json!({}),
189                output_schema: None,
190                annotations: None,
191                icons: Vec::new(),
192                meta: None,
193            },
194            McpTool {
195                name: "tool2".to_string(),
196                title: None,
197                description: Some("Tool 2".to_string()),
198                input_schema: serde_json::json!({}),
199                output_schema: None,
200                annotations: None,
201                icons: Vec::new(),
202                meta: None,
203            },
204        ];
205
206        let wrappers = create_mcp_tools("test", tools, manager);
207
208        assert_eq!(wrappers.len(), 2);
209        assert_eq!(wrappers[0].name(), "mcp__test__tool1");
210        assert_eq!(wrappers[1].name(), "mcp__test__tool2");
211    }
212
213    fn annotated_tool(annotations: Option<McpToolAnnotations>) -> McpTool {
214        McpTool {
215            name: "fixture".to_string(),
216            title: None,
217            description: None,
218            input_schema: serde_json::json!({"type": "object"}),
219            output_schema: None,
220            annotations,
221            icons: Vec::new(),
222            meta: None,
223        }
224    }
225
226    #[test]
227    fn closed_world_read_only_annotation_does_not_escalate_confirmation() {
228        let manager = Arc::new(McpManager::new());
229        let wrapper = McpToolWrapper::new(
230            "use_fixture".to_string(),
231            annotated_tool(Some(McpToolAnnotations {
232                read_only_hint: Some(true),
233                destructive_hint: Some(false),
234                idempotent_hint: Some(true),
235                open_world_hint: Some(false),
236                ..Default::default()
237            })),
238            manager,
239        );
240
241        assert!(!wrapper.requires_confirmation(&serde_json::json!({})));
242    }
243
244    #[test]
245    fn unknown_open_world_mutating_and_submit_tools_escalate_confirmation() {
246        let cases = [
247            None,
248            Some(McpToolAnnotations {
249                read_only_hint: Some(true),
250                open_world_hint: Some(true),
251                ..Default::default()
252            }),
253            Some(McpToolAnnotations {
254                read_only_hint: Some(false),
255                open_world_hint: Some(false),
256                ..Default::default()
257            }),
258            Some(McpToolAnnotations {
259                read_only_hint: Some(true),
260                open_world_hint: Some(false),
261                additional: HashMap::from([(
262                    "x-a3s-risk".to_string(),
263                    serde_json::json!("submit"),
264                )]),
265                ..Default::default()
266            }),
267        ];
268
269        for annotations in cases {
270            let wrapper = McpToolWrapper::new(
271                "use_fixture".to_string(),
272                annotated_tool(annotations),
273                Arc::new(McpManager::new()),
274            );
275            assert!(wrapper.requires_confirmation(&serde_json::json!({})));
276        }
277    }
278}