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

bamboo_engine/runtime/managers/adapters/
tool.rs
