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        let mut runtime_state = session
64            .agent_runtime_state
65            .clone()
66            .unwrap_or_else(|| bamboo_domain::AgentRuntimeState::new(session_id));
67
68        // Mirror the live pipeline's #30 biased-cancel wrap so a cancel issued
69        // DURING tool execution (e.g. a long foreground Bash run) is honored on
70        // this adapter path too, not only between rounds. `biased` checks
71        // cancellation first; on cancel the in-flight tool futures are dropped
72        // (true cancellation — foreground Bash is kill_on_drop). The per-tool
73        // `tokio::time::timeout` inside `execute_round_tool_calls` is preserved;
74        // cancel is strictly an additional early-exit. #104.
75        let result = tokio::select! {
76            biased;
77            _ = cancel.cancelled() => return Err(AgentError::Cancelled),
78            result = crate::runtime::runner::tool_execution::execute_round_tool_calls(
79                tool_calls,
80                &frame,
81                session,
82                &mut runtime_state,
83                task_context,
84                config
85                    .summarization_model_name
86                    .as_deref()
87                    .or(config.background_model_name.as_deref()),
88                config
89                    .summarization_model_provider
90                    .as_ref()
91                    .or(config.background_model_provider.as_ref()),
92                tool_schemas,
93            ) => result?,
94        };
95        if !config.hook_runner.is_empty() {
96            session.agent_runtime_state = Some(runtime_state);
97        }
98
99        Ok(ToolRoundResult {
100            awaiting_clarification: result.awaiting_clarification,
101            should_break: false,
102            tool_calls_count: tool_calls.len(),
103        })
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110    use bamboo_agent_core::tools::{FunctionCall, ToolError, ToolExecutionContext, ToolResult};
111    use bamboo_agent_core::Message;
112    use bamboo_llm::provider::LLMStream;
113
114    // Stubs that PANIC if invoked — they prove the biased cancel arm returns BEFORE
115    // any LLM or tool work runs.
116    struct PanicProvider;
117    #[async_trait]
118    impl bamboo_llm::LLMProvider for PanicProvider {
119        async fn chat_stream(
120            &self,
121            _messages: &[Message],
122            _tools: &[ToolSchema],
123            _max_output_tokens: Option<u32>,
124            _model: &str,
125        ) -> bamboo_llm::provider::Result<LLMStream> {
126            panic!("LLM must not be invoked when the run is already cancelled");
127        }
128    }
129
130    struct PanicExecutor;
131    #[async_trait]
132    impl ToolExecutor for PanicExecutor {
133        async fn execute(&self, _call: &ToolCall) -> Result<ToolResult, ToolError> {
134            panic!("tools must not run when the run is already cancelled");
135        }
136        async fn execute_with_context(
137            &self,
138            call: &ToolCall,
139            _ctx: ToolExecutionContext<'_>,
140        ) -> Result<ToolResult, ToolError> {
141            self.execute(call).await
142        }
143        fn list_tools(&self) -> Vec<ToolSchema> {
144            Vec::new()
145        }
146    }
147
148    #[tokio::test]
149    async fn execute_tool_calls_short_circuits_to_cancelled_when_token_already_fired() {
150        let mgr = DefaultToolManager::new(Arc::new(PanicExecutor), Arc::new(PanicProvider));
151        let (event_tx, _rx) = mpsc::channel(8);
152        let mut session = Session::new("s1", "model");
153        let config = AgentLoopConfig::default();
154        let mut task_context = None;
155        // A non-empty tool call: without the cancel short-circuit, execute_round
156        // would route it to PanicExecutor and the test would panic.
157        let tool_calls = vec![ToolCall {
158            id: "c1".to_string(),
159            tool_type: "function".to_string(),
160            function: FunctionCall {
161                name: "anything".to_string(),
162                arguments: "{}".to_string(),
163            },
164        }];
165
166        let cancel = CancellationToken::new();
167        cancel.cancel(); // already cancelled before the call
168
169        let result = mgr
170            .execute_tool_calls(
171                &tool_calls,
172                &event_tx,
173                None,
174                "s1",
175                "r1",
176                0,
177                &mut session,
178                &config,
179                &mut task_context,
180                &[],
181                &cancel,
182            )
183            .await;
184
185        assert!(
186            matches!(result, Err(AgentError::Cancelled)),
187            "a pre-cancelled token returns Cancelled before touching tools/LLM; got {result:?}"
188        );
189    }
190}
bamboo_engine/runtime/managers/adapters/tool.rs

bamboo_engine/runtime/managers/adapters/
tool.rs
