Skip to main content

recursive/
agent.rs

1//! Agent loop. The whole kernel.
2//!
3//! Tiny on purpose: receive a goal, ask the model what to do, run any tool
4//! calls, feed results back, repeat until the model stops requesting tools
5//! or we hit the step budget. Everything interesting (which model, which
6//! tools, what system prompt) is injected by the caller.
7//!
8//! The loop emits `StepEvent`s through a channel so a UI/CLI/log layer can
9//! observe progress without coupling to the agent's internals.
10
11use std::sync::Arc;
12use tokio::sync::mpsc;
13use tracing::{debug, info, warn};
14
15use crate::error::{Error, Result};
16use crate::llm::{Completion, LlmProvider, ToolCall};
17use crate::message::Message;
18use crate::tools::ToolRegistry;
19
20/// Threshold for consecutive identical failing tool calls before declaring stuck.
21const STUCK_THRESHOLD: usize = 3;
22
23#[derive(Debug, Clone)]
24pub enum StepEvent {
25    AssistantText {
26        text: String,
27        step: usize,
28    },
29    ToolCall {
30        call: ToolCall,
31        step: usize,
32    },
33    ToolResult {
34        id: String,
35        name: String,
36        output: String,
37        step: usize,
38    },
39    Finished {
40        reason: FinishReason,
41        steps: usize,
42    },
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum FinishReason {
47    NoMoreToolCalls,
48    BudgetExceeded,
49    ProviderStop(String),
50    Stuck {
51        repeated_call: String,
52        repeats: usize,
53    },
54}
55
56#[derive(Debug, Clone)]
57pub struct AgentOutcome {
58    pub final_message: Option<String>,
59    pub transcript: Vec<Message>,
60    pub steps: usize,
61    pub finish: FinishReason,
62}
63
64pub struct Agent {
65    llm: Arc<dyn LlmProvider>,
66    tools: ToolRegistry,
67    transcript: Vec<Message>,
68    max_steps: usize,
69    events: Option<mpsc::UnboundedSender<StepEvent>>,
70}
71
72impl Agent {
73    pub fn builder() -> AgentBuilder {
74        AgentBuilder::default()
75    }
76
77    /// Drive the loop until the model stops calling tools, or the budget is exhausted.
78    pub async fn run(&mut self, goal: impl Into<String>) -> Result<AgentOutcome> {
79        let goal = goal.into();
80        info!(target: "recursive::agent", goal = %truncate(&goal, 200), "agent run starting");
81        self.transcript.push(Message::user(goal));
82
83        let mut final_message: Option<String> = None;
84        let specs = self.tools.specs();
85
86        // Tracking for anti-stuck heuristic
87        let mut last_call_key: Option<(String, String)> = None;
88        let mut consecutive_errors: usize = 0;
89
90        for step in 1..=self.max_steps {
91            debug!(target: "recursive::agent", step, "calling llm");
92            let completion: Completion = self.llm.complete(&self.transcript, &specs).await?;
93
94            if !completion.content.is_empty() {
95                self.emit(StepEvent::AssistantText {
96                    text: completion.content.clone(),
97                    step,
98                });
99                final_message = Some(completion.content.clone());
100            }
101
102            if completion.tool_calls.is_empty() {
103                self.transcript
104                    .push(Message::assistant(completion.content.clone()));
105                let finish = match completion.finish_reason {
106                    Some(r) if r != "stop" && r != "end_turn" => FinishReason::ProviderStop(r),
107                    _ => FinishReason::NoMoreToolCalls,
108                };
109                self.emit(StepEvent::Finished {
110                    reason: finish.clone(),
111                    steps: step,
112                });
113                return Ok(AgentOutcome {
114                    final_message,
115                    transcript: std::mem::take(&mut self.transcript),
116                    steps: step,
117                    finish,
118                });
119            }
120
121            self.transcript.push(Message::assistant_with_tool_calls(
122                completion.content.clone(),
123                completion.tool_calls.clone(),
124            ));
125
126            for call in completion.tool_calls.iter() {
127                self.emit(StepEvent::ToolCall {
128                    call: call.clone(),
129                    step,
130                });
131                let result = match self.tools.invoke(&call.name, call.arguments.clone()).await {
132                    Ok(output) => output,
133                    Err(err) => format!("ERROR: {err}"),
134                };
135
136                // Anti-stuck heuristic: track identical failing calls
137                let call_key = (
138                    call.name.clone(),
139                    serde_json::to_string(&call.arguments).unwrap_or_default(),
140                );
141                let is_error = result.starts_with("ERROR:");
142
143                if is_error {
144                    if last_call_key == Some(call_key.clone()) {
145                        consecutive_errors += 1;
146                    } else {
147                        consecutive_errors = 1;
148                    }
149                } else {
150                    consecutive_errors = 0;
151                }
152
153                last_call_key = Some(call_key);
154
155                // Check if stuck threshold reached
156                if consecutive_errors >= STUCK_THRESHOLD {
157                    let repeated_call = call.name.clone();
158                    let repeats = consecutive_errors;
159                    let finish = FinishReason::Stuck {
160                        repeated_call,
161                        repeats,
162                    };
163                    self.emit(StepEvent::Finished {
164                        reason: finish.clone(),
165                        steps: step,
166                    });
167                    return Ok(AgentOutcome {
168                        final_message,
169                        transcript: std::mem::take(&mut self.transcript),
170                        steps: step,
171                        finish,
172                    });
173                }
174
175                self.emit(StepEvent::ToolResult {
176                    id: call.id.clone(),
177                    name: call.name.clone(),
178                    output: result.clone(),
179                    step,
180                });
181                self.transcript
182                    .push(Message::tool_result(call.id.clone(), result));
183            }
184        }
185
186        warn!(target: "recursive::agent", "step budget exceeded");
187        self.emit(StepEvent::Finished {
188            reason: FinishReason::BudgetExceeded,
189            steps: self.max_steps,
190        });
191        Err(Error::StepBudgetExceeded(self.max_steps))
192    }
193
194    pub fn transcript(&self) -> &[Message] {
195        &self.transcript
196    }
197
198    fn emit(&self, event: StepEvent) {
199        if let Some(tx) = &self.events {
200            let _ = tx.send(event);
201        }
202    }
203}
204
205#[derive(Default)]
206pub struct AgentBuilder {
207    llm: Option<Arc<dyn LlmProvider>>,
208    tools: ToolRegistry,
209    system: Option<String>,
210    max_steps: Option<usize>,
211    events: Option<mpsc::UnboundedSender<StepEvent>>,
212}
213
214impl AgentBuilder {
215    pub fn llm(mut self, llm: Arc<dyn LlmProvider>) -> Self {
216        self.llm = Some(llm);
217        self
218    }
219    pub fn tools(mut self, tools: ToolRegistry) -> Self {
220        self.tools = tools;
221        self
222    }
223    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
224        self.system = Some(prompt.into());
225        self
226    }
227    pub fn max_steps(mut self, n: usize) -> Self {
228        self.max_steps = Some(n);
229        self
230    }
231    pub fn events(mut self, tx: mpsc::UnboundedSender<StepEvent>) -> Self {
232        self.events = Some(tx);
233        self
234    }
235    pub fn build(self) -> Result<Agent> {
236        let llm = self
237            .llm
238            .ok_or_else(|| Error::Config("agent: missing llm provider".into()))?;
239        let mut transcript = Vec::new();
240        if let Some(sys) = self.system {
241            transcript.push(Message::system(sys));
242        }
243        Ok(Agent {
244            llm,
245            tools: self.tools,
246            transcript,
247            max_steps: self.max_steps.unwrap_or(32),
248            events: self.events,
249        })
250    }
251}
252
253fn truncate(s: &str, n: usize) -> String {
254    if s.chars().count() <= n {
255        s.to_string()
256    } else {
257        let mut out: String = s.chars().take(n).collect();
258        out.push_str("...");
259        out
260    }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266    use crate::llm::{Completion, MockProvider, ToolCall};
267    use crate::tools::Tool;
268    use async_trait::async_trait;
269    use serde_json::{json, Value};
270
271    struct Adder;
272
273    #[async_trait]
274    impl Tool for Adder {
275        fn spec(&self) -> crate::llm::ToolSpec {
276            crate::llm::ToolSpec {
277                name: "add".into(),
278                description: "add a and b".into(),
279                parameters: json!({"type":"object","properties":{"a":{"type":"integer"},"b":{"type":"integer"}}}),
280            }
281        }
282        async fn execute(&self, args: Value) -> Result<String> {
283            let a = args["a"].as_i64().unwrap_or(0);
284            let b = args["b"].as_i64().unwrap_or(0);
285            Ok((a + b).to_string())
286        }
287    }
288
289    #[tokio::test]
290    async fn terminates_when_model_emits_no_tool_calls() {
291        let llm = Arc::new(MockProvider::new(vec![Completion {
292            content: "done".into(),
293            tool_calls: vec![],
294            finish_reason: Some("stop".into()),
295        }]));
296        let mut agent = Agent::builder().llm(llm).build().unwrap();
297        let out = agent.run("hi").await.unwrap();
298        assert_eq!(out.final_message.as_deref(), Some("done"));
299        assert_eq!(out.steps, 1);
300        assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
301    }
302
303    #[tokio::test]
304    async fn runs_a_tool_then_completes() {
305        let llm = Arc::new(MockProvider::new(vec![
306            Completion {
307                content: "let me add".into(),
308                tool_calls: vec![ToolCall {
309                    id: "c1".into(),
310                    name: "add".into(),
311                    arguments: json!({"a":2,"b":3}),
312                }],
313                finish_reason: Some("tool_calls".into()),
314            },
315            Completion {
316                content: "5".into(),
317                tool_calls: vec![],
318                finish_reason: Some("stop".into()),
319            },
320        ]));
321        let tools = ToolRegistry::new().register(Arc::new(Adder));
322        let mut agent = Agent::builder().llm(llm).tools(tools).build().unwrap();
323        let out = agent.run("what is 2+3?").await.unwrap();
324        assert_eq!(out.final_message.as_deref(), Some("5"));
325        assert_eq!(out.steps, 2);
326        // transcript should be: user, assistant(tool_call), tool, assistant("5")
327        assert_eq!(out.transcript.len(), 4);
328    }
329
330    #[tokio::test]
331    async fn reports_step_budget_exceeded() {
332        let mut script = Vec::new();
333        for _ in 0..10 {
334            script.push(Completion {
335                content: "".into(),
336                tool_calls: vec![ToolCall {
337                    id: "x".into(),
338                    name: "add".into(),
339                    arguments: json!({"a":1,"b":1}),
340                }],
341                finish_reason: Some("tool_calls".into()),
342            });
343        }
344        let llm = Arc::new(MockProvider::new(script));
345        let tools = ToolRegistry::new().register(Arc::new(Adder));
346        let mut agent = Agent::builder()
347            .llm(llm)
348            .tools(tools)
349            .max_steps(3)
350            .build()
351            .unwrap();
352        let err = agent.run("loop").await.unwrap_err();
353        assert!(matches!(err, Error::StepBudgetExceeded(3)));
354    }
355
356    #[tokio::test]
357    async fn unknown_tool_returns_error_to_model_not_abort() {
358        let llm = Arc::new(MockProvider::new(vec![
359            Completion {
360                content: "".into(),
361                tool_calls: vec![ToolCall {
362                    id: "c1".into(),
363                    name: "nope".into(),
364                    arguments: json!({}),
365                }],
366                finish_reason: Some("tool_calls".into()),
367            },
368            Completion {
369                content: "ok i give up".into(),
370                tool_calls: vec![],
371                finish_reason: Some("stop".into()),
372            },
373        ]));
374        let mut agent = Agent::builder().llm(llm).build().unwrap();
375        let out = agent.run("call a missing tool").await.unwrap();
376        // tool message must contain the error so the model can recover
377        let tool_msg = out
378            .transcript
379            .iter()
380            .find(|m| m.role == crate::message::Role::Tool)
381            .unwrap();
382        assert!(tool_msg.content.contains("ERROR"));
383        assert_eq!(out.final_message.as_deref(), Some("ok i give up"));
384    }
385
386    #[tokio::test]
387    async fn emits_events_in_order() {
388        let llm = Arc::new(MockProvider::new(vec![
389            Completion {
390                content: "thinking".into(),
391                tool_calls: vec![ToolCall {
392                    id: "c1".into(),
393                    name: "add".into(),
394                    arguments: json!({"a":1,"b":1}),
395                }],
396                finish_reason: Some("tool_calls".into()),
397            },
398            Completion {
399                content: "two".into(),
400                tool_calls: vec![],
401                finish_reason: Some("stop".into()),
402            },
403        ]));
404        let tools = ToolRegistry::new().register(Arc::new(Adder));
405        let (tx, mut rx) = mpsc::unbounded_channel();
406        let mut agent = Agent::builder()
407            .llm(llm)
408            .tools(tools)
409            .events(tx)
410            .build()
411            .unwrap();
412        agent.run("add").await.unwrap();
413        let mut kinds = Vec::new();
414        while let Ok(e) = rx.try_recv() {
415            kinds.push(match e {
416                StepEvent::AssistantText { .. } => "text",
417                StepEvent::ToolCall { .. } => "call",
418                StepEvent::ToolResult { .. } => "result",
419                StepEvent::Finished { .. } => "done",
420            });
421        }
422        assert_eq!(kinds, vec!["text", "call", "result", "text", "done"]);
423    }
424
425    #[tokio::test]
426    async fn stops_when_repeated_call_keeps_erroring() {
427        // MockProvider scripted to call a non-existent tool 4 times
428        let mut script = Vec::new();
429        for i in 0..4 {
430            script.push(Completion {
431                content: "".into(),
432                tool_calls: vec![ToolCall {
433                    id: format!("c{}", i),
434                    name: "UnknownTool".into(),
435                    arguments: json!({"arg": "value"}),
436                }],
437                finish_reason: Some("tool_calls".into()),
438            });
439        }
440        let llm = Arc::new(MockProvider::new(script));
441        let mut agent = Agent::builder().llm(llm).max_steps(10).build().unwrap();
442        let out = agent.run("call unknown tool").await.unwrap();
443
444        // Should be stuck after 3 consecutive errors
445        assert!(matches!(out.finish, FinishReason::Stuck { .. }));
446        if let FinishReason::Stuck {
447            repeated_call,
448            repeats,
449        } = &out.finish
450        {
451            assert_eq!(repeated_call, "UnknownTool");
452            assert_eq!(*repeats, 3);
453        }
454    }
455
456    #[tokio::test]
457    async fn does_not_trigger_when_args_differ() {
458        // MockProvider scripted to call same tool with different args each time
459        let mut script = Vec::new();
460        for i in 0..3 {
461            script.push(Completion {
462                content: "".into(),
463                tool_calls: vec![ToolCall {
464                    id: format!("c{}", i),
465                    name: "add".into(),
466                    arguments: json!({"a": i, "b": i}),
467                }],
468                finish_reason: Some("tool_calls".into()),
469            });
470        }
471        let llm = Arc::new(MockProvider::new(script));
472        let tools = ToolRegistry::new().register(Arc::new(Adder));
473        // Set max_steps low so test terminates with budget
474        let mut agent = Agent::builder()
475            .llm(llm)
476            .tools(tools)
477            .max_steps(3)
478            .build()
479            .unwrap();
480        let err = agent.run("add with different args").await.unwrap_err();
481
482        // Should hit budget, not stuck (args differ each time)
483        assert!(matches!(err, Error::StepBudgetExceeded(3)));
484    }
485}