bamboo_engine/mcp/
executor.rs

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

bamboo_engine/mcp/
executor.rs
