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

bamboo_engine/runtime/managers/adapters/
tool.rs
