bamboo_server_tools/
overlay_executor.rs

1use async_trait::async_trait;
2
3use bamboo_agent_core::tools::{
4    normalize_tool_name, parse_tool_args_best_effort, Tool, ToolCall, ToolCtx, ToolError,
5    ToolExecutionContext, ToolExecutor, ToolOutcome, ToolResult, ToolSchema,
6};
7use bamboo_tools::normalize_tool_ref;
8
9/// Tool executor that overlays a single tool on top of an existing executor.
10///
11/// This is used to add server-only tools (like `SubAgent`) without mutating the
12/// underlying built-in/MCP executor.
13pub struct OverlayToolExecutor {
14    base: std::sync::Arc<dyn ToolExecutor>,
15    overlay: std::sync::Arc<dyn Tool>,
16}
17
18impl OverlayToolExecutor {
19    pub fn new(base: std::sync::Arc<dyn ToolExecutor>, overlay: std::sync::Arc<dyn Tool>) -> Self {
20        Self { base, overlay }
21    }
22}
23
24#[async_trait]
25impl ToolExecutor for OverlayToolExecutor {
26    async fn execute(&self, call: &ToolCall) -> Result<ToolResult, ToolError> {
27        self.execute_with_context(call, ToolExecutionContext::none(&call.id))
28            .await
29    }
30
31    async fn execute_with_context(
32        &self,
33        call: &ToolCall,
34        ctx: ToolExecutionContext<'_>,
35    ) -> Result<ToolResult, ToolError> {
36        let name = normalize_tool_name(&call.function.name);
37        let is_overlay_call = name == self.overlay.name()
38            || normalize_tool_ref(name)
39                .as_deref()
40                .is_some_and(|normalized| normalized == self.overlay.name());
41        if is_overlay_call {
42            let args_raw = call.function.arguments.trim();
43            let (args, parse_warning) = parse_tool_args_best_effort(&call.function.arguments);
44            if let Some(warning) = parse_warning {
45                tracing::warn!(
46                    "Overlay tool argument parsing fallback applied: tool_call_id={}, tool_name={}, args_len={}, warning={}",
47                    call.id,
48                    call.function.name,
49                    args_raw.len(),
50                    warning
51                );
52            }
53            return self
54                .overlay
55                .invoke(args, ctx.to_tool_ctx())
56                .await
57                .map(|outcome| outcome.into_tool_result());
58        }
59        self.base.execute_with_context(call, ctx).await
60    }
61
62    async fn execute_with_context_outcome(
63        &self,
64        call: &ToolCall,
65        ctx: ToolExecutionContext<'_>,
66    ) -> Result<ToolOutcome, ToolError> {
67        let name = normalize_tool_name(&call.function.name);
68        let is_overlay_call = name == self.overlay.name()
69            || normalize_tool_ref(name)
70                .as_deref()
71                .is_some_and(|normalized| normalized == self.overlay.name());
72        if is_overlay_call {
73            let (args, _) = parse_tool_args_best_effort(&call.function.arguments);
74            return self.overlay.invoke(args, ctx.to_tool_ctx()).await;
75        }
76        self.base.execute_with_context_outcome(call, ctx).await
77    }
78
79    fn list_tools(&self) -> Vec<ToolSchema> {
80        let mut tools = self.base.list_tools();
81
82        // Ensure overlay tool is present exactly once.
83        let overlay_schema = self.overlay.to_schema();
84        let overlay_name = overlay_schema.function.name.clone();
85        tools.retain(|t| t.function.name != overlay_name);
86        tools.push(overlay_schema);
87
88        tools.sort_by_key(|t| t.function.name.clone());
89        tools
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    use serde_json::json;
98
99    use bamboo_agent_core::tools::FunctionCall;
100
101    struct BaseExecutor;
102
103    #[async_trait]
104    impl ToolExecutor for BaseExecutor {
105        async fn execute(&self, call: &ToolCall) -> Result<ToolResult, ToolError> {
106            Err(ToolError::Execution(format!(
107                "base executor called for {}",
108                call.function.name
109            )))
110        }
111
112        async fn execute_with_context(
113            &self,
114            call: &ToolCall,
115            _ctx: ToolExecutionContext<'_>,
116        ) -> Result<ToolResult, ToolError> {
117            self.execute(call).await
118        }
119
120        fn list_tools(&self) -> Vec<ToolSchema> {
121            Vec::new()
122        }
123    }
124
125    struct SubAgentOverlayTool;
126
127    #[async_trait]
128    impl Tool for SubAgentOverlayTool {
129        fn name(&self) -> &str {
130            "SubAgent"
131        }
132
133        fn description(&self) -> &str {
134            "overlay sub agent"
135        }
136
137        fn parameters_schema(&self) -> serde_json::Value {
138            json!({"type":"object","properties":{}})
139        }
140
141        async fn invoke(
142            &self,
143            _args: serde_json::Value,
144            _ctx: ToolCtx,
145        ) -> Result<ToolOutcome, ToolError> {
146            Ok(ToolOutcome::Completed(ToolResult {
147                success: true,
148                result: "overlay".to_string(),
149                display_preference: None,
150                images: Vec::new(),
151            }))
152        }
153    }
154
155    fn make_call(name: &str) -> ToolCall {
156        ToolCall {
157            id: "call_1".to_string(),
158            tool_type: "function".to_string(),
159            function: FunctionCall {
160                name: name.to_string(),
161                arguments: "{}".to_string(),
162            },
163        }
164    }
165
166    #[tokio::test]
167    async fn overlay_executor_routes_spawn_alias_to_overlay_tool() {
168        let overlay = OverlayToolExecutor::new(
169            std::sync::Arc::new(BaseExecutor),
170            std::sync::Arc::new(SubAgentOverlayTool),
171        );
172
173        let result = overlay
174            .execute(&make_call("sub_task"))
175            .await
176            .expect("spawn alias should route to overlay");
177
178        assert!(result.success);
179        assert_eq!(result.result, "overlay");
180    }
181
182    #[tokio::test]
183    async fn overlay_executor_keeps_non_overlay_calls_on_base_executor() {
184        let overlay = OverlayToolExecutor::new(
185            std::sync::Arc::new(BaseExecutor),
186            std::sync::Arc::new(SubAgentOverlayTool),
187        );
188
189        let err = overlay
190            .execute(&make_call("Read"))
191            .await
192            .expect_err("non-overlay call should stay on base executor");
193
194        assert!(
195            matches!(err, ToolError::Execution(msg) if msg.contains("base executor called for Read"))
196        );
197    }
198}
bamboo_server_tools/overlay_executor.rs

bamboo_server_tools/
overlay_executor.rs
