bamboo_agent/agent/mcp/
executor.rs

1use crate::agent::core::tools::{
2    ToolCall, ToolError, ToolExecutionContext, ToolExecutor, ToolResult, ToolSchema,
3};
4use async_trait::async_trait;
5use std::sync::Arc;
6use tracing::{debug, error};
7
8use crate::agent::mcp::error::McpError;
9use crate::agent::mcp::manager::McpServerManager;
10use crate::agent::mcp::tool_index::ToolIndex;
11use crate::agent::mcp::types::McpContentItem;
12
13/// MCP tool executor that delegates to the MCP server manager
14pub struct McpToolExecutor {
15    manager: Arc<McpServerManager>,
16    index: Arc<ToolIndex>,
17}
18
19impl McpToolExecutor {
20    pub fn new(manager: Arc<McpServerManager>, index: Arc<ToolIndex>) -> Self {
21        Self { manager, index }
22    }
23
24    /// Convert MCP result to string representation
25    fn format_result_content(content: &[McpContentItem]) -> String {
26        content
27            .iter()
28            .map(|item| match item {
29                McpContentItem::Text { text } => text.clone(),
30                McpContentItem::Image { data, mime_type } => {
31                    format!("[Image: {} ({} bytes)]", mime_type, data.len())
32                }
33                McpContentItem::Resource { resource } => {
34                    if let Some(text) = &resource.text {
35                        format!("[Resource {}]: {}", resource.uri, text)
36                    } else {
37                        format!("[Resource {}]", resource.uri)
38                    }
39                }
40            })
41            .collect::<Vec<_>>()
42            .join("\n")
43    }
44}
45
46#[async_trait]
47impl ToolExecutor for McpToolExecutor {
48    async fn execute(&self, call: &ToolCall) -> std::result::Result<ToolResult, ToolError> {
49        let tool_name = &call.function.name;
50
51        // Lookup the tool alias
52        let alias = match self.index.lookup(tool_name) {
53            Some(alias) => alias,
54            None => {
55                return Err(ToolError::NotFound(format!(
56                    "MCP tool '{}' not found",
57                    tool_name
58                )));
59            }
60        };
61
62        debug!(
63            "Executing MCP tool: {} (server: {}, original: {})",
64            tool_name, alias.server_id, alias.original_name
65        );
66
67        // Parse arguments
68        let args: serde_json::Value = serde_json::from_str(&call.function.arguments)
69            .map_err(|e| ToolError::InvalidArguments(format!("Invalid JSON: {}", e)))?;
70
71        // Execute via manager
72        match self
73            .manager
74            .call_tool(&alias.server_id, &alias.original_name, args)
75            .await
76        {
77            Ok(result) => {
78                if result.is_error {
79                    let error_text = Self::format_result_content(&result.content);
80                    Ok(ToolResult {
81                        success: false,
82                        result: error_text,
83                        display_preference: None,
84                    })
85                } else {
86                    let content = Self::format_result_content(&result.content);
87                    Ok(ToolResult {
88                        success: true,
89                        result: content,
90                        display_preference: None,
91                    })
92                }
93            }
94            Err(McpError::ServerNotFound(id)) => Err(ToolError::NotFound(format!(
95                "MCP server '{}' not found",
96                id
97            ))),
98            Err(McpError::ToolNotFound(name)) => {
99                Err(ToolError::NotFound(format!("Tool '{}' not found", name)))
100            }
101            Err(e) => {
102                error!("MCP tool execution failed: {}", e);
103                Err(ToolError::Execution(format!("MCP error: {}", e)))
104            }
105        }
106    }
107
108    fn list_tools(&self) -> Vec<ToolSchema> {
109        self.index
110            .all_aliases()
111            .into_iter()
112            .filter_map(|alias| {
113                // Get tool info from manager
114                self.manager
115                    .get_tool_info(&alias.server_id, &alias.original_name)
116                    .map(|tool| ToolSchema {
117                        schema_type: "function".to_string(),
118                        function: crate::agent::core::tools::FunctionSchema {
119                            name: alias.alias,
120                            description: tool.description,
121                            parameters: tool.parameters,
122                        },
123                    })
124            })
125            .collect()
126    }
127}
128
129/// Composite tool executor that tries built-in tools first, then MCP
130pub struct CompositeToolExecutor {
131    builtin: Arc<dyn ToolExecutor>,
132    mcp: Arc<dyn ToolExecutor>,
133}
134
135impl CompositeToolExecutor {
136    pub fn new(builtin: Arc<dyn ToolExecutor>, mcp: Arc<dyn ToolExecutor>) -> Self {
137        Self { builtin, mcp }
138    }
139}
140
141#[async_trait]
142impl ToolExecutor for CompositeToolExecutor {
143    async fn execute(&self, call: &ToolCall) -> std::result::Result<ToolResult, ToolError> {
144        // Try built-in first
145        match self.builtin.execute(call).await {
146            Ok(result) => return Ok(result),
147            Err(ToolError::NotFound(_)) => {
148                // Fall through to MCP
149            }
150            Err(e) => return Err(e),
151        }
152
153        // Try MCP
154        self.mcp.execute(call).await
155    }
156
157    async fn execute_with_context(
158        &self,
159        call: &ToolCall,
160        ctx: ToolExecutionContext<'_>,
161    ) -> std::result::Result<ToolResult, ToolError> {
162        // Try built-in first (preserve context for streaming tools).
163        match self.builtin.execute_with_context(call, ctx).await {
164            Ok(result) => return Ok(result),
165            Err(ToolError::NotFound(_)) => {
166                // Fall through to MCP
167            }
168            Err(e) => return Err(e),
169        }
170
171        // Try MCP (context ignored by default).
172        self.mcp.execute_with_context(call, ctx).await
173    }
174
175    fn list_tools(&self) -> Vec<ToolSchema> {
176        let mut tools = self.builtin.list_tools();
177        tools.extend(self.mcp.list_tools());
178        tools
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use crate::agent::core::tools::{FunctionCall, FunctionSchema};
186    use crate::agent::mcp::types::McpContentItem;
187    use mockall::mock;
188    use mockall::predicate::*;
189
190    // Mock McpTransport for testing
191    mock! {
192        pub ToolExecutor {}
193
194        #[async_trait]
195        impl ToolExecutor for ToolExecutor {
196            async fn execute(&self, call: &ToolCall) -> std::result::Result<ToolResult, ToolError>;
197            fn list_tools(&self) -> Vec<ToolSchema>;
198        }
199    }
200
201    fn create_test_tool_call(name: &str, args: &str) -> ToolCall {
202        ToolCall {
203            id: "test-id".to_string(),
204            tool_type: "function".to_string(),
205            function: FunctionCall {
206                name: name.to_string(),
207                arguments: args.to_string(),
208            },
209        }
210    }
211
212    #[test]
213    fn test_format_result_text() {
214        let content = vec![
215            McpContentItem::Text {
216                text: "Hello".to_string(),
217            },
218            McpContentItem::Text {
219                text: "World".to_string(),
220            },
221        ];
222        let result = McpToolExecutor::format_result_content(&content);
223        assert_eq!(result, "Hello\nWorld");
224    }
225
226    #[test]
227    fn test_format_result_image() {
228        let content = vec![McpContentItem::Image {
229            data: "base64imagedata".to_string(),
230            mime_type: "image/png".to_string(),
231        }];
232        let result = McpToolExecutor::format_result_content(&content);
233        assert_eq!(result, "[Image: image/png (15 bytes)]");
234    }
235
236    #[test]
237    fn test_format_result_resource_with_text() {
238        let content = vec![McpContentItem::Resource {
239            resource: crate::agent::mcp::types::McpResource {
240                uri: "file:///test.txt".to_string(),
241                mime_type: Some("text/plain".to_string()),
242                text: Some("File content".to_string()),
243                blob: None,
244            },
245        }];
246        let result = McpToolExecutor::format_result_content(&content);
247        assert_eq!(result, "[Resource file:///test.txt]: File content");
248    }
249
250    #[test]
251    fn test_format_result_resource_without_text() {
252        let content = vec![McpContentItem::Resource {
253            resource: crate::agent::mcp::types::McpResource {
254                uri: "file:///test.bin".to_string(),
255                mime_type: None,
256                text: None,
257                blob: Some("base64data".to_string()),
258            },
259        }];
260        let result = McpToolExecutor::format_result_content(&content);
261        assert_eq!(result, "[Resource file:///test.bin]");
262    }
263
264    #[test]
265    fn test_format_result_mixed() {
266        let content = vec![
267            McpContentItem::Text {
268                text: "Result:".to_string(),
269            },
270            McpContentItem::Image {
271                data: "img".to_string(),
272                mime_type: "image/png".to_string(),
273            },
274        ];
275        let result = McpToolExecutor::format_result_content(&content);
276        assert!(result.contains("Result:"));
277        assert!(result.contains("[Image:"));
278    }
279
280    #[tokio::test]
281    async fn test_composite_executor_fallback() {
282        let mut mock_builtin = MockToolExecutor::new();
283        let mut mock_mcp = MockToolExecutor::new();
284
285        // Built-in returns NotFound, so it should fall through to MCP
286        mock_builtin
287            .expect_execute()
288            .returning(|_| Err(ToolError::NotFound("not found".to_string())));
289
290        mock_mcp.expect_execute().returning(|_| {
291            Ok(ToolResult {
292                success: true,
293                result: "MCP result".to_string(),
294                display_preference: None,
295            })
296        });
297
298        mock_builtin.expect_list_tools().returning(|| vec![]);
299        mock_mcp.expect_list_tools().returning(|| vec![]);
300
301        let composite = CompositeToolExecutor::new(Arc::new(mock_builtin), Arc::new(mock_mcp));
302
303        let call = create_test_tool_call("test_tool", "{}");
304        let result = composite.execute(&call).await.unwrap();
305        assert!(result.success);
306        assert_eq!(result.result, "MCP result");
307    }
308
309    #[tokio::test]
310    async fn test_composite_executor_builtin_success() {
311        let mut mock_builtin = MockToolExecutor::new();
312        let mock_mcp = MockToolExecutor::new();
313
314        // Built-in succeeds, MCP should not be called
315        mock_builtin.expect_execute().returning(|_| {
316            Ok(ToolResult {
317                success: true,
318                result: "Built-in result".to_string(),
319                display_preference: None,
320            })
321        });
322
323        mock_builtin.expect_list_tools().returning(|| {
324            vec![ToolSchema {
325                schema_type: "function".to_string(),
326                function: FunctionSchema {
327                    name: "builtin_tool".to_string(),
328                    description: "A built-in tool".to_string(),
329                    parameters: serde_json::json!({}),
330                },
331            }]
332        });
333
334        let composite = CompositeToolExecutor::new(Arc::new(mock_builtin), Arc::new(mock_mcp));
335
336        let call = create_test_tool_call("test_tool", "{}");
337        let result = composite.execute(&call).await.unwrap();
338        assert!(result.success);
339        assert_eq!(result.result, "Built-in result");
340    }
341
342    #[tokio::test]
343    async fn test_composite_executor_builtin_error() {
344        let mut mock_builtin = MockToolExecutor::new();
345        let mock_mcp = MockToolExecutor::new();
346
347        // Built-in returns error (not NotFound), should propagate
348        mock_builtin
349            .expect_execute()
350            .returning(|_| Err(ToolError::Execution("Built-in error".to_string())));
351
352        mock_builtin.expect_list_tools().returning(|| {
353            vec![ToolSchema {
354                schema_type: "function".to_string(),
355                function: FunctionSchema {
356                    name: "builtin_tool".to_string(),
357                    description: "A built-in tool".to_string(),
358                    parameters: serde_json::json!({}),
359                },
360            }]
361        });
362
363        let composite = CompositeToolExecutor::new(Arc::new(mock_builtin), Arc::new(mock_mcp));
364
365        let call = create_test_tool_call("test_tool", "{}");
366        let result = composite.execute(&call).await;
367        assert!(result.is_err());
368        match result.unwrap_err() {
369            ToolError::Execution(msg) => assert_eq!(msg, "Built-in error"),
370            _ => panic!("Expected Execution error"),
371        }
372    }
373
374    #[test]
375    fn test_composite_list_tools() {
376        let mut mock_builtin = MockToolExecutor::new();
377        let mut mock_mcp = MockToolExecutor::new();
378
379        mock_builtin.expect_list_tools().returning(|| {
380            vec![ToolSchema {
381                schema_type: "function".to_string(),
382                function: FunctionSchema {
383                    name: "builtin_tool".to_string(),
384                    description: "Built-in tool".to_string(),
385                    parameters: serde_json::json!({}),
386                },
387            }]
388        });
389
390        mock_mcp.expect_list_tools().returning(|| {
391            vec![ToolSchema {
392                schema_type: "function".to_string(),
393                function: FunctionSchema {
394                    name: "mcp_tool".to_string(),
395                    description: "MCP tool".to_string(),
396                    parameters: serde_json::json!({}),
397                },
398            }]
399        });
400
401        let composite = CompositeToolExecutor::new(Arc::new(mock_builtin), Arc::new(mock_mcp));
402
403        let tools = composite.list_tools();
404        assert_eq!(tools.len(), 2);
405        assert_eq!(tools[0].function.name, "builtin_tool");
406        assert_eq!(tools[1].function.name, "mcp_tool");
407    }
408}
bamboo_agent/agent/mcp/executor.rs

bamboo_agent/agent/mcp/
executor.rs
