1use std::sync::Arc;
12
13use serde::{Deserialize, Serialize};
14use tokio::sync::mpsc;
15use tracing::{debug, info, warn};
16
17use crate::compact::Compactor;
18use crate::error::{Error, Result};
19use crate::hooks::{Hook, HookAction, HookEvent, HookRegistry};
20use crate::llm::{Completion, LlmProvider, StreamSender, TokenUsage, ToolCall};
21use crate::message::Message;
22use crate::tools::ToolRegistry;
23
24const STUCK_THRESHOLD: usize = 3;
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
30#[serde(rename_all = "snake_case")]
31pub enum PermissionDecision {
32 Allow,
34 Deny(String),
36 Transform(serde_json::Value),
38}
39
40pub type PermissionHook = Arc<dyn Fn(&str, &serde_json::Value) -> PermissionDecision + Send + Sync>;
49
50const TRIM_PLACEHOLDER: &str = "[older tool output trimmed to fit budget]";
51
52#[derive(Debug, Clone, PartialEq, Default)]
54pub enum PlanningMode {
55 #[default]
57 Immediate,
58 PlanFirst,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
63#[serde(tag = "kind", rename_all = "snake_case")]
64#[non_exhaustive]
65pub enum StepEvent {
66 AssistantText { text: String, step: usize },
72 ToolCall { call: ToolCall, step: usize },
78 Latency { step: usize, llm_ms: u64 },
83 ToolResult {
90 id: String,
91 name: String,
92 output: String,
93 step: usize,
94 },
95 Usage { usage: TokenUsage, step: usize },
101 PartialToken { text: String, step: usize },
107 Compacted {
115 removed: usize,
116 kept: usize,
117 summary_chars: usize,
118 step: usize,
119 },
120 Finished { reason: FinishReason, steps: usize },
127 PlanProposed {
129 plan_text: String,
131 tool_calls: Vec<ToolCall>,
133 },
134 PlanConfirmed,
136 PlanRejected { reason: String },
138}
139
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
141#[serde(tag = "kind", rename_all = "snake_case")]
142#[non_exhaustive]
156pub enum FinishReason {
157 NoMoreToolCalls,
159 BudgetExceeded,
161 ProviderStop(String),
163 Stuck {
165 repeated_call: String,
166 repeats: usize,
167 },
168 TranscriptLimit { chars: usize, limit: usize },
170 PlanPending,
172}
173
174#[derive(Debug, Clone)]
175pub struct AgentOutcome {
176 pub final_message: Option<String>,
177 pub transcript: Vec<Message>,
178 pub steps: usize,
179 pub finish: FinishReason,
180 pub total_usage: TokenUsage,
181 pub total_llm_latency_ms: u64,
182}
183
184pub struct Agent {
214 llm: Arc<dyn LlmProvider>,
215 tools: ToolRegistry,
216 transcript: Vec<Message>,
217 max_steps: usize,
218 max_transcript_chars: Option<usize>,
219 events: Option<mpsc::UnboundedSender<StepEvent>>,
220 streaming: bool,
221 total_llm_latency_ms: u64,
222 compactor: Option<Compactor>,
223 permission_hook: Option<PermissionHook>,
224 hooks: HookRegistry,
225 planning_mode: PlanningMode,
226 plan_buffer: Option<Vec<ToolCall>>,
227 plan_confirmed: bool,
228}
229
230impl Agent {
231 pub fn builder() -> AgentBuilder {
232 AgentBuilder::default()
233 }
234
235 pub fn set_transcript(&mut self, transcript: Vec<Message>) {
238 self.transcript = transcript;
239 }
240
241 pub fn set_events(&mut self, tx: Option<mpsc::UnboundedSender<StepEvent>>) {
245 self.events = tx;
246 }
247
248 pub fn confirm_plan(&mut self) {
250 self.plan_confirmed = true;
251 self.emit(StepEvent::PlanConfirmed);
252 }
253
254 pub fn reject_plan(&mut self, reason: &str) {
256 if let Some(calls) = self.plan_buffer.take() {
258 for call in &calls {
259 let result = format!(
260 "ERROR: Plan rejected. Reason: {reason}. \
261 The tool call {}({}) was not executed. \
262 Please revise your approach.",
263 call.name,
264 serde_json::to_string(&call.arguments).unwrap_or_default()
265 );
266 self.transcript
267 .push(Message::tool_result(call.id.clone(), result));
268 }
269 }
270 self.emit(StepEvent::PlanRejected {
271 reason: reason.into(),
272 });
273 }
274
275 async fn execute_tool_calls(
280 &mut self,
281 calls: &[ToolCall],
282 _step: usize,
283 ) -> Vec<(String, String, String, serde_json::Value)> {
284 let mut results: Vec<(String, String, String, serde_json::Value)> = Vec::new();
285
286 struct PendingCall {
288 id: String,
289 name: String,
290 args: serde_json::Value,
291 }
292 let mut pending: Vec<PendingCall> = Vec::new();
293
294 for call in calls {
295 let effective_args = if let Some(ref hook) = self.permission_hook {
296 match hook(&call.name, &call.arguments) {
297 PermissionDecision::Allow => call.arguments.clone(),
298 PermissionDecision::Deny(reason) => {
299 let result = format!("ERROR: {reason}");
300 results.push((
301 call.id.clone(),
302 call.name.clone(),
303 result,
304 call.arguments.clone(),
305 ));
306 continue;
307 }
308 PermissionDecision::Transform(new_args) => new_args,
309 }
310 } else {
311 call.arguments.clone()
312 };
313
314 let hook_action = self.hooks.dispatch(HookEvent::PreToolCall {
316 name: &call.name,
317 args: &effective_args,
318 });
319 match hook_action {
320 HookAction::Skip => {
321 let result = "ERROR: tool call skipped by hook".to_string();
322 results.push((
323 call.id.clone(),
324 call.name.clone(),
325 result,
326 call.arguments.clone(),
327 ));
328 continue;
329 }
330 HookAction::Error(msg) => {
331 let result = format!("ERROR: {msg}");
332 results.push((
333 call.id.clone(),
334 call.name.clone(),
335 result,
336 call.arguments.clone(),
337 ));
338 continue;
339 }
340 HookAction::Continue => {}
341 }
342
343 pending.push(PendingCall {
344 id: call.id.clone(),
345 name: call.name.clone(),
346 args: effective_args,
347 });
348 }
349
350 let mut i = 0;
352 while i < pending.len() {
353 if self.tools.is_readonly(&pending[i].name) {
354 let batch_start = i;
356 while i < pending.len() && self.tools.is_readonly(&pending[i].name) {
357 i += 1;
358 }
359 let batch: Vec<PendingCall> = pending.drain(batch_start..i).collect();
360 i = batch_start;
361
362 let mut join_set = tokio::task::JoinSet::new();
363 for pc in &batch {
364 let name = pc.name.clone();
365 let args = pc.args.clone();
366 let tools = self.tools.clone();
367 join_set.spawn(async move {
368 let tool_start = std::time::Instant::now();
369 let result = match tools.invoke(&name, args).await {
370 Ok(output) => output,
371 Err(err) => format!("ERROR: {err}"),
372 };
373 let duration_ms = tool_start.elapsed().as_millis() as u64;
374 (name, result, duration_ms)
375 });
376 }
377
378 let mut batch_results: Vec<(String, String, u64)> = Vec::new();
379 while let Some(res) = join_set.join_next().await {
380 let (name, result, duration_ms) = res.unwrap();
381 batch_results.push((name, result, duration_ms));
382 }
383
384 for pc in &batch {
385 let (_, result, duration_ms) = batch_results
386 .iter()
387 .find(|(n, _, _)| n == &pc.name)
388 .unwrap();
389 results.push((
390 pc.id.clone(),
391 pc.name.clone(),
392 result.clone(),
393 pc.args.clone(),
394 ));
395 self.hooks.dispatch(HookEvent::PostToolCall {
396 name: &pc.name,
397 args: &pc.args,
398 result,
399 duration_ms: *duration_ms,
400 });
401 }
402 } else {
403 let pc = pending.remove(i);
404 let tool_start = std::time::Instant::now();
405 let result = match self.tools.invoke(&pc.name, pc.args.clone()).await {
406 Ok(output) => output,
407 Err(err) => format!("ERROR: {err}"),
408 };
409 let duration_ms = tool_start.elapsed().as_millis() as u64;
410 results.push((
411 pc.id.clone(),
412 pc.name.clone(),
413 result.clone(),
414 pc.args.clone(),
415 ));
416 self.hooks.dispatch(HookEvent::PostToolCall {
417 name: &pc.name,
418 args: &pc.args,
419 result: &result,
420 duration_ms,
421 });
422 }
423 }
424
425 results
426 }
427
428 #[tracing::instrument(skip(self), fields(goal))]
429 pub async fn run(&mut self, goal: impl Into<String>) -> Result<AgentOutcome> {
430 let goal = goal.into();
431 info!(target: "recursive::agent", goal = %truncate(&goal, 200), "agent run starting");
432 self.transcript.push(Message::user(goal.clone()));
433 self.hooks.dispatch(HookEvent::SessionStart { goal: &goal });
434
435 let mut final_message: Option<String> = None;
436 let specs = self.tools.specs();
437
438 let mut last_call_key: Option<(String, String)> = None;
440 let mut consecutive_errors: usize = 0;
441
442 let mut total_usage = TokenUsage::default();
444 self.total_llm_latency_ms = 0;
446
447 for step in 1..=self.max_steps {
448 let step_span = tracing::info_span!("agent.step", step);
449 let _guard = step_span.enter();
450 if let Some(limit) = self.max_transcript_chars {
453 self.maybe_trim_transcript(limit, step);
454 let chars: usize = self.transcript.iter().map(|m| m.content.len()).sum();
455 if chars >= limit {
456 let finish = FinishReason::TranscriptLimit { chars, limit };
457 self.emit(StepEvent::Finished {
458 reason: finish.clone(),
459 steps: step,
460 });
461 return Ok(AgentOutcome {
462 final_message,
463 transcript: std::mem::take(&mut self.transcript),
464 steps: step,
465 finish,
466 total_usage,
467 total_llm_latency_ms: self.total_llm_latency_ms,
468 });
469 }
470 }
471
472 self.hooks.dispatch(HookEvent::PreCompact {
474 transcript_len: self.transcript.iter().map(|m| m.content.len()).sum(),
475 });
476 self.maybe_compact(step).await?;
477
478 if self.plan_confirmed {
481 self.plan_confirmed = false;
482 if let Some(calls) = self.plan_buffer.take() {
483 let results = self.execute_tool_calls(&calls, step).await;
484 for (id, name, output, _args) in results {
485 self.emit(StepEvent::ToolResult {
486 id: id.clone(),
487 name: name.clone(),
488 output: output.clone(),
489 step,
490 });
491 self.transcript.push(Message::tool_result(id, output));
492 }
493 continue;
494 }
495 }
496
497 debug!(target: "recursive::agent", step, "calling llm");
498 let start = std::time::Instant::now();
499 let completion: Completion = if self.streaming {
500 let (delta_tx, mut delta_rx) = mpsc::unbounded_channel::<String>();
502 let stream_tx: Option<StreamSender> = Some(delta_tx);
503 let events_tx = self.events.clone();
505 tokio::spawn(async move {
506 while let Some(text) = delta_rx.recv().await {
507 if let Some(ref tx) = events_tx {
508 let _ = tx.send(StepEvent::PartialToken { text, step });
509 }
510 }
511 });
512 self.llm.stream(&self.transcript, &specs, stream_tx).await?
513 } else {
514 self.llm.complete(&self.transcript, &specs).await?
515 };
516 let llm_ms = start.elapsed().as_millis() as u64;
517 self.total_llm_latency_ms = self.total_llm_latency_ms.saturating_add(llm_ms);
518 self.emit(StepEvent::Latency { step, llm_ms });
519
520 if let Some(u) = completion.usage {
522 total_usage = total_usage.accumulate(u);
523 self.emit(StepEvent::Usage { usage: u, step });
524 }
525
526 if !completion.content.is_empty() {
527 self.emit(StepEvent::AssistantText {
528 text: completion.content.clone(),
529 step,
530 });
531 final_message = Some(completion.content.clone());
532 }
533
534 if completion.tool_calls.is_empty() {
535 if matches!(completion.finish_reason.as_deref(), Some("length")) {
541 self.emit(StepEvent::Finished {
542 reason: FinishReason::ProviderStop("length".into()),
543 steps: step,
544 });
545 return Err(Error::ProviderTruncated("length".into()));
546 }
547
548 self.transcript
549 .push(Message::assistant(completion.content.clone()));
550 if completion.reasoning_content.is_some() {
552 if let Some(msg) = self.transcript.last_mut() {
553 msg.reasoning_content = completion.reasoning_content.clone();
554 }
555 }
556 let finish = match completion.finish_reason {
557 Some(r) if r != "stop" && r != "end_turn" => FinishReason::ProviderStop(r),
558 _ => FinishReason::NoMoreToolCalls,
559 };
560 self.emit(StepEvent::Finished {
561 reason: finish.clone(),
562 steps: step,
563 });
564 let outcome = AgentOutcome {
565 final_message,
566 transcript: std::mem::take(&mut self.transcript),
567 steps: step,
568 finish,
569 total_usage,
570 total_llm_latency_ms: self.total_llm_latency_ms,
571 };
572 self.hooks
573 .dispatch(HookEvent::SessionEnd { outcome: &outcome });
574 return Ok(outcome);
575 }
576
577 self.transcript.push(Message::assistant_with_tool_calls(
578 completion.content.clone(),
579 completion.tool_calls.clone(),
580 ));
581 if completion.reasoning_content.is_some() {
583 if let Some(msg) = self.transcript.last_mut() {
584 msg.reasoning_content = completion.reasoning_content.clone();
585 }
586 }
587
588 for call in &completion.tool_calls {
590 self.emit(StepEvent::ToolCall {
591 call: call.clone(),
592 step,
593 });
594 }
595
596 if self.planning_mode == PlanningMode::PlanFirst && self.plan_buffer.is_none() {
599 self.plan_buffer = Some(completion.tool_calls.clone());
600
601 let plan_text = completion
603 .tool_calls
604 .iter()
605 .map(|tc| {
606 let args_str = serde_json::to_string(&tc.arguments).unwrap_or_default();
607 format!(" - {}({})", tc.name, args_str)
608 })
609 .collect::<Vec<_>>()
610 .join("\n");
611 let plan_text = format!(
612 "The agent proposes the following steps:\n{}\n\nConfirm or reject this plan.",
613 plan_text
614 );
615
616 self.emit(StepEvent::PlanProposed {
617 plan_text: plan_text.clone(),
618 tool_calls: completion.tool_calls.clone(),
619 });
620
621 return Ok(AgentOutcome {
623 final_message: Some(plan_text),
624 transcript: std::mem::take(&mut self.transcript),
625 steps: step,
626 finish: FinishReason::PlanPending,
627 total_usage: TokenUsage::default(),
628 total_llm_latency_ms: self.total_llm_latency_ms,
629 });
630 }
631
632 let results = self.execute_tool_calls(&completion.tool_calls, step).await;
636
637 for (id, name, result, args) in &results {
639 self.emit(StepEvent::ToolResult {
640 id: id.clone(),
641 name: name.clone(),
642 output: result.clone(),
643 step,
644 });
645
646 let call_key = (
648 name.clone(),
649 serde_json::to_string(args).unwrap_or_default(),
650 );
651 let is_error = result.starts_with("ERROR:");
652
653 if is_error {
654 if last_call_key == Some(call_key.clone()) {
655 consecutive_errors += 1;
656 } else {
657 consecutive_errors = 1;
658 }
659 } else {
660 consecutive_errors = 0;
661 }
662
663 last_call_key = Some(call_key);
664
665 if consecutive_errors >= STUCK_THRESHOLD {
667 let repeated_call = name.clone();
668 let repeats = consecutive_errors;
669 let finish = FinishReason::Stuck {
670 repeated_call,
671 repeats,
672 };
673 self.emit(StepEvent::Finished {
674 reason: finish.clone(),
675 steps: step,
676 });
677 return Ok(AgentOutcome {
678 final_message,
679 transcript: std::mem::take(&mut self.transcript),
680 steps: step,
681 finish,
682 total_usage,
683 total_llm_latency_ms: self.total_llm_latency_ms,
684 });
685 }
686
687 self.transcript
688 .push(Message::tool_result(id.clone(), result.clone()));
689 }
690 }
691
692 warn!(target: "recursive::agent", "step budget exceeded");
693 let finish = FinishReason::BudgetExceeded;
694 self.emit(StepEvent::Finished {
695 reason: finish.clone(),
696 steps: self.max_steps,
697 });
698 let outcome = AgentOutcome {
699 final_message,
700 transcript: std::mem::take(&mut self.transcript),
701 steps: self.max_steps,
702 finish,
703 total_usage,
704 total_llm_latency_ms: self.total_llm_latency_ms,
705 };
706 self.hooks
707 .dispatch(HookEvent::SessionEnd { outcome: &outcome });
708 Ok(outcome)
709 }
710
711 async fn maybe_compact(&mut self, step: usize) -> Result<()> {
719 let compactor = match &self.compactor {
720 Some(c) => c,
721 None => return Ok(()),
722 };
723
724 let chars = Compactor::estimate_chars(&self.transcript);
725 if chars < compactor.threshold_chars {
726 return Ok(());
727 }
728
729 let min_messages = compactor.keep_recent_n + 2;
731 if self.transcript.len() < min_messages {
732 return Ok(());
733 }
734
735 let summary_msg = compactor
736 .compact(self.llm.as_ref(), &self.transcript)
737 .await?;
738 let summary_chars = summary_msg.content.len();
739
740 let keep = compactor.keep_recent_n;
743 let mut split = self.transcript.len().saturating_sub(keep);
744
745 while split > 0 && matches!(self.transcript[split].role, crate::message::Role::Tool) {
754 split -= 1;
755 }
756
757 let removed = split;
758 let kept = self.transcript.len() - split;
759
760 self.transcript.drain(..split);
762 self.transcript.insert(0, summary_msg);
763
764 self.hooks.dispatch(HookEvent::PostCompact {
765 removed,
766 summary_chars,
767 });
768
769 self.emit(StepEvent::Compacted {
770 removed,
771 kept,
772 summary_chars,
773 step,
774 });
775
776 Ok(())
777 }
778
779 fn maybe_trim_transcript(&mut self, limit: usize, step: usize) {
780 let mut chars: usize = self.transcript.iter().map(|m| m.content.len()).sum();
781 if chars < limit {
782 return;
783 }
784
785 let mut trimmed_count: usize = 0;
786 let placeholder_len = TRIM_PLACEHOLDER.len();
787
788 for msg in self.transcript.iter_mut().skip(1) {
791 if msg.role == crate::message::Role::Tool && msg.content.len() > 200 {
792 let old_len = msg.content.len();
793 msg.content = TRIM_PLACEHOLDER.to_string();
794 trimmed_count += 1;
795 chars = chars
797 .saturating_sub(old_len)
798 .saturating_add(placeholder_len);
799 if chars < limit {
800 break;
801 }
802 }
803 }
804
805 if trimmed_count > 0 {
806 let note = format!(
807 "[trimmed {} old tool result{} to fit budget]",
808 trimmed_count,
809 if trimmed_count == 1 { "" } else { "s" }
810 );
811 self.emit(StepEvent::AssistantText { text: note, step });
812 }
813 }
814
815 pub fn transcript(&self) -> &[Message] {
816 &self.transcript
817 }
818
819 fn emit(&self, event: StepEvent) {
820 if let Some(tx) = &self.events {
821 let _ = tx.send(event);
822 }
823 }
824}
825
826#[derive(Default)]
841pub struct AgentBuilder {
842 llm: Option<Arc<dyn LlmProvider>>,
843 tools: ToolRegistry,
844 system: Option<String>,
845 max_steps: Option<usize>,
846 max_transcript_chars: Option<usize>,
847 events: Option<mpsc::UnboundedSender<StepEvent>>,
848 seed: Vec<Message>,
849 streaming: bool,
850 compactor: Option<Compactor>,
851 permission_hook: Option<PermissionHook>,
852 hooks: HookRegistry,
853 planning_mode: PlanningMode,
854}
855
856impl AgentBuilder {
857 pub fn llm(mut self, llm: Arc<dyn LlmProvider>) -> Self {
862 self.llm = Some(llm);
863 self
864 }
865 pub fn tools(mut self, tools: ToolRegistry) -> Self {
871 self.tools = tools;
872 self
873 }
874 pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
880 self.system = Some(prompt.into());
881 self
882 }
883 pub fn max_steps(mut self, n: usize) -> Self {
889 self.max_steps = Some(n);
890 self
891 }
892 pub fn max_transcript_chars(mut self, n: usize) -> Self {
904 self.max_transcript_chars = Some(n);
905 self
906 }
907 pub fn events(mut self, tx: mpsc::UnboundedSender<StepEvent>) -> Self {
914 self.events = Some(tx);
915 self
916 }
917 pub fn seed_transcript(mut self, messages: Vec<Message>) -> Self {
922 self.seed = messages;
923 self
924 }
925 pub fn streaming(mut self, enabled: bool) -> Self {
931 self.streaming = enabled;
932 self
933 }
934 pub fn compactor(mut self, compactor: Compactor) -> Self {
940 self.compactor = Some(compactor);
941 self
942 }
943 pub fn permission_hook<F>(mut self, hook: F) -> Self
949 where
950 F: Fn(&str, &serde_json::Value) -> PermissionDecision + Send + Sync + 'static,
951 {
952 self.permission_hook = Some(Arc::new(hook));
953 self
954 }
955
956 pub fn permission_hook_opt(mut self, hook: Option<PermissionHook>) -> Self {
962 self.permission_hook = hook;
963 self
964 }
965 pub fn hook(mut self, hook: Arc<dyn Hook>) -> Self {
970 self.hooks.register(hook);
971 self
972 }
973 pub fn planning_mode(mut self, mode: PlanningMode) -> Self {
979 self.planning_mode = mode;
980 self
981 }
982 pub fn build(self) -> Result<Agent> {
983 let llm = self.llm.ok_or_else(|| Error::Config {
984 message: "agent: missing llm provider".into(),
985 })?;
986 let mut transcript = Vec::new();
987 if let Some(sys) = self.system {
988 transcript.push(Message::system(sys));
989 }
990 transcript.extend(self.seed);
991 Ok(Agent {
992 llm,
993 tools: self.tools,
994 transcript,
995 max_steps: self.max_steps.unwrap_or(32),
996 max_transcript_chars: self.max_transcript_chars,
997 events: self.events,
998 streaming: self.streaming,
999 total_llm_latency_ms: 0,
1000 compactor: self.compactor,
1001 permission_hook: self.permission_hook,
1002 hooks: self.hooks,
1003 planning_mode: self.planning_mode,
1004 plan_buffer: None,
1005 plan_confirmed: false,
1006 })
1007 }
1008}
1009
1010fn truncate(s: &str, n: usize) -> String {
1011 if s.chars().count() <= n {
1012 s.to_string()
1013 } else {
1014 let mut out: String = s.chars().take(n).collect();
1015 out.push_str("...");
1016 out
1017 }
1018}
1019
1020#[cfg(test)]
1021mod tests {
1022 use super::*;
1023 use crate::llm::{Completion, MockProvider, TokenUsage, ToolCall};
1024 use crate::tools::Tool;
1025 use async_trait::async_trait;
1026 use serde_json::{json, Value};
1027
1028 struct Adder;
1029
1030 #[async_trait]
1031 impl Tool for Adder {
1032 fn spec(&self) -> crate::llm::ToolSpec {
1033 crate::llm::ToolSpec {
1034 name: "add".into(),
1035 description: "add a and b".into(),
1036 parameters: json!({"type":"object","properties":{"a":{"type":"integer"},"b":{"type":"integer"}}}),
1037 }
1038 }
1039 async fn execute(&self, args: Value) -> Result<String> {
1040 let a = args["a"].as_i64().unwrap_or(0);
1041 let b = args["b"].as_i64().unwrap_or(0);
1042 Ok((a + b).to_string())
1043 }
1044 }
1045
1046 #[tokio::test]
1047 async fn terminates_when_model_emits_no_tool_calls() {
1048 let llm = Arc::new(MockProvider::new(vec![Completion {
1049 content: "done".into(),
1050 tool_calls: vec![],
1051 finish_reason: Some("stop".into()),
1052 usage: None,
1053 reasoning_content: None,
1054 }]));
1055 let mut agent = Agent::builder().llm(llm).build().unwrap();
1056 let out = agent.run("hi").await.unwrap();
1057 assert_eq!(out.final_message.as_deref(), Some("done"));
1058 assert_eq!(out.steps, 1);
1059 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
1060 }
1061
1062 #[tokio::test]
1063 async fn runs_a_tool_then_completes() {
1064 let llm = Arc::new(MockProvider::new(vec![
1065 Completion {
1066 content: "let me add".into(),
1067 tool_calls: vec![ToolCall {
1068 id: "c1".into(),
1069 name: "add".into(),
1070 arguments: json!({"a":2,"b":3}),
1071 }],
1072 finish_reason: Some("tool_calls".into()),
1073 usage: None,
1074 reasoning_content: None,
1075 },
1076 Completion {
1077 content: "5".into(),
1078 tool_calls: vec![],
1079 finish_reason: Some("stop".into()),
1080 usage: None,
1081 reasoning_content: None,
1082 },
1083 ]));
1084 let tools = ToolRegistry::local().register(Arc::new(Adder));
1085 let mut agent = Agent::builder().llm(llm).tools(tools).build().unwrap();
1086 let out = agent.run("what is 2+3?").await.unwrap();
1087 assert_eq!(out.final_message.as_deref(), Some("5"));
1088 assert_eq!(out.steps, 2);
1089 assert_eq!(out.transcript.len(), 4);
1091 }
1092
1093 #[tokio::test]
1094 async fn reports_step_budget_exceeded() {
1095 let mut script = Vec::new();
1096 for _ in 0..10 {
1097 script.push(Completion {
1098 content: "".into(),
1099 tool_calls: vec![ToolCall {
1100 id: "x".into(),
1101 name: "add".into(),
1102 arguments: json!({"a":1,"b":1}),
1103 }],
1104 finish_reason: Some("tool_calls".into()),
1105 usage: None,
1106 reasoning_content: None,
1107 });
1108 }
1109 let llm = Arc::new(MockProvider::new(script));
1110 let tools = ToolRegistry::local().register(Arc::new(Adder));
1111 let mut agent = Agent::builder()
1112 .llm(llm)
1113 .tools(tools)
1114 .max_steps(3)
1115 .build()
1116 .unwrap();
1117 let out = agent.run("loop").await.unwrap();
1118 assert!(matches!(out.finish, FinishReason::BudgetExceeded));
1119 assert_eq!(out.steps, 3);
1120 assert!(!out.transcript.is_empty());
1123 }
1124
1125 #[tokio::test]
1126 async fn unknown_tool_returns_error_to_model_not_abort() {
1127 let llm = Arc::new(MockProvider::new(vec![
1128 Completion {
1129 content: "".into(),
1130 tool_calls: vec![ToolCall {
1131 id: "c1".into(),
1132 name: "nope".into(),
1133 arguments: json!({}),
1134 }],
1135 finish_reason: Some("tool_calls".into()),
1136 usage: None,
1137 reasoning_content: None,
1138 },
1139 Completion {
1140 content: "ok i give up".into(),
1141 tool_calls: vec![],
1142 finish_reason: Some("stop".into()),
1143 usage: None,
1144 reasoning_content: None,
1145 },
1146 ]));
1147 let mut agent = Agent::builder().llm(llm).build().unwrap();
1148 let out = agent.run("call a missing tool").await.unwrap();
1149 let tool_msg = out
1151 .transcript
1152 .iter()
1153 .find(|m| m.role == crate::message::Role::Tool)
1154 .unwrap();
1155 assert!(tool_msg.content.contains("ERROR"));
1156 assert_eq!(out.final_message.as_deref(), Some("ok i give up"));
1157 }
1158
1159 #[tokio::test]
1160 async fn emits_events_in_order() {
1161 let llm = Arc::new(MockProvider::new(vec![
1162 Completion {
1163 content: "thinking".into(),
1164 tool_calls: vec![ToolCall {
1165 id: "c1".into(),
1166 name: "add".into(),
1167 arguments: json!({"a":1,"b":1}),
1168 }],
1169 finish_reason: Some("tool_calls".into()),
1170 usage: None,
1171 reasoning_content: None,
1172 },
1173 Completion {
1174 content: "two".into(),
1175 tool_calls: vec![],
1176 finish_reason: Some("stop".into()),
1177 usage: None,
1178 reasoning_content: None,
1179 },
1180 ]));
1181 let tools = ToolRegistry::local().register(Arc::new(Adder));
1182 let (tx, mut rx) = mpsc::unbounded_channel();
1183 let mut agent = Agent::builder()
1184 .llm(llm)
1185 .tools(tools)
1186 .events(tx)
1187 .build()
1188 .unwrap();
1189 agent.run("add").await.unwrap();
1190 let mut kinds = Vec::new();
1191 while let Ok(e) = rx.try_recv() {
1192 kinds.push(match e {
1193 StepEvent::AssistantText { .. } => "text",
1194 StepEvent::ToolCall { .. } => "call",
1195 StepEvent::ToolResult { .. } => "result",
1196 StepEvent::Finished { .. } => "done",
1197 StepEvent::Usage { .. } => "usage",
1198 StepEvent::Latency { .. } => "latency",
1199 StepEvent::PartialToken { .. } => "partial",
1200 StepEvent::Compacted { .. } => "compacted",
1201 StepEvent::PlanProposed { .. } => "plan_proposed",
1202 StepEvent::PlanConfirmed => "plan_confirmed",
1203 StepEvent::PlanRejected { .. } => "plan_rejected",
1204 });
1205 }
1206 assert_eq!(
1207 kinds,
1208 vec!["latency", "text", "call", "result", "latency", "text", "done"]
1209 );
1210 }
1211
1212 #[tokio::test]
1213 async fn stops_when_repeated_call_keeps_erroring() {
1214 let mut script = Vec::new();
1216 for i in 0..4 {
1217 script.push(Completion {
1218 content: "".into(),
1219 tool_calls: vec![ToolCall {
1220 id: format!("c{}", i),
1221 name: "UnknownTool".into(),
1222 arguments: json!({"arg": "value"}),
1223 }],
1224 finish_reason: Some("tool_calls".into()),
1225 usage: None,
1226 reasoning_content: None,
1227 });
1228 }
1229 let llm = Arc::new(MockProvider::new(script));
1230 let mut agent = Agent::builder().llm(llm).max_steps(10).build().unwrap();
1231 let out = agent.run("call unknown tool").await.unwrap();
1232
1233 assert!(matches!(out.finish, FinishReason::Stuck { .. }));
1235 if let FinishReason::Stuck {
1236 repeated_call,
1237 repeats,
1238 } = &out.finish
1239 {
1240 assert_eq!(repeated_call, "UnknownTool");
1241 assert_eq!(*repeats, 3);
1242 }
1243 }
1244
1245 #[tokio::test]
1246 async fn truncated_response_surfaces_as_error() {
1247 let llm = Arc::new(MockProvider::new(vec![Completion {
1251 content: "I was going to say more but ran out of".into(),
1252 tool_calls: vec![],
1253 finish_reason: Some("length".into()),
1254 usage: None,
1255 reasoning_content: None,
1256 }]));
1257 let mut agent = Agent::builder().llm(llm).build().unwrap();
1258 let err = agent.run("hi").await.unwrap_err();
1259 assert!(matches!(err, Error::ProviderTruncated(ref s) if s == "length"));
1260 }
1261
1262 #[tokio::test]
1263 async fn does_not_trigger_when_args_differ() {
1264 let mut script = Vec::new();
1266 for i in 0..3 {
1267 script.push(Completion {
1268 content: "".into(),
1269 tool_calls: vec![ToolCall {
1270 id: format!("c{}", i),
1271 name: "add".into(),
1272 arguments: json!({"a": i, "b": i}),
1273 }],
1274 finish_reason: Some("tool_calls".into()),
1275 usage: None,
1276 reasoning_content: None,
1277 });
1278 }
1279 let llm = Arc::new(MockProvider::new(script));
1280 let tools = ToolRegistry::local().register(Arc::new(Adder));
1281 let mut agent = Agent::builder()
1283 .llm(llm)
1284 .tools(tools)
1285 .max_steps(3)
1286 .build()
1287 .unwrap();
1288 let out = agent.run("add with different args").await.unwrap();
1289
1290 assert!(matches!(out.finish, FinishReason::BudgetExceeded));
1292 assert_eq!(out.steps, 3);
1293 }
1294
1295 #[tokio::test]
1312 async fn budget_exceeded_yields_writable_transcript() {
1313 use crate::TranscriptFile;
1314
1315 let mut script = Vec::new();
1316 for i in 0..10 {
1317 script.push(Completion {
1318 content: String::new(),
1319 tool_calls: vec![ToolCall {
1320 id: format!("t{i}"),
1321 name: "adder".into(),
1322 arguments: json!({"a": i, "b": i + 1}),
1323 }],
1324 finish_reason: Some("tool_calls".into()),
1325 usage: None,
1326 reasoning_content: None,
1327 });
1328 }
1329 let llm = Arc::new(MockProvider::new(script));
1330 let tools = ToolRegistry::local().register(Arc::new(Adder));
1331 let mut agent = Agent::builder()
1332 .llm(llm)
1333 .tools(tools)
1334 .max_steps(3)
1335 .build()
1336 .unwrap();
1337 let out = agent.run("loop").await.unwrap();
1338
1339 assert!(matches!(out.finish, FinishReason::BudgetExceeded));
1340 assert!(
1341 !out.transcript.is_empty(),
1342 "transcript must survive BudgetExceeded for auto-resume"
1343 );
1344
1345 let tmp = tempfile::NamedTempFile::new().unwrap();
1346 let file =
1347 TranscriptFile::new(out.transcript.clone(), out.steps, Some("mock-model".into()));
1348 file.write_to(tmp.path()).unwrap();
1349
1350 let restored = TranscriptFile::read_from(tmp.path()).unwrap();
1351 assert_eq!(
1352 restored.messages().len(),
1353 out.transcript.len(),
1354 "round-trip transcript length must match"
1355 );
1356 }
1357
1358 #[tokio::test]
1359 async fn accumulates_usage_across_llm_calls() {
1360 let u1 = TokenUsage {
1361 prompt_tokens: 10,
1362 completion_tokens: 5,
1363 total_tokens: 15,
1364 cache_hit_tokens: 0,
1365 cache_miss_tokens: 0,
1366 };
1367 let u2 = TokenUsage {
1368 prompt_tokens: 10,
1369 completion_tokens: 5,
1370 total_tokens: 15,
1371 cache_hit_tokens: 0,
1372 cache_miss_tokens: 0,
1373 };
1374 let llm = Arc::new(MockProvider::new(vec![
1375 Completion {
1376 content: "step 1".into(),
1377 tool_calls: vec![ToolCall {
1378 id: "c1".into(),
1379 name: "add".into(),
1380 arguments: json!({"a":1,"b":1}),
1381 }],
1382 finish_reason: Some("tool_calls".into()),
1383 usage: Some(u1),
1384 reasoning_content: None,
1385 },
1386 Completion {
1387 content: "step 2".into(),
1388 tool_calls: vec![],
1389 finish_reason: Some("stop".into()),
1390 usage: Some(u2),
1391 reasoning_content: None,
1392 },
1393 ]));
1394 let tools = ToolRegistry::local().register(Arc::new(Adder));
1395 let mut agent = Agent::builder().llm(llm).tools(tools).build().unwrap();
1396 let out = agent.run("add twice").await.unwrap();
1397
1398 assert_eq!(out.total_usage.prompt_tokens, 20);
1399 assert_eq!(out.total_usage.completion_tokens, 10);
1400 assert_eq!(out.total_usage.total_tokens, 30);
1401 }
1402
1403 #[tokio::test]
1404 async fn outcome_has_zero_usage_when_provider_never_reports() {
1405 let llm = Arc::new(MockProvider::new(vec![Completion {
1406 content: "done".into(),
1407 tool_calls: vec![],
1408 finish_reason: Some("stop".into()),
1409 usage: None,
1410 reasoning_content: None,
1411 }]));
1412 let mut agent = Agent::builder().llm(llm).build().unwrap();
1413 let out = agent.run("hi").await.unwrap();
1414
1415 assert_eq!(out.total_usage, TokenUsage::default());
1416 }
1417
1418 #[tokio::test]
1419 async fn step_event_usage_emitted_per_llm_call() {
1420 let u = TokenUsage {
1421 prompt_tokens: 10,
1422 completion_tokens: 5,
1423 total_tokens: 15,
1424 cache_hit_tokens: 0,
1425 cache_miss_tokens: 0,
1426 };
1427 let llm = Arc::new(MockProvider::new(vec![Completion {
1428 content: "first".into(),
1429 tool_calls: vec![],
1430 finish_reason: Some("stop".into()),
1431 usage: Some(u),
1432 reasoning_content: None,
1433 }]));
1434 let (tx, mut rx) = mpsc::unbounded_channel();
1435 let mut agent = Agent::builder().llm(llm).events(tx).build().unwrap();
1436 agent.run("hi").await.unwrap();
1437
1438 let mut usage_events = 0;
1439 while let Ok(e) = rx.try_recv() {
1440 if matches!(e, StepEvent::Usage { .. }) {
1441 usage_events += 1;
1442 }
1443 }
1444 assert_eq!(usage_events, 1);
1445 }
1446
1447 #[tokio::test]
1448 async fn transcript_limit_stops_loop() {
1449 let mut script = Vec::new();
1454 for _ in 0..30 {
1455 script.push(Completion {
1456 content: "x".into(),
1457 tool_calls: vec![ToolCall {
1458 id: "c1".into(),
1459 name: "add".into(),
1460 arguments: json!({"a":1,"b":1}),
1461 }],
1462 finish_reason: Some("tool_calls".into()),
1463 usage: None,
1464 reasoning_content: None,
1465 });
1466 }
1467 let llm = Arc::new(MockProvider::new(script));
1468 let tools = ToolRegistry::local().register(Arc::new(Adder));
1469 let mut agent = Agent::builder()
1470 .llm(llm)
1471 .tools(tools)
1472 .max_transcript_chars(50)
1473 .max_steps(100)
1474 .build()
1475 .unwrap();
1476 let out = agent.run("hi").await.unwrap();
1477 assert!(matches!(out.finish, FinishReason::TranscriptLimit { .. }));
1478 if let FinishReason::TranscriptLimit { chars, limit } = &out.finish {
1479 assert!(*chars >= 50);
1480 assert_eq!(*limit, 50);
1481 }
1482 }
1483
1484 #[tokio::test]
1485 async fn transcript_limit_unset_runs_to_completion() {
1486 let llm = Arc::new(MockProvider::new(vec![
1487 Completion {
1488 content: "let me add".into(),
1489 tool_calls: vec![ToolCall {
1490 id: "c1".into(),
1491 name: "add".into(),
1492 arguments: json!({"a":2,"b":3}),
1493 }],
1494 finish_reason: Some("tool_calls".into()),
1495 usage: None,
1496 reasoning_content: None,
1497 },
1498 Completion {
1499 content: "5".into(),
1500 tool_calls: vec![],
1501 finish_reason: Some("stop".into()),
1502 usage: None,
1503 reasoning_content: None,
1504 },
1505 ]));
1506 let tools = ToolRegistry::local().register(Arc::new(Adder));
1507 let mut agent = Agent::builder().llm(llm).tools(tools).build().unwrap();
1508 let out = agent.run("what is 2+3?").await.unwrap();
1509 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
1510 }
1511
1512 #[tokio::test]
1513 async fn transcript_limit_is_checked_before_llm_call() {
1514 let llm = Arc::new(MockProvider::new(vec![]));
1517 let mut agent = Agent::builder()
1518 .llm(llm)
1519 .max_transcript_chars(10)
1520 .build()
1521 .unwrap();
1522 let out = agent
1523 .run("a very long goal that exceeds the limit")
1524 .await
1525 .unwrap();
1526 assert!(matches!(out.finish, FinishReason::TranscriptLimit { .. }));
1527 if let FinishReason::TranscriptLimit { chars, limit } = &out.finish {
1528 assert!(*chars >= 10);
1529 assert_eq!(*limit, 10);
1530 }
1531 assert_eq!(out.steps, 1);
1533 }
1534
1535 #[tokio::test]
1536 async fn compaction_triggers_with_low_threshold() {
1537 let llm = Arc::new(MockProvider::new(vec![
1543 Completion {
1544 content: "first call".into(),
1545 tool_calls: vec![ToolCall {
1546 id: "c1".into(),
1547 name: "add".into(),
1548 arguments: json!({"a":1,"b":1}),
1549 }],
1550 finish_reason: Some("tool_calls".into()),
1551 usage: None,
1552 reasoning_content: None,
1553 },
1554 Completion {
1555 content: "second call".into(),
1556 tool_calls: vec![ToolCall {
1557 id: "c1".into(),
1558 name: "add".into(),
1559 arguments: json!({"a":1,"b":1}),
1560 }],
1561 finish_reason: Some("tool_calls".into()),
1562 usage: None,
1563 reasoning_content: None,
1564 },
1565 Completion {
1567 content: "Summary: added numbers, tests pass.".into(),
1568 tool_calls: vec![],
1569 finish_reason: Some("stop".into()),
1570 usage: None,
1571 reasoning_content: None,
1572 },
1573 Completion {
1575 content: "done".into(),
1576 tool_calls: vec![],
1577 finish_reason: Some("stop".into()),
1578 usage: None,
1579 reasoning_content: None,
1580 },
1581 ]));
1582 let tools = ToolRegistry::local().register(Arc::new(Adder));
1583 let (tx, mut rx) = mpsc::unbounded_channel();
1584 let compactor = crate::compact::Compactor::new(10).keep_recent_n(2);
1585 let mut agent = Agent::builder()
1586 .llm(llm)
1587 .tools(tools)
1588 .events(tx)
1589 .compactor(compactor)
1590 .max_steps(10)
1591 .build()
1592 .unwrap();
1593 let out = agent.run("hi").await.unwrap();
1594 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
1595
1596 let mut compacted_count = 0;
1598 while let Ok(e) = rx.try_recv() {
1599 if matches!(e, StepEvent::Compacted { .. }) {
1600 compacted_count += 1;
1601 }
1602 }
1603 assert_eq!(compacted_count, 1, "expected exactly one Compacted event");
1604
1605 let summary_msgs: Vec<&Message> = out
1607 .transcript
1608 .iter()
1609 .filter(|m| m.role == crate::message::Role::System)
1610 .collect();
1611 assert!(
1612 !summary_msgs.is_empty(),
1613 "expected at least one system message (the summary)"
1614 );
1615 assert!(
1616 summary_msgs
1617 .iter()
1618 .any(|m| m.content.contains("[compacted:")),
1619 "expected a system message with compacted header"
1620 );
1621 }
1622
1623 #[tokio::test]
1624 async fn compaction_disabled_by_default() {
1625 let llm = Arc::new(MockProvider::new(vec![
1627 Completion {
1628 content: "let me add".into(),
1629 tool_calls: vec![ToolCall {
1630 id: "c1".into(),
1631 name: "add".into(),
1632 arguments: json!({"a":1,"b":1}),
1633 }],
1634 finish_reason: Some("tool_calls".into()),
1635 usage: None,
1636 reasoning_content: None,
1637 },
1638 Completion {
1639 content: "done".into(),
1640 tool_calls: vec![],
1641 finish_reason: Some("stop".into()),
1642 usage: None,
1643 reasoning_content: None,
1644 },
1645 ]));
1646 let tools = ToolRegistry::local().register(Arc::new(Adder));
1647 let (tx, mut rx) = mpsc::unbounded_channel();
1648 let mut agent = Agent::builder()
1649 .llm(llm)
1650 .tools(tools)
1651 .events(tx)
1652 .max_steps(10)
1653 .build()
1654 .unwrap();
1655 let out = agent.run("hi").await.unwrap();
1656 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
1657
1658 let mut compacted_count = 0;
1660 while let Ok(e) = rx.try_recv() {
1661 if matches!(e, StepEvent::Compacted { .. }) {
1662 compacted_count += 1;
1663 }
1664 }
1665 assert_eq!(
1666 compacted_count, 0,
1667 "expected no Compacted events by default"
1668 );
1669 }
1670
1671 #[tokio::test]
1690 async fn compaction_keeps_tool_calls_paired_with_results() {
1691 use crate::message::Role;
1692 let llm = Arc::new(MockProvider::new(vec![
1693 Completion {
1694 content: "looking...".to_string(),
1695 tool_calls: vec![ToolCall {
1696 id: "call-1".to_string(),
1697 name: "adder".to_string(),
1698 arguments: json!({"a": 1, "b": 2}),
1699 }],
1700 finish_reason: None,
1701 usage: None,
1702 reasoning_content: None,
1703 },
1704 Completion {
1706 content: "Summary of older messages.".to_string(),
1707 tool_calls: vec![],
1708 finish_reason: Some("stop".to_string()),
1709 usage: None,
1710 reasoning_content: None,
1711 },
1712 Completion {
1713 content: "done".to_string(),
1714 tool_calls: vec![],
1715 finish_reason: Some("stop".to_string()),
1716 usage: None,
1717 reasoning_content: None,
1718 },
1719 ]));
1720 let tools = ToolRegistry::local().register(Arc::new(Adder));
1721 let compactor = crate::compact::Compactor::new(10).keep_recent_n(1);
1724 let mut agent = Agent::builder()
1725 .llm(llm)
1726 .tools(tools)
1727 .max_steps(5)
1728 .compactor(compactor)
1729 .build()
1730 .unwrap();
1731
1732 let out = agent.run("test").await.unwrap();
1733 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
1734
1735 let first = &out.transcript[0];
1741 assert!(
1742 matches!(first.role, Role::System) && first.content.contains("[compacted:"),
1743 "compaction never fired — transcript[0] = {:?}, content={:?}",
1744 first.role,
1745 &first.content.chars().take(60).collect::<String>()
1746 );
1747
1748 for (i, m) in out.transcript.iter().enumerate() {
1751 if m.role == Role::Tool {
1752 assert!(
1753 i > 0,
1754 "Tool message at index 0 (right after summary) is impossible"
1755 );
1756 let prev = &out.transcript[i - 1];
1757 assert!(
1758 matches!(prev.role, Role::Assistant) && !prev.tool_calls.is_empty(),
1759 "tool message at index {i} is orphaned — previous message \
1760 has role={:?}, tool_calls={}",
1761 prev.role,
1762 prev.tool_calls.len()
1763 );
1764 }
1765 }
1766 }
1767
1768 #[test]
1769 fn step_event_serializes_with_kind_tag() {
1770 let ev = StepEvent::AssistantText {
1771 text: "hello".into(),
1772 step: 1,
1773 };
1774 let json = serde_json::to_string(&ev).unwrap();
1775 assert!(json.contains(r#""kind":"assistant_text""#));
1776 let back: StepEvent = serde_json::from_str(&json).unwrap();
1777 assert!(
1778 matches!(back, StepEvent::AssistantText { text, step } if text == "hello" && step == 1)
1779 );
1780 }
1781
1782 #[test]
1783 fn step_event_tool_call_uses_snake_case() {
1784 let ev = StepEvent::ToolCall {
1785 call: ToolCall {
1786 id: "c1".into(),
1787 name: "read_file".into(),
1788 arguments: json!({"path": "foo.txt"}),
1789 },
1790 step: 2,
1791 };
1792 let json = serde_json::to_string(&ev).unwrap();
1793 assert!(json.contains(r#""kind":"tool_call""#));
1794 }
1795
1796 #[test]
1797 fn step_event_latency_serializes_with_kind_tag() {
1798 let ev = StepEvent::Latency {
1799 step: 3,
1800 llm_ms: 42,
1801 };
1802 let json = serde_json::to_string(&ev).unwrap();
1803 assert!(json.contains(r#""kind":"latency""#));
1804 assert!(json.contains(r#""step":3"#));
1805 assert!(json.contains(r#""llm_ms":42"#));
1806 let back: StepEvent = serde_json::from_str(&json).unwrap();
1807 assert!(matches!(back, StepEvent::Latency { step, llm_ms } if step == 3 && llm_ms == 42));
1808 }
1809
1810 #[test]
1811 fn finish_reason_serializes_with_kind_tag() {
1812 let fr = FinishReason::Stuck {
1813 repeated_call: "read_file".into(),
1814 repeats: 3,
1815 };
1816 let json = serde_json::to_string(&fr).unwrap();
1817 assert!(json.contains(r#""kind":"stuck""#));
1818 let back: FinishReason = serde_json::from_str(&json).unwrap();
1819 assert_eq!(back, fr);
1820 }
1821
1822 #[test]
1823 fn finish_reason_transcript_limit_roundtrips() {
1824 let fr = FinishReason::TranscriptLimit {
1825 chars: 4096,
1826 limit: 2048,
1827 };
1828 let json = serde_json::to_string(&fr).unwrap();
1829 let back: FinishReason = serde_json::from_str(&json).unwrap();
1830 assert_eq!(back, fr);
1831 }
1832
1833 #[tokio::test]
1834 async fn seeded_transcript_lands_before_new_goal() {
1835 let seed = vec![
1836 Message::system("sys".to_string()),
1837 Message::user("old goal".to_string()),
1838 Message::assistant("old reply".to_string()),
1839 ];
1840 let llm = Arc::new(MockProvider::new(vec![Completion {
1841 content: "fresh reply".into(),
1842 tool_calls: vec![],
1843 finish_reason: Some("stop".into()),
1844 usage: None,
1845 reasoning_content: None,
1846 }]));
1847 let mut agent = Agent::builder()
1848 .llm(llm)
1849 .seed_transcript(seed)
1850 .build()
1851 .unwrap();
1852 let out = agent.run("new goal").await.unwrap();
1853 assert_eq!(out.transcript.len(), 5);
1855 assert_eq!(out.transcript[0].content, "sys");
1856 assert_eq!(out.transcript[1].content, "old goal");
1857 assert_eq!(out.transcript[2].content, "old reply");
1858 assert_eq!(out.transcript[3].content, "new goal");
1859 assert_eq!(out.transcript[4].content, "fresh reply");
1860 }
1861
1862 #[tokio::test]
1863 async fn seed_transcript_does_not_emit_events_for_seed() {
1864 let seed = vec![
1865 Message::user("old goal".to_string()),
1866 Message::assistant("old reply".to_string()),
1867 ];
1868 let llm = Arc::new(MockProvider::new(vec![Completion {
1869 content: "fresh".into(),
1870 tool_calls: vec![],
1871 finish_reason: Some("stop".into()),
1872 usage: None,
1873 reasoning_content: None,
1874 }]));
1875 let (tx, mut rx) = mpsc::unbounded_channel();
1876 let mut agent = Agent::builder()
1877 .llm(llm)
1878 .seed_transcript(seed)
1879 .events(tx)
1880 .build()
1881 .unwrap();
1882 agent.run("new goal").await.unwrap();
1883 let mut kinds: Vec<&'static str> = Vec::new();
1884 while let Ok(ev) = rx.try_recv() {
1885 kinds.push(match ev {
1886 StepEvent::AssistantText { .. } => "text",
1887 StepEvent::Finished { .. } => "done",
1888 StepEvent::Latency { .. } => "latency",
1889 _ => "other",
1890 });
1891 }
1892 assert_eq!(kinds, vec!["latency", "text", "done"]);
1894 }
1895
1896 #[tokio::test]
1897 async fn emits_latency_event_per_llm_call() {
1898 let llm = Arc::new(MockProvider::new(vec![
1900 Completion {
1901 content: "step 1".into(),
1902 tool_calls: vec![ToolCall {
1903 id: "c1".into(),
1904 name: "add".into(),
1905 arguments: json!({"a":1,"b":1}),
1906 }],
1907 finish_reason: Some("tool_calls".into()),
1908 usage: None,
1909 reasoning_content: None,
1910 },
1911 Completion {
1912 content: "step 2".into(),
1913 tool_calls: vec![],
1914 finish_reason: Some("stop".into()),
1915 usage: None,
1916 reasoning_content: None,
1917 },
1918 ]));
1919 let tools = ToolRegistry::local().register(Arc::new(Adder));
1920 let (tx, mut rx) = mpsc::unbounded_channel();
1921 let mut agent = Agent::builder()
1922 .llm(llm)
1923 .tools(tools)
1924 .events(tx)
1925 .build()
1926 .unwrap();
1927 let out = agent.run("add twice").await.unwrap();
1928
1929 let _: u64 = out.total_llm_latency_ms;
1931
1932 let mut latency_count = 0;
1934 while let Ok(e) = rx.try_recv() {
1935 if matches!(e, StepEvent::Latency { .. }) {
1936 latency_count += 1;
1937 }
1938 }
1939 assert_eq!(latency_count, 2);
1941 }
1942
1943 #[tokio::test]
1946 async fn trims_old_tool_result_to_fit_budget() {
1947 let llm = Arc::new(MockProvider::new(vec![
1951 Completion {
1952 content: "let me run a tool".into(),
1953 tool_calls: vec![ToolCall {
1954 id: "c1".into(),
1955 name: "big".into(),
1956 arguments: json!({}),
1957 }],
1958 finish_reason: Some("tool_calls".into()),
1959 usage: None,
1960 reasoning_content: None,
1961 },
1962 Completion {
1963 content: "done".into(),
1964 tool_calls: vec![],
1965 finish_reason: Some("stop".into()),
1966 usage: None,
1967 reasoning_content: None,
1968 },
1969 ]));
1970 struct BigResultTool;
1972 #[async_trait]
1973 impl Tool for BigResultTool {
1974 fn spec(&self) -> crate::llm::ToolSpec {
1975 crate::llm::ToolSpec {
1976 name: "big".into(),
1977 description: "returns a big result".into(),
1978 parameters: json!({"type":"object"}),
1979 }
1980 }
1981 async fn execute(&self, _args: Value) -> Result<String> {
1982 Ok("x".repeat(500))
1983 }
1984 }
1985 let tools = ToolRegistry::local().register(Arc::new(BigResultTool));
1986 let mut agent = Agent::builder()
1994 .llm(llm)
1995 .tools(tools)
1996 .max_transcript_chars(100)
1997 .max_steps(10)
1998 .build()
1999 .unwrap();
2000 let out = agent.run("hi").await.unwrap();
2001 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2003 let tool_msgs: Vec<&Message> = out
2005 .transcript
2006 .iter()
2007 .filter(|m| m.role == crate::message::Role::Tool)
2008 .collect();
2009 assert!(!tool_msgs.is_empty());
2010 for msg in &tool_msgs {
2011 assert_eq!(msg.content, TRIM_PLACEHOLDER);
2012 }
2013 }
2014
2015 #[tokio::test]
2016 async fn transcript_limit_fires_when_trimming_not_enough() {
2017 let llm = Arc::new(MockProvider::new(vec![
2020 Completion {
2021 content: "x".into(),
2022 tool_calls: vec![ToolCall {
2023 id: "c1".into(),
2024 name: "add".into(),
2025 arguments: json!({"a":1,"b":1}),
2026 }],
2027 finish_reason: Some("tool_calls".into()),
2028 usage: None,
2029 reasoning_content: None,
2030 },
2031 Completion {
2032 content: "y".into(),
2033 tool_calls: vec![],
2034 finish_reason: Some("stop".into()),
2035 usage: None,
2036 reasoning_content: None,
2037 },
2038 ]));
2039 let tools = ToolRegistry::local().register(Arc::new(Adder));
2040 let mut agent = Agent::builder()
2042 .llm(llm)
2043 .tools(tools)
2044 .max_transcript_chars(1)
2045 .max_steps(10)
2046 .build()
2047 .unwrap();
2048 let out = agent.run("hi").await.unwrap();
2049 assert!(matches!(out.finish, FinishReason::TranscriptLimit { .. }));
2050 }
2051
2052 #[tokio::test]
2057 async fn permission_hook_allow_passes_args_unchanged() {
2058 let llm = Arc::new(MockProvider::new(vec![
2060 Completion {
2061 content: "let me add".into(),
2062 tool_calls: vec![ToolCall {
2063 id: "c1".into(),
2064 name: "add".into(),
2065 arguments: json!({"a": 2, "b": 3}),
2066 }],
2067 finish_reason: Some("tool_calls".into()),
2068 usage: None,
2069 reasoning_content: None,
2070 },
2071 Completion {
2072 content: "5".into(),
2073 tool_calls: vec![],
2074 finish_reason: Some("stop".into()),
2075 usage: None,
2076 reasoning_content: None,
2077 },
2078 ]));
2079 let tools = ToolRegistry::local().register(Arc::new(Adder));
2080 let mut agent = Agent::builder()
2081 .llm(llm)
2082 .tools(tools)
2083 .permission_hook(|_name, _args| PermissionDecision::Allow)
2084 .build()
2085 .unwrap();
2086 let out = agent.run("what is 2+3?").await.unwrap();
2087 assert_eq!(out.final_message.as_deref(), Some("5"));
2088 assert_eq!(out.steps, 2);
2089 }
2090
2091 #[tokio::test]
2092 async fn permission_hook_deny_returns_error_to_model() {
2093 let llm = Arc::new(MockProvider::new(vec![
2096 Completion {
2097 content: "let me add".into(),
2098 tool_calls: vec![ToolCall {
2099 id: "c1".into(),
2100 name: "add".into(),
2101 arguments: json!({"a": 2, "b": 3}),
2102 }],
2103 finish_reason: Some("tool_calls".into()),
2104 usage: None,
2105 reasoning_content: None,
2106 },
2107 Completion {
2108 content: "i see the error".into(),
2109 tool_calls: vec![],
2110 finish_reason: Some("stop".into()),
2111 usage: None,
2112 reasoning_content: None,
2113 },
2114 ]));
2115 let tools = ToolRegistry::local().register(Arc::new(Adder));
2116 let mut agent = Agent::builder()
2117 .llm(llm)
2118 .tools(tools)
2119 .permission_hook(|_name, _args| PermissionDecision::Deny("not allowed".into()))
2120 .build()
2121 .unwrap();
2122 let out = agent.run("add numbers").await.unwrap();
2123 let tool_msgs: Vec<&Message> = out
2125 .transcript
2126 .iter()
2127 .filter(|m| m.role == crate::message::Role::Tool)
2128 .collect();
2129 assert_eq!(tool_msgs.len(), 1);
2130 assert!(tool_msgs[0].content.contains("not allowed"));
2131 assert_eq!(out.final_message.as_deref(), Some("i see the error"));
2133 }
2134
2135 #[tokio::test]
2136 async fn permission_hook_transform_replaces_args() {
2137 let llm = Arc::new(MockProvider::new(vec![
2140 Completion {
2141 content: "let me add".into(),
2142 tool_calls: vec![ToolCall {
2143 id: "c1".into(),
2144 name: "add".into(),
2145 arguments: json!({"a": 100, "b": 200}),
2146 }],
2147 finish_reason: Some("tool_calls".into()),
2148 usage: None,
2149 reasoning_content: None,
2150 },
2151 Completion {
2152 content: "done".into(),
2153 tool_calls: vec![],
2154 finish_reason: Some("stop".into()),
2155 usage: None,
2156 reasoning_content: None,
2157 },
2158 ]));
2159 let tools = ToolRegistry::local().register(Arc::new(Adder));
2160 let mut agent = Agent::builder()
2162 .llm(llm)
2163 .tools(tools)
2164 .permission_hook(|_name, _args| PermissionDecision::Transform(json!({"a": 1, "b": 2})))
2165 .build()
2166 .unwrap();
2167 let out = agent.run("add numbers").await.unwrap();
2168 let tool_msgs: Vec<&Message> = out
2170 .transcript
2171 .iter()
2172 .filter(|m| m.role == crate::message::Role::Tool)
2173 .collect();
2174 assert_eq!(tool_msgs.len(), 1);
2175 assert_eq!(tool_msgs[0].content, "3");
2176 }
2177
2178 #[tokio::test]
2179 async fn default_no_hook_is_unchanged() {
2180 let llm = Arc::new(MockProvider::new(vec![Completion {
2183 content: "done".into(),
2184 tool_calls: vec![],
2185 finish_reason: Some("stop".into()),
2186 usage: None,
2187 reasoning_content: None,
2188 }]));
2189 let mut agent = Agent::builder().llm(llm).build().unwrap();
2190 let out = agent.run("hi").await.unwrap();
2191 assert_eq!(out.final_message.as_deref(), Some("done"));
2192 assert_eq!(out.steps, 1);
2193 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2194 }
2195 #[test]
2198 fn planning_mode_default_is_immediate() {
2199 assert_eq!(PlanningMode::default(), PlanningMode::Immediate);
2200 }
2201
2202 #[test]
2203 fn planning_mode_variants_are_distinct() {
2204 assert_ne!(PlanningMode::Immediate, PlanningMode::PlanFirst);
2205 }
2206
2207 #[tokio::test]
2208 async fn builder_with_planfirst_succeeds() {
2209 let llm = Arc::new(MockProvider::new(vec![Completion {
2210 content: "done".into(),
2211 tool_calls: vec![],
2212 finish_reason: Some("stop".into()),
2213 usage: None,
2214 reasoning_content: None,
2215 }]));
2216 let agent = Agent::builder()
2217 .llm(llm)
2218 .planning_mode(PlanningMode::PlanFirst)
2219 .build()
2220 .unwrap();
2221 assert_eq!(agent.planning_mode, PlanningMode::PlanFirst);
2222 assert!(agent.plan_buffer.is_none());
2223 }
2224
2225 #[tokio::test]
2226 async fn immediate_mode_runs_normally() {
2227 let llm = Arc::new(MockProvider::new(vec![Completion {
2228 content: "done".into(),
2229 tool_calls: vec![],
2230 finish_reason: Some("stop".into()),
2231 usage: None,
2232 reasoning_content: None,
2233 }]));
2234 let mut agent = Agent::builder().llm(llm).build().unwrap();
2235 let outcome = agent.run("test").await.unwrap();
2236 assert_eq!(outcome.finish, FinishReason::NoMoreToolCalls);
2237 }
2238
2239 #[test]
2240 fn plan_proposed_can_be_constructed() {
2241 let event = StepEvent::PlanProposed {
2242 plan_text: "Step 1: read file, Step 2: edit file".into(),
2243 tool_calls: vec![ToolCall {
2244 id: "call_1".into(),
2245 name: "read_file".into(),
2246 arguments: json!({"path": "test.txt"}),
2247 }],
2248 };
2249 match &event {
2250 StepEvent::PlanProposed {
2251 plan_text,
2252 tool_calls,
2253 } => {
2254 assert!(plan_text.contains("read file"));
2255 assert_eq!(tool_calls.len(), 1);
2256 }
2257 _ => panic!("wrong variant"),
2258 }
2259 }
2260
2261 #[test]
2262 fn plan_confirmed_can_be_constructed() {
2263 let event = StepEvent::PlanConfirmed;
2264 assert!(matches!(event, StepEvent::PlanConfirmed));
2265 }
2266
2267 #[test]
2268 fn plan_rejected_can_be_constructed() {
2269 let event = StepEvent::PlanRejected {
2270 reason: "too many steps".into(),
2271 };
2272 match &event {
2273 StepEvent::PlanRejected { reason } => {
2274 assert_eq!(reason, "too many steps");
2275 }
2276 _ => panic!("wrong variant"),
2277 }
2278 }
2279
2280 #[test]
2281 fn confirm_plan_does_not_panic() {
2282 let mut agent = Agent {
2283 llm: Arc::new(MockProvider::new(vec![])),
2284 tools: ToolRegistry::local(),
2285 transcript: vec![],
2286 max_steps: 32,
2287 max_transcript_chars: None,
2288 events: None,
2289 streaming: false,
2290 total_llm_latency_ms: 0,
2291 compactor: None,
2292 permission_hook: None,
2293 hooks: HookRegistry::new(),
2294 planning_mode: PlanningMode::Immediate,
2295 plan_buffer: Some(vec![]),
2296 plan_confirmed: false,
2297 };
2298 agent.confirm_plan();
2299 assert!(agent.plan_confirmed);
2300 }
2301
2302 #[test]
2303 fn reject_plan_does_not_panic() {
2304 let mut agent = Agent {
2305 llm: Arc::new(MockProvider::new(vec![])),
2306 tools: ToolRegistry::local(),
2307 transcript: vec![],
2308 max_steps: 32,
2309 max_transcript_chars: None,
2310 events: None,
2311 streaming: false,
2312 total_llm_latency_ms: 0,
2313 compactor: None,
2314 permission_hook: None,
2315 hooks: HookRegistry::new(),
2316 planning_mode: PlanningMode::Immediate,
2317 plan_buffer: Some(vec![]),
2318 plan_confirmed: false,
2319 };
2320 agent.reject_plan("bad plan");
2321 assert!(agent.plan_buffer.is_none());
2322 assert!(!agent.plan_confirmed);
2323 }
2324}
2325#[cfg(test)]
2330mod tracing_tests {
2331 use crate::llm::Completion;
2332 use crate::llm::MockProvider;
2333 use crate::Agent;
2334 use std::sync::Arc;
2335 use tracing_test::traced_test;
2336
2337 #[traced_test]
2338 #[tokio::test]
2339 async fn agent_run_creates_span() {
2340 let llm = Arc::new(MockProvider::new(vec![Completion {
2341 content: "done".into(),
2342 tool_calls: vec![],
2343 finish_reason: Some("stop".into()),
2344 usage: None,
2345 reasoning_content: None,
2346 }]));
2347 let mut agent = Agent::builder().llm(llm).build().unwrap();
2348 agent.run("test goal").await.unwrap();
2349
2350 assert!(logs_contain("run:"));
2352 }
2353
2354 #[traced_test]
2355 #[tokio::test]
2356 async fn agent_step_spans_nested_under_run() {
2357 let llm = Arc::new(MockProvider::new(vec![Completion {
2358 content: "step 1".into(),
2359 tool_calls: vec![],
2360 finish_reason: Some("stop".into()),
2361 usage: None,
2362 reasoning_content: None,
2363 }]));
2364 let mut agent = Agent::builder().llm(llm).build().unwrap();
2365 agent.run("test").await.unwrap();
2366
2367 assert!(logs_contain("run:"));
2369 assert!(logs_contain("step="));
2370 }
2371}
2372
2373#[cfg(test)]
2379mod parallel_tests {
2380 use super::*;
2381 use crate::llm::{Completion, MockProvider, ToolCall};
2382 use crate::tools::Tool;
2383 use crate::ToolSpec;
2384 use async_trait::async_trait;
2385 use serde_json::{json, Value};
2386 use std::sync::atomic::{AtomicUsize, Ordering};
2387 use std::sync::Arc;
2388 use std::time::Duration;
2389
2390 struct SlowReadOnly {
2392 counter: Arc<AtomicUsize>,
2393 delay_ms: u64,
2394 }
2395
2396 #[async_trait]
2397 impl Tool for SlowReadOnly {
2398 fn spec(&self) -> ToolSpec {
2399 ToolSpec {
2400 name: "slow_read".into(),
2401 description: "a slow read-only tool".into(),
2402 parameters: json!({"type": "object"}),
2403 }
2404 }
2405 fn is_readonly(&self) -> bool {
2406 true
2407 }
2408 async fn execute(&self, _args: Value) -> Result<String> {
2409 let order = self.counter.fetch_add(1, Ordering::SeqCst);
2410 tokio::time::sleep(Duration::from_millis(self.delay_ms)).await;
2411 Ok(format!("read-{}", order))
2412 }
2413 }
2414
2415 struct TrackingWrite {
2417 counter: Arc<AtomicUsize>,
2418 }
2419
2420 #[async_trait]
2421 impl Tool for TrackingWrite {
2422 fn spec(&self) -> ToolSpec {
2423 ToolSpec {
2424 name: "track_write".into(),
2425 description: "a tracked write tool".into(),
2426 parameters: json!({"type": "object"}),
2427 }
2428 }
2429 async fn execute(&self, _args: Value) -> Result<String> {
2430 let order = self.counter.fetch_add(1, Ordering::SeqCst);
2431 Ok(format!("write-{}", order))
2432 }
2433 }
2434
2435 #[tokio::test]
2436 async fn read_only_tools_execute_in_parallel() {
2437 let counter = Arc::new(AtomicUsize::new(0));
2440 let tool1 = Arc::new(SlowReadOnly {
2441 counter: counter.clone(),
2442 delay_ms: 100,
2443 });
2444 let tool2 = Arc::new(SlowReadOnly {
2445 counter: counter.clone(),
2446 delay_ms: 100,
2447 });
2448
2449 let tools = ToolRegistry::local().register(tool1).register(tool2);
2450
2451 let llm = Arc::new(MockProvider::new(vec![
2452 Completion {
2453 content: "reading...".into(),
2454 tool_calls: vec![
2455 ToolCall {
2456 id: "c1".into(),
2457 name: "slow_read".into(),
2458 arguments: json!({}),
2459 },
2460 ToolCall {
2461 id: "c2".into(),
2462 name: "slow_read".into(),
2463 arguments: json!({}),
2464 },
2465 ],
2466 finish_reason: Some("tool_calls".into()),
2467 usage: None,
2468 reasoning_content: None,
2469 },
2470 Completion {
2471 content: "done".into(),
2472 tool_calls: vec![],
2473 finish_reason: Some("stop".into()),
2474 usage: None,
2475 reasoning_content: None,
2476 },
2477 ]));
2478
2479 let start = std::time::Instant::now();
2480 let mut agent = Agent::builder()
2481 .llm(llm)
2482 .tools(tools)
2483 .max_steps(5)
2484 .build()
2485 .unwrap();
2486 let out = agent.run("read in parallel").await.unwrap();
2487 let elapsed = start.elapsed();
2488
2489 assert!(
2491 elapsed.as_millis() < 150,
2492 "parallel execution took {}ms, expected <150ms",
2493 elapsed.as_millis()
2494 );
2495 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2496
2497 let tool_msgs: Vec<&Message> = out
2499 .transcript
2500 .iter()
2501 .filter(|m| m.role == crate::message::Role::Tool)
2502 .collect();
2503 assert_eq!(tool_msgs.len(), 2);
2504 }
2505
2506 #[tokio::test]
2507 async fn write_tools_execute_sequentially() {
2508 let counter = Arc::new(AtomicUsize::new(0));
2511 let tool1 = Arc::new(TrackingWrite {
2512 counter: counter.clone(),
2513 });
2514 let tool2 = Arc::new(TrackingWrite {
2515 counter: counter.clone(),
2516 });
2517
2518 let tools = ToolRegistry::local().register(tool1).register(tool2);
2519
2520 let llm = Arc::new(MockProvider::new(vec![
2521 Completion {
2522 content: "writing...".into(),
2523 tool_calls: vec![
2524 ToolCall {
2525 id: "c1".into(),
2526 name: "track_write".into(),
2527 arguments: json!({}),
2528 },
2529 ToolCall {
2530 id: "c2".into(),
2531 name: "track_write".into(),
2532 arguments: json!({}),
2533 },
2534 ],
2535 finish_reason: Some("tool_calls".into()),
2536 usage: None,
2537 reasoning_content: None,
2538 },
2539 Completion {
2540 content: "done".into(),
2541 tool_calls: vec![],
2542 finish_reason: Some("stop".into()),
2543 usage: None,
2544 reasoning_content: None,
2545 },
2546 ]));
2547
2548 let mut agent = Agent::builder()
2549 .llm(llm)
2550 .tools(tools)
2551 .max_steps(5)
2552 .build()
2553 .unwrap();
2554 let out = agent.run("write sequentially").await.unwrap();
2555
2556 let tool_msgs: Vec<&Message> = out
2558 .transcript
2559 .iter()
2560 .filter(|m| m.role == crate::message::Role::Tool)
2561 .collect();
2562 assert_eq!(tool_msgs.len(), 2);
2563 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2564 }
2565
2566 #[tokio::test]
2567 async fn mixed_read_write_preserves_order() {
2568 let counter = Arc::new(AtomicUsize::new(0));
2572 let read_tool = Arc::new(SlowReadOnly {
2573 counter: counter.clone(),
2574 delay_ms: 10,
2575 });
2576 let write_tool = Arc::new(TrackingWrite {
2577 counter: counter.clone(),
2578 });
2579
2580 let tools = ToolRegistry::local()
2581 .register(read_tool)
2583 .register(write_tool);
2584
2585 let llm = Arc::new(MockProvider::new(vec![
2586 Completion {
2587 content: "mixed...".into(),
2588 tool_calls: vec![
2589 ToolCall {
2590 id: "c1".into(),
2591 name: "slow_read".into(),
2592 arguments: json!({}),
2593 },
2594 ToolCall {
2595 id: "c2".into(),
2596 name: "track_write".into(),
2597 arguments: json!({}),
2598 },
2599 ToolCall {
2600 id: "c3".into(),
2601 name: "slow_read".into(),
2602 arguments: json!({}),
2603 },
2604 ],
2605 finish_reason: Some("tool_calls".into()),
2606 usage: None,
2607 reasoning_content: None,
2608 },
2609 Completion {
2610 content: "done".into(),
2611 tool_calls: vec![],
2612 finish_reason: Some("stop".into()),
2613 usage: None,
2614 reasoning_content: None,
2615 },
2616 ]));
2617
2618 let mut agent = Agent::builder()
2619 .llm(llm)
2620 .tools(tools)
2621 .max_steps(5)
2622 .build()
2623 .unwrap();
2624 let out = agent.run("mixed").await.unwrap();
2625
2626 let tool_msgs: Vec<&Message> = out
2628 .transcript
2629 .iter()
2630 .filter(|m| m.role == crate::message::Role::Tool)
2631 .collect();
2632 assert_eq!(tool_msgs.len(), 3);
2633 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2634 }
2635
2636 #[tokio::test]
2637 async fn parallel_read_only_with_unknown_tool() {
2638 let counter = Arc::new(AtomicUsize::new(0));
2641 let read_tool = Arc::new(SlowReadOnly {
2642 counter: counter.clone(),
2643 delay_ms: 10,
2644 });
2645
2646 let tools = ToolRegistry::local().register(read_tool);
2647
2648 let llm = Arc::new(MockProvider::new(vec![
2649 Completion {
2650 content: "mixed...".into(),
2651 tool_calls: vec![
2652 ToolCall {
2653 id: "c1".into(),
2654 name: "slow_read".into(),
2655 arguments: json!({}),
2656 },
2657 ToolCall {
2658 id: "c2".into(),
2659 name: "unknown_tool".into(),
2660 arguments: json!({}),
2661 },
2662 ],
2663 finish_reason: Some("tool_calls".into()),
2664 usage: None,
2665 reasoning_content: None,
2666 },
2667 Completion {
2668 content: "done".into(),
2669 tool_calls: vec![],
2670 finish_reason: Some("stop".into()),
2671 usage: None,
2672 reasoning_content: None,
2673 },
2674 ]));
2675
2676 let mut agent = Agent::builder()
2677 .llm(llm)
2678 .tools(tools)
2679 .max_steps(5)
2680 .build()
2681 .unwrap();
2682 let out = agent.run("mixed with unknown").await.unwrap();
2683
2684 let tool_msgs: Vec<&Message> = out
2686 .transcript
2687 .iter()
2688 .filter(|m| m.role == crate::message::Role::Tool)
2689 .collect();
2690 assert_eq!(tool_msgs.len(), 2);
2691 assert!(tool_msgs[1].content.contains("ERROR"));
2693 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2694 }
2695
2696 #[tokio::test]
2697 async fn mixed_read_write_executes_sequentially() {
2698 let counter = Arc::new(AtomicUsize::new(0));
2702 let read_tool = Arc::new(SlowReadOnly {
2703 counter: counter.clone(),
2704 delay_ms: 100,
2705 });
2706 let write_tool = Arc::new(TrackingWrite {
2707 counter: counter.clone(),
2708 });
2709
2710 let tools = ToolRegistry::local()
2711 .register(read_tool)
2712 .register(write_tool);
2713
2714 let llm = Arc::new(MockProvider::new(vec![
2715 Completion {
2716 content: "mixed rw...".into(),
2717 tool_calls: vec![
2718 ToolCall {
2719 id: "c1".into(),
2720 name: "slow_read".into(),
2721 arguments: json!({}),
2722 },
2723 ToolCall {
2724 id: "c2".into(),
2725 name: "track_write".into(),
2726 arguments: json!({}),
2727 },
2728 ToolCall {
2729 id: "c3".into(),
2730 name: "slow_read".into(),
2731 arguments: json!({}),
2732 },
2733 ],
2734 finish_reason: Some("tool_calls".into()),
2735 usage: None,
2736 reasoning_content: None,
2737 },
2738 Completion {
2739 content: "done".into(),
2740 tool_calls: vec![],
2741 finish_reason: Some("stop".into()),
2742 usage: None,
2743 reasoning_content: None,
2744 },
2745 ]));
2746
2747 let start = std::time::Instant::now();
2748 let mut agent = Agent::builder()
2749 .llm(llm)
2750 .tools(tools)
2751 .max_steps(5)
2752 .build()
2753 .unwrap();
2754 let out = agent.run("mixed read write").await.unwrap();
2755 let elapsed = start.elapsed();
2756
2757 assert!(
2759 elapsed.as_millis() >= 200,
2760 "mixed read/write took {}ms, expected >=200ms (write breaks parallel batching)",
2761 elapsed.as_millis()
2762 );
2763
2764 let tool_msgs: Vec<&Message> = out
2765 .transcript
2766 .iter()
2767 .filter(|m| m.role == crate::message::Role::Tool)
2768 .collect();
2769 assert_eq!(tool_msgs.len(), 3);
2770 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2771 }
2772
2773 #[tokio::test]
2774 async fn single_tool_call_unaffected_by_parallel_mode() {
2775 let counter = Arc::new(AtomicUsize::new(0));
2778 let read_tool = Arc::new(SlowReadOnly {
2779 counter: counter.clone(),
2780 delay_ms: 10,
2781 });
2782
2783 let tools = ToolRegistry::local().register(read_tool);
2784
2785 let llm = Arc::new(MockProvider::new(vec![
2786 Completion {
2787 content: "single call...".into(),
2788 tool_calls: vec![ToolCall {
2789 id: "c1".into(),
2790 name: "slow_read".into(),
2791 arguments: json!({}),
2792 }],
2793 finish_reason: Some("tool_calls".into()),
2794 usage: None,
2795 reasoning_content: None,
2796 },
2797 Completion {
2798 content: "done".into(),
2799 tool_calls: vec![],
2800 finish_reason: Some("stop".into()),
2801 usage: None,
2802 reasoning_content: None,
2803 },
2804 ]));
2805
2806 let mut agent = Agent::builder()
2807 .llm(llm)
2808 .tools(tools)
2809 .max_steps(5)
2810 .build()
2811 .unwrap();
2812 let out = agent.run("single tool call").await.unwrap();
2813
2814 let tool_msgs: Vec<&Message> = out
2815 .transcript
2816 .iter()
2817 .filter(|m| m.role == crate::message::Role::Tool)
2818 .collect();
2819 assert_eq!(tool_msgs.len(), 1);
2820 assert!(tool_msgs[0].content.contains("read-0"));
2821 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2822 }
2823
2824 #[tokio::test]
2825 async fn multiple_write_tools_preserve_order() {
2826 let counter = Arc::new(AtomicUsize::new(0));
2828 let write_tool = Arc::new(TrackingWrite {
2829 counter: counter.clone(),
2830 });
2831
2832 let tools = ToolRegistry::local().register(write_tool);
2833
2834 let llm = Arc::new(MockProvider::new(vec![
2835 Completion {
2836 content: "writing three...".into(),
2837 tool_calls: vec![
2838 ToolCall {
2839 id: "c1".into(),
2840 name: "track_write".into(),
2841 arguments: json!({}),
2842 },
2843 ToolCall {
2844 id: "c2".into(),
2845 name: "track_write".into(),
2846 arguments: json!({}),
2847 },
2848 ToolCall {
2849 id: "c3".into(),
2850 name: "track_write".into(),
2851 arguments: json!({}),
2852 },
2853 ],
2854 finish_reason: Some("tool_calls".into()),
2855 usage: None,
2856 reasoning_content: None,
2857 },
2858 Completion {
2859 content: "done".into(),
2860 tool_calls: vec![],
2861 finish_reason: Some("stop".into()),
2862 usage: None,
2863 reasoning_content: None,
2864 },
2865 ]));
2866
2867 let mut agent = Agent::builder()
2868 .llm(llm)
2869 .tools(tools)
2870 .max_steps(5)
2871 .build()
2872 .unwrap();
2873 let out = agent.run("write three times").await.unwrap();
2874
2875 let tool_msgs: Vec<&Message> = out
2876 .transcript
2877 .iter()
2878 .filter(|m| m.role == crate::message::Role::Tool)
2879 .collect();
2880 assert_eq!(tool_msgs.len(), 3);
2881 assert!(tool_msgs[0].content.contains("write-0"));
2883 assert!(tool_msgs[1].content.contains("write-1"));
2884 assert!(tool_msgs[2].content.contains("write-2"));
2885 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2886 }
2887
2888 #[tokio::test]
2889 async fn no_tool_calls_completes_immediately() {
2890 let llm = Arc::new(MockProvider::new(vec![Completion {
2893 content: "nothing to do".into(),
2894 tool_calls: vec![],
2895 finish_reason: Some("stop".into()),
2896 usage: None,
2897 reasoning_content: None,
2898 }]));
2899
2900 let mut agent = Agent::builder().llm(llm).max_steps(5).build().unwrap();
2901 let out = agent.run("no tools needed").await.unwrap();
2902
2903 assert_eq!(out.finish, FinishReason::NoMoreToolCalls);
2904 assert_eq!(out.final_message.as_deref(), Some("nothing to do"));
2905 assert_eq!(out.steps, 1);
2906 let tool_msgs: Vec<&Message> = out
2908 .transcript
2909 .iter()
2910 .filter(|m| m.role == crate::message::Role::Tool)
2911 .collect();
2912 assert_eq!(tool_msgs.len(), 0);
2913 }
2914}