bamboo_engine/runtime/managers/adapters/
tool.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use bamboo_agent_core::tools::{ToolCall, ToolExecutor, ToolSchema};
5use bamboo_agent_core::{AgentError, AgentEvent, Session};
6use bamboo_llm::LLMProvider;
7use bamboo_metrics::MetricsCollector;
8use tokio::sync::mpsc;
9use tokio_util::sync::CancellationToken;
10
11use crate::runtime::config::AgentLoopConfig;
12use crate::runtime::managers::tool::{ToolManager, ToolRoundResult};
13use crate::runtime::task_context::TaskLoopContext;
14
15/// Default tool manager that delegates to existing runner functions.
16pub struct DefaultToolManager {
17    tools: Arc<dyn ToolExecutor>,
18    llm: Arc<dyn LLMProvider>,
19}
20
21impl DefaultToolManager {
22    pub fn new(tools: Arc<dyn ToolExecutor>, llm: Arc<dyn LLMProvider>) -> Self {
23        Self { tools, llm }
24    }
25}
26
27#[async_trait]
28impl ToolManager for DefaultToolManager {
29    fn resolve_tool_schemas(&self, config: &AgentLoopConfig, session: &Session) -> Vec<ToolSchema> {
30        crate::runtime::runner::session_setup::tool_schemas::resolve_available_tool_schemas_for_session(
31            config,
32            self.tools.as_ref(),
33            session,
34        )
35    }
36
37    #[allow(clippy::too_many_arguments)]
38    async fn execute_tool_calls(
39        &self,
40        tool_calls: &[ToolCall],
41        event_tx: &mpsc::Sender<AgentEvent>,
42        metrics_collector: Option<&MetricsCollector>,
43        session_id: &str,
44        round_id: &str,
45        round: usize,
46        session: &mut Session,
47        config: &AgentLoopConfig,
48        task_context: &mut Option<TaskLoopContext>,
49        tool_schemas: &[ToolSchema],
50        cancel: &CancellationToken,
51    ) -> Result<ToolRoundResult, AgentError> {
52        let frame = crate::runtime::runner::round_frame::RoundFrame {
53            session_id,
54            round_id,
55            turn: round,
56            debug_enabled: false,
57            event_tx,
58            metrics_collector,
59            config,
60            llm: &self.llm,
61            tools: &self.tools,
62        };
63
64        // Mirror the live pipeline's #30 biased-cancel wrap so a cancel issued
65        // DURING tool execution (e.g. a long foreground Bash run) is honored on
66        // this adapter path too, not only between rounds. `biased` checks
67        // cancellation first; on cancel the in-flight tool futures are dropped
68        // (true cancellation — foreground Bash is kill_on_drop). The per-tool
69        // `tokio::time::timeout` inside `execute_round_tool_calls` is preserved;
70        // cancel is strictly an additional early-exit. #104.
71        let result = tokio::select! {
72            biased;
73            _ = cancel.cancelled() => return Err(AgentError::Cancelled),
74            result = crate::runtime::runner::tool_execution::execute_round_tool_calls(
75                tool_calls,
76                &frame,
77                session,
78                task_context,
79                config
80                    .summarization_model_name
81                    .as_deref()
82                    .or(config.background_model_name.as_deref()),
83                config
84                    .summarization_model_provider
85                    .as_ref()
86                    .or(config.background_model_provider.as_ref()),
87                tool_schemas,
88            ) => result?,
89        };
90
91        Ok(ToolRoundResult {
92            awaiting_clarification: result.awaiting_clarification,
93            should_break: false,
94            tool_calls_count: tool_calls.len(),
95        })
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use bamboo_agent_core::tools::{FunctionCall, ToolError, ToolExecutionContext, ToolResult};
103    use bamboo_agent_core::Message;
104    use bamboo_llm::provider::LLMStream;
105
106    // Stubs that PANIC if invoked — they prove the biased cancel arm returns BEFORE
107    // any LLM or tool work runs.
108    struct PanicProvider;
109    #[async_trait]
110    impl bamboo_llm::LLMProvider for PanicProvider {
111        async fn chat_stream(
112            &self,
113            _messages: &[Message],
114            _tools: &[ToolSchema],
115            _max_output_tokens: Option<u32>,
116            _model: &str,
117        ) -> bamboo_llm::provider::Result<LLMStream> {
118            panic!("LLM must not be invoked when the run is already cancelled");
119        }
120    }
121
122    struct PanicExecutor;
123    #[async_trait]
124    impl ToolExecutor for PanicExecutor {
125        async fn execute(&self, _call: &ToolCall) -> Result<ToolResult, ToolError> {
126            panic!("tools must not run when the run is already cancelled");
127        }
128        async fn execute_with_context(
129            &self,
130            call: &ToolCall,
131            _ctx: ToolExecutionContext<'_>,
132        ) -> Result<ToolResult, ToolError> {
133            self.execute(call).await
134        }
135        fn list_tools(&self) -> Vec<ToolSchema> {
136            Vec::new()
137        }
138    }
139
140    #[tokio::test]
141    async fn execute_tool_calls_short_circuits_to_cancelled_when_token_already_fired() {
142        let mgr = DefaultToolManager::new(Arc::new(PanicExecutor), Arc::new(PanicProvider));
143        let (event_tx, _rx) = mpsc::channel(8);
144        let mut session = Session::new("s1", "model");
145        let config = AgentLoopConfig::default();
146        let mut task_context = None;
147        // A non-empty tool call: without the cancel short-circuit, execute_round
148        // would route it to PanicExecutor and the test would panic.
149        let tool_calls = vec![ToolCall {
150            id: "c1".to_string(),
151            tool_type: "function".to_string(),
152            function: FunctionCall {
153                name: "anything".to_string(),
154                arguments: "{}".to_string(),
155            },
156        }];
157
158        let cancel = CancellationToken::new();
159        cancel.cancel(); // already cancelled before the call
160
161        let result = mgr
162            .execute_tool_calls(
163                &tool_calls,
164                &event_tx,
165                None,
166                "s1",
167                "r1",
168                0,
169                &mut session,
170                &config,
171                &mut task_context,
172                &[],
173                &cancel,
174            )
175            .await;
176
177        assert!(
178            matches!(result, Err(AgentError::Cancelled)),
179            "a pre-cancelled token returns Cancelled before touching tools/LLM; got {result:?}"
180        );
181    }
182}
bamboo_engine/runtime/managers/adapters/tool.rs

bamboo_engine/runtime/managers/adapters/
tool.rs
