Skip to main content

bamboo_agent/agent/core/tools/
result_handler.rs

1use std::collections::VecDeque;
2use std::sync::Arc;
3
4use tokio::sync::mpsc;
5
6use crate::agent::core::composition::CompositionExecutor;
7use crate::agent::core::tools::executor::execute_tool_call_with_context;
8use crate::agent::core::tools::{
9    convert_from_standard_result, AgenticToolResult, ToolCall, ToolError, ToolExecutionContext,
10    ToolExecutor, ToolResult,
11};
12use crate::agent::core::{AgentEvent, Message, Session};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum ToolHandlingOutcome {
16    Continue,
17    AwaitingClarification,
18}
19
20pub const MAX_SUB_ACTIONS: usize = 64;
21
22pub fn parse_tool_args(arguments: &str) -> std::result::Result<serde_json::Value, ToolError> {
23    let args_raw = arguments.trim();
24
25    if args_raw.is_empty() {
26        return Ok(serde_json::json!({}));
27    }
28
29    serde_json::from_str(args_raw)
30        .map_err(|error| ToolError::InvalidArguments(format!("Invalid JSON arguments: {error}")))
31}
32
33pub fn try_parse_agentic_result(result: &ToolResult) -> Option<AgenticToolResult> {
34    if result.result.trim_start().starts_with('{') {
35        if let Ok(parsed) = serde_json::from_str::<AgenticToolResult>(&result.result) {
36            return Some(parsed);
37        }
38    }
39
40    match result.display_preference.as_deref() {
41        Some("clarification") | Some("actions_needed") => {
42            Some(convert_from_standard_result(result.clone()))
43        }
44        _ => None,
45    }
46}
47
48pub async fn handle_tool_result_with_agentic_support(
49    result: &ToolResult,
50    tool_call: &ToolCall,
51    event_tx: &mpsc::Sender<AgentEvent>,
52    session: &mut Session,
53    tools: &dyn ToolExecutor,
54    composition_executor: Option<Arc<CompositionExecutor>>,
55) -> ToolHandlingOutcome {
56    let Some(agentic_result) = try_parse_agentic_result(result) else {
57        session.add_message(Message::tool_result(
58            tool_call.id.clone(),
59            result.result.clone(),
60        ));
61        return ToolHandlingOutcome::Continue;
62    };
63
64    match agentic_result {
65        AgenticToolResult::Success { result } => {
66            session.add_message(Message::tool_result(tool_call.id.clone(), result));
67            ToolHandlingOutcome::Continue
68        }
69        AgenticToolResult::Error { error } => {
70            let _ = event_tx
71                .send(AgentEvent::ToolError {
72                    tool_call_id: tool_call.id.clone(),
73                    error: error.clone(),
74                })
75                .await;
76
77            session.add_message(Message::tool_result(
78                tool_call.id.clone(),
79                format!("Error: {error}"),
80            ));
81
82            ToolHandlingOutcome::Continue
83        }
84        AgenticToolResult::NeedClarification { question, options } => {
85            send_clarification_request(event_tx, question.clone(), options).await;
86
87            session.add_message(Message::tool_result(
88                tool_call.id.clone(),
89                format!("Clarification needed: {question}"),
90            ));
91
92            ToolHandlingOutcome::AwaitingClarification
93        }
94        AgenticToolResult::NeedMoreActions { actions, reason } => {
95            session.add_message(Message::tool_result(
96                tool_call.id.clone(),
97                format!(
98                    "Need more actions: {reason} ({} actions pending)",
99                    actions.len()
100                ),
101            ));
102
103            execute_sub_actions(&actions, event_tx, session, tools, composition_executor).await
104        }
105    }
106}
107
108pub async fn send_clarification_request(
109    event_tx: &mpsc::Sender<AgentEvent>,
110    question: String,
111    options: Option<Vec<String>>,
112) {
113    let _ = event_tx
114        .send(AgentEvent::NeedClarification { question, options })
115        .await;
116}
117
118pub async fn execute_sub_actions(
119    actions: &[ToolCall],
120    event_tx: &mpsc::Sender<AgentEvent>,
121    session: &mut Session,
122    tools: &dyn ToolExecutor,
123    composition_executor: Option<Arc<CompositionExecutor>>,
124) -> ToolHandlingOutcome {
125    let mut pending: VecDeque<ToolCall> = actions.iter().cloned().collect();
126    let mut processed = 0usize;
127
128    while let Some(action) = pending.pop_front() {
129        if processed >= MAX_SUB_ACTIONS {
130            let error = format!("Reached max sub-action limit ({MAX_SUB_ACTIONS})");
131            let _ = event_tx
132                .send(AgentEvent::ToolError {
133                    tool_call_id: action.id.clone(),
134                    error: error.clone(),
135                })
136                .await;
137            session.add_message(Message::tool_result(action.id.clone(), error));
138            return ToolHandlingOutcome::Continue;
139        }
140
141        processed += 1;
142
143        let args =
144            parse_tool_args(&action.function.arguments).unwrap_or_else(|_| serde_json::json!({}));
145
146        let _ = event_tx
147            .send(AgentEvent::ToolStart {
148                tool_call_id: action.id.clone(),
149                tool_name: action.function.name.clone(),
150                arguments: args,
151            })
152            .await;
153
154        let tool_ctx = ToolExecutionContext {
155            session_id: Some(&session.id),
156            tool_call_id: &action.id,
157            event_tx: Some(event_tx),
158        };
159
160        match execute_tool_call_with_context(&action, tools, composition_executor.clone(), tool_ctx)
161            .await
162        {
163            Ok(result) => {
164                let _ = event_tx
165                    .send(AgentEvent::ToolComplete {
166                        tool_call_id: action.id.clone(),
167                        result: result.clone(),
168                    })
169                    .await;
170
171                match try_parse_agentic_result(&result) {
172                    Some(AgenticToolResult::Success { result }) => {
173                        session.add_message(Message::tool_result(action.id.clone(), result));
174                    }
175                    Some(AgenticToolResult::Error { error }) => {
176                        let _ = event_tx
177                            .send(AgentEvent::ToolError {
178                                tool_call_id: action.id.clone(),
179                                error: error.clone(),
180                            })
181                            .await;
182                        session.add_message(Message::tool_result(
183                            action.id.clone(),
184                            format!("Error: {error}"),
185                        ));
186                    }
187                    Some(AgenticToolResult::NeedClarification { question, options }) => {
188                        send_clarification_request(event_tx, question.clone(), options).await;
189                        session.add_message(Message::tool_result(
190                            action.id.clone(),
191                            format!("Clarification needed: {question}"),
192                        ));
193                        return ToolHandlingOutcome::AwaitingClarification;
194                    }
195                    Some(AgenticToolResult::NeedMoreActions {
196                        actions: next_actions,
197                        reason,
198                    }) => {
199                        session.add_message(Message::tool_result(
200                            action.id.clone(),
201                            format!(
202                                "Need more actions: {reason} ({} actions pending)",
203                                next_actions.len()
204                            ),
205                        ));
206                        pending.extend(next_actions);
207                    }
208                    None => {
209                        session.add_message(Message::tool_result(
210                            action.id.clone(),
211                            result.result.clone(),
212                        ));
213                    }
214                }
215            }
216            Err(error) => {
217                let error_msg = error.to_string();
218                let _ = event_tx
219                    .send(AgentEvent::ToolError {
220                        tool_call_id: action.id.clone(),
221                        error: error_msg.clone(),
222                    })
223                    .await;
224                session.add_message(Message::tool_result(
225                    action.id.clone(),
226                    format!("Error: {error_msg}"),
227                ));
228            }
229        }
230    }
231
232    ToolHandlingOutcome::Continue
233}
234
235#[cfg(test)]
236mod tests {
237    use async_trait::async_trait;
238    use std::collections::HashMap;
239    use std::sync::Arc;
240    use tokio::sync::mpsc;
241
242    use crate::agent::core::tools::{FunctionCall, ToolSchema};
243
244    use super::*;
245
246    struct StaticExecutor {
247        results: HashMap<String, ToolResult>,
248    }
249
250    impl StaticExecutor {
251        fn new(results: HashMap<String, ToolResult>) -> Self {
252            Self { results }
253        }
254    }
255
256    #[async_trait]
257    impl ToolExecutor for StaticExecutor {
258        async fn execute(
259            &self,
260            call: &ToolCall,
261        ) -> crate::agent::core::tools::executor::Result<ToolResult> {
262            self.results
263                .get(&call.function.name)
264                .cloned()
265                .ok_or_else(|| ToolError::NotFound(call.function.name.clone()))
266        }
267
268        fn list_tools(&self) -> Vec<ToolSchema> {
269            Vec::new()
270        }
271    }
272
273    fn make_tool_call(id: &str, name: &str, arguments: &str) -> ToolCall {
274        ToolCall {
275            id: id.to_string(),
276            tool_type: "function".to_string(),
277            function: FunctionCall {
278                name: name.to_string(),
279                arguments: arguments.to_string(),
280            },
281        }
282    }
283
284    #[tokio::test]
285    async fn need_clarification_sends_event() {
286        let (event_tx, mut event_rx) = mpsc::channel(8);
287        let tools: Arc<dyn ToolExecutor> = Arc::new(StaticExecutor::new(HashMap::new()));
288        let mut session = Session::new("s1", "test-model");
289        let tool_call = make_tool_call("call_parent", "smart_tool", "{}");
290
291        let result = ToolResult {
292            success: true,
293            result: serde_json::to_string(&AgenticToolResult::NeedClarification {
294                question: "Which file should I inspect?".to_string(),
295                options: Some(vec!["src/main.rs".to_string(), "src/lib.rs".to_string()]),
296            })
297            .unwrap(),
298            display_preference: None,
299        };
300
301        let outcome = handle_tool_result_with_agentic_support(
302            &result,
303            &tool_call,
304            &event_tx,
305            &mut session,
306            tools.as_ref(),
307            None,
308        )
309        .await;
310
311        assert_eq!(outcome, ToolHandlingOutcome::AwaitingClarification);
312
313        let event = event_rx.recv().await.expect("missing clarification event");
314        match event {
315            AgentEvent::NeedClarification { question, options } => {
316                assert_eq!(question, "Which file should I inspect?");
317                assert_eq!(
318                    options,
319                    Some(vec!["src/main.rs".to_string(), "src/lib.rs".to_string()])
320                );
321            }
322            other => panic!("unexpected event: {other:?}"),
323        }
324    }
325
326    #[tokio::test]
327    async fn need_more_actions_executes_sub_actions() {
328        let (event_tx, mut event_rx) = mpsc::channel(16);
329        let sub_action = make_tool_call("call_sub", "sub_tool", "{}");
330        let parent_call = make_tool_call("call_parent", "smart_tool", "{}");
331
332        let mut results = HashMap::new();
333        results.insert(
334            "sub_tool".to_string(),
335            ToolResult {
336                success: true,
337                result: "sub-action-done".to_string(),
338                display_preference: None,
339            },
340        );
341        let tools: Arc<dyn ToolExecutor> = Arc::new(StaticExecutor::new(results));
342        let mut session = Session::new("s2", "test-model");
343
344        let result = ToolResult {
345            success: true,
346            result: serde_json::to_string(&AgenticToolResult::NeedMoreActions {
347                actions: vec![sub_action],
348                reason: "Need workspace context".to_string(),
349            })
350            .unwrap(),
351            display_preference: None,
352        };
353
354        let outcome = handle_tool_result_with_agentic_support(
355            &result,
356            &parent_call,
357            &event_tx,
358            &mut session,
359            tools.as_ref(),
360            None,
361        )
362        .await;
363
364        assert_eq!(outcome, ToolHandlingOutcome::Continue);
365        assert!(session
366            .messages
367            .iter()
368            .any(
369                |message| message.tool_call_id.as_deref() == Some("call_sub")
370                    && message.content == "sub-action-done"
371            ));
372
373        let mut saw_sub_start = false;
374        let mut saw_sub_complete = false;
375
376        while let Ok(event) = event_rx.try_recv() {
377            match event {
378                AgentEvent::ToolStart { tool_call_id, .. } if tool_call_id == "call_sub" => {
379                    saw_sub_start = true;
380                }
381                AgentEvent::ToolComplete { tool_call_id, .. } if tool_call_id == "call_sub" => {
382                    saw_sub_complete = true;
383                }
384                _ => {}
385            }
386        }
387
388        assert!(saw_sub_start);
389        assert!(saw_sub_complete);
390    }
391
392    #[test]
393    fn parse_tool_args_rejects_invalid_json() {
394        let error = parse_tool_args("not-json").expect_err("invalid json should fail");
395        assert!(matches!(error, ToolError::InvalidArguments(_)));
396    }
397}