1use 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
20const 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 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 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 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 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 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 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 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 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 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 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 assert!(matches!(err, Error::StepBudgetExceeded(3)));
484 }
485}